M-Labs
/
artiq
mirror of https://github.com/m-labs/artiq.git
2
0
Fork 2
Code Issues Projects Releases Wiki Activity

compiler: Slight array op implementation cleanup [nfc] 

array_unaryop_funcs was never used; since the mangled names
are unique, a single dictionary would be nicer for overrides
anyway.s
This commit is contained in:
David Nadlinger 2020-08-09 17:03:56 +01:00
parent 33d931a5b7
commit 8e262acd1e
1 changed files with 12 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
artiq/compiler/transforms/artiq_ir_generator.py
View File

@ -88,10 +88,8 @@ class ARTIQIRGenerator(algorithm.Visitor):
necessary. They are kept track of in global dictionaries, with a mangled name necessary. They are kept track of in global dictionaries, with a mangled name
containing types and operations as key: containing types and operations as key:
:ivar array_unaryop_funcs: the map from mangled name to implementation of unary :ivar array_op_funcs: the map from mangled name to implementation of
operations for arrays operations on/between arrays
:ivar array_binop_funcs: the map from mangled name to implementation of binary
operations between arrays
""" """
_size_type = builtins.TInt32() _size_type = builtins.TInt32()
@ -120,8 +118,7 @@ class ARTIQIRGenerator(algorithm.Visitor):
self.function_map = dict() self.function_map = dict()
self.variable_map = dict() self.variable_map = dict()
self.method_map = defaultdict(lambda: []) self.method_map = defaultdict(lambda: [])
self.array_unaryop_funcs = dict() self.array_op_funcs = dict()
self.array_binop_funcs = dict()
def annotate_calls(self, devirtualization): def annotate_calls(self, devirtualization):
for var_node in devirtualization.variable_map: for var_node in devirtualization.variable_map:
@ -1392,10 +1389,10 @@ class ARTIQIRGenerator(algorithm.Visitor):
def _get_array_unaryop(self, name, make_op, result_type, arg_type): def _get_array_unaryop(self, name, make_op, result_type, arg_type):
name = "_array_{}_{}".format( name = "_array_{}_{}".format(
name, self._mangle_arrayop_types([result_type, arg_type])) name, self._mangle_arrayop_types([result_type, arg_type]))
if name not in self.array_unaryop_funcs: if name not in self.array_op_funcs:
self.array_binop_funcs[name] = self._make_array_unaryop( self.array_op_funcs[name] = self._make_array_unaryop(
name, make_op, result_type, arg_type) name, make_op, result_type, arg_type)
return self.array_binop_funcs[name] return self.array_op_funcs[name]
def visit_UnaryOpT(self, node): def visit_UnaryOpT(self, node):
if isinstance(node.op, ast.Not): if isinstance(node.op, ast.Not):
@ -1523,7 +1520,7 @@ class ARTIQIRGenerator(algorithm.Visitor):
name = "_array_{}_{}".format( name = "_array_{}_{}".format(
type(op).__name__, type(op).__name__,
self._mangle_arrayop_types([result_type, lhs_type, rhs_type])) self._mangle_arrayop_types([result_type, lhs_type, rhs_type]))
if name not in self.array_binop_funcs: if name not in self.array_op_funcs:
def body_gen(result, lhs, rhs): def body_gen(result, lhs, rhs):
# At this point, shapes are assumed to match; could just pass buffer # At this point, shapes are assumed to match; could just pass buffer
@ -1561,9 +1558,9 @@ class ARTIQIRGenerator(algorithm.Visitor):
ir.Constant(0, self._size_type), lambda index: self.append( ir.Constant(0, self._size_type), lambda index: self.append(
ir.Compare(ast.Lt(loc=None), index, num_total_elts)), loop_gen) ir.Compare(ast.Lt(loc=None), index, num_total_elts)), loop_gen)
self.array_binop_funcs[name] = self._make_array_binop( self.array_op_funcs[name] = self._make_array_binop(
name, result_type, lhs_type, rhs_type, body_gen) name, result_type, lhs_type, rhs_type, body_gen)
return self.array_binop_funcs[name] return self.array_op_funcs[name]
def _invoke_arrayop(self, func, params): def _invoke_arrayop(self, func, params):
closure = self.append( closure = self.append(
@ -1585,7 +1582,7 @@ class ARTIQIRGenerator(algorithm.Visitor):
def _get_matmult(self, result_type, lhs_type, rhs_type): def _get_matmult(self, result_type, lhs_type, rhs_type):
name = "_array_MatMult_" + self._mangle_arrayop_types( name = "_array_MatMult_" + self._mangle_arrayop_types(
[result_type, lhs_type, rhs_type]) [result_type, lhs_type, rhs_type])
if name not in self.array_binop_funcs: if name not in self.array_op_funcs:
def body_gen(result, lhs, rhs): def body_gen(result, lhs, rhs):
assert builtins.is_array(result.type), \ assert builtins.is_array(result.type), \
@ -1656,9 +1653,9 @@ class ARTIQIRGenerator(algorithm.Visitor):
ir.Constant(0, self._size_type), lambda index: self.append( ir.Constant(0, self._size_type), lambda index: self.append(
ir.Compare(ast.Lt(loc=None), index, num_rows)), row_loop) ir.Compare(ast.Lt(loc=None), index, num_rows)), row_loop)
self.array_binop_funcs[name] = self._make_array_binop( self.array_op_funcs[name] = self._make_array_binop(
name, result_type, lhs_type, rhs_type, body_gen) name, result_type, lhs_type, rhs_type, body_gen)
return self.array_binop_funcs[name] return self.array_op_funcs[name]
def _get_matmult_shapes(self, lhs, rhs): def _get_matmult_shapes(self, lhs, rhs):
lhs_shape = self.append(ir.GetAttr(lhs, "shape")) lhs_shape = self.append(ir.GetAttr(lhs, "shape"))