class Simplify(Transform): def __init__(self): transform.Transform.__init__(self) # associate var names with any immutable values # they are bound to self.bindings = ScopedDict() # which expressions have already been computed # and stored in some variable? self.available_expressions = ScopedDict() def pre_apply(self, fn): ma = TypeBasedMutabilityAnalysis() # which types have elements that might # change between two accesses? self.mutable_types = ma.visit_fn(fn) self.use_counts = use_count(fn) def immutable_type(self, t): return t not in self.mutable_types _immutable_classes = set([Const, Var, Closure, ClosureElt, Tuple, TupleProj, Cast, PrimCall, TypedFn, UntypedFn, ArrayView, Slice, Map, Reduce, Scan, OuterMap, IndexMap, IndexReduce, IndexScan, ]) def immutable(self, expr): """ TODO: make all this mutability/immutability stuff sane """ klass = expr.__class__ result = (klass in self._immutable_classes and (all(self.immutable(c) for c in expr.children()))) or \ (klass is Attribute and isinstance(expr.type, ImmutableT)) return result def temp(self, expr, name = None, use_count = 1): """ Wrapper around Codegen.assign_name which also updates bindings and use_counts """ if self.is_simple(expr): return expr else: new_var = self.assign_name(expr, name = name) self.bindings[new_var.name] = expr self.use_counts[new_var.name] = use_count return new_var def transform_expr(self, expr): if self.is_simple(expr): if expr.type == NoneType: return none else: return Transform.transform_expr(self, expr) stored = self.available_expressions.get(expr) if stored is not None: return stored return Transform.transform_expr(self, expr) def transform_Var(self, expr): t = expr.type if t.__class__ is NoneT: return none elif t.__class__ is SliceT and \ t.start_type == NoneType and \ t.stop_type == NoneType and \ t.step_type == NoneType: return slice_none name = expr.name prev_expr = expr while name in self.bindings: prev_expr = expr expr = self.bindings[name] if expr.__class__ is Var: name = expr.name else: break c = expr.__class__ if c is Var or c is Const: return expr else: return prev_expr def transform_Cast(self, expr): v = self.transform_expr(expr.value) if v.type == expr.type: return v elif v.__class__ is Const and isinstance(expr.type, ScalarT): return Const(expr.type.dtype.type(v.value), type = expr.type) elif self.is_simple(v): expr.value = v return expr else: expr.value = self.assign_name(v) return expr def transform_Attribute(self, expr): v = self.transform_expr(expr.value) if v.__class__ is Var and v.name in self.bindings: stored_v = self.bindings[v.name] c = stored_v.__class__ if c is Var or c is Struct: v = stored_v elif c is ArrayView: if expr.name == 'shape': return self.transform_expr(stored_v.shape) elif expr.name == 'strides': return self.transform_expr(stored_v.strides) elif expr.name == 'data': return self.transform_expr(stored_v.data) elif c is AllocArray: if expr.name == 'shape': return self.transform_expr(stored_v.shape) elif c is Slice: if expr.name == "start": return self.transform_expr(stored_v.start) elif expr.name == "stop": return self.transform_expr(stored_v.stop) else: assert expr.name == "step", "Unexpected attribute for slice: %s" % expr.name return self.transform_expr(stored_v.step) if v.__class__ is Struct: idx = v.type.field_pos(expr.name) return v.args[idx] elif v.__class__ is not Var: v = self.temp(v, "struct") if expr.value == v: return expr else: return Attribute(value = v, name = expr.name, type = expr.type) def transform_Closure(self, expr): expr.args = tuple(self.transform_simple_exprs(expr.args)) return expr def transform_Tuple(self, expr): expr.elts = tuple( self.transform_simple_exprs(expr.elts)) return expr def transform_TupleProj(self, expr): idx = expr.index assert isinstance(idx, int), \ "TupleProj index must be an integer, got: " + str(idx) new_tuple = self.transform_expr(expr.tuple) if new_tuple.__class__ is Var and new_tuple.name in self.bindings: tuple_expr = self.bindings[new_tuple.name] if tuple_expr.__class__ is Tuple: assert idx < len(tuple_expr.elts), \ "Too few elements in tuple %s : %s, elts = %s" % (expr, tuple_expr.type, tuple_expr.elts) return tuple_expr.elts[idx] elif tuple_expr.__class__ is Struct: assert idx < len(tuple_expr.args), \ "Too few args in closure %s : %s, elts = %s" % (expr, tuple_expr.type, tuple_expr.elts) return tuple_expr.args[idx] #if not self.is_simple(new_tuple): # complex_expr = new_tuple # new_tuple = self.assign_name(complex_expr, "tuple") # print "MADE ME A NEW TUPLE", complex_expr, new_tuple expr.tuple = new_tuple return expr def transform_ClosureElt(self, expr): idx = expr.index assert isinstance(idx, int), \ "ClosureElt index must be an integer, got: " + str(idx) new_closure = self.transform_expr(expr.closure) if new_closure.__class__ is Var and new_closure.name in self.bindings: closure_expr = self.bindings[new_closure.name] if closure_expr.__class__ is Closure: return closure_expr.args[idx] if not self.is_simple(new_closure): new_closure = self.assign_name(new_closure, "closure") expr.closure = new_closure return expr def transform_Call(self, expr): fn = self.transform_expr(expr.fn) args = self.transform_simple_exprs(expr.args) if fn.type.__class__ is ClosureT: closure_elts = self.closure_elts(fn) combined_args = tuple(closure_elts) + tuple(args) if fn.type.fn.__class__ is TypedFn: fn = fn.type.fn else: assert isinstance(fn.type.fn, UntypedFn) from .. type_inference import specialize fn = specialize(fn, get_types(combined_args)) assert fn.return_type == expr.type return Call(fn, combined_args, type = fn.return_type) else: expr.fn = fn expr.args = args return expr def transform_if_simple_expr(self, expr): if isinstance(expr, Expr): return self.transform_simple_expr(expr) else: return expr def transform_simple_expr(self, expr, name = None): if name is None: name = "temp" result = self.transform_expr(expr) if not self.is_simple(result): return self.assign_name(result, name) else: return result def transform_simple_exprs(self, args): return [self.transform_simple_expr(x) for x in args] def transform_Array(self, expr): expr.elts = tuple(self.transform_simple_exprs(expr.elts)) return expr def transform_Slice(self, expr): expr.start = self.transform_simple_expr(expr.start) expr.stop = self.transform_simple_expr(expr.stop) expr.step = self.transform_simple_expr(expr.step) return expr def transform_index_expr(self, expr): if expr.__class__ is Tuple: new_elts = [] for elt in expr.elts: new_elt = self.transform_expr(elt) if not self.is_simple(new_elt) and new_elt.type.__class__ is not SliceT: new_elt = self.temp(new_elt, "index_tuple_elt") new_elts.append(new_elt) expr.elts = tuple(new_elts) return expr else: return self.transform_expr(expr) def transform_Index(self, expr): expr.value = self.transform_expr(expr.value) expr.index = self.transform_index_expr(expr.index) if expr.value.__class__ is Array and expr.index.__class__ is Const: assert isinstance(expr.index.value, (int, long)) and \ len(expr.value.elts) > expr.index.value return expr.value.elts[expr.index.value] # take expressions like "a[i][j]" and turn them into "a[i,j]" if expr.value.__class__ is Index: base_array = expr.value.value if isinstance(base_array.type, ArrayT): base_index = expr.value.index if isinstance(base_index.type, TupleT): indices = self.tuple_elts(base_index) else: assert isinstance(base_index.type, ScalarT), \ "Unexpected index type %s : %s in %s" % (base_index, base_index.type, expr) indices = [base_index] if isinstance(expr.index.type, TupleT): indices = tuple(indices) + tuple(self.tuple_elts(expr.index)) else: assert isinstance(expr.index.type, ScalarT), \ "Unexpected index type %s : %s in %s" % (expr.index, expr.index.type, expr) indices = tuple(indices) + (expr.index,) expr = self.index(base_array, self.tuple(indices)) return self.transform_expr(expr) if expr.value.__class__ is not Var: expr.value = self.temp(expr.value, "array") return expr def transform_Struct(self, expr): new_args = self.transform_simple_exprs(expr.args) return syntax.Struct(new_args, type = expr.type) def transform_Select(self, expr): cond = self.transform_expr(expr.cond) trueval = self.transform_expr(expr.true_value) falseval = self.transform_expr(expr.false_value) if is_true(cond): return trueval elif is_false(cond): return falseval elif trueval == falseval: return trueval else: expr.cond = cond expr.false_value = falseval expr.true_value = trueval return expr def transform_PrimCall(self, expr): args = self.transform_simple_exprs(expr.args) prim = expr.prim if all_constants(args): return syntax.Const(value = prim.fn(*collect_constants(args)), type = expr.type) if len(args) == 1: x = args[0] if prim == prims.logical_not: if is_false(x): return true elif is_true(x): return false if len(args) == 2: x,y = args if prim == prims.add: if is_zero(x): return y elif is_zero(y): return x if y.__class__ is Const and y.value < 0: expr.prim = prims.subtract expr.args = (x, Const(value = -y.value, type = y.type)) return expr elif x.__class__ is Const and x.value < 0: expr.prim = prims.subtract expr.args = (y, Const(value = -x.value, type = x.type)) return expr elif prim == prims.subtract: if is_zero(y): return x elif is_zero(x) and y.__class__ is Var: stored = self.bindings.get(y.name) # 0 - (a * b) --> -a * b |or| a * -b if stored and stored.__class__ is PrimCall and stored.prim == prims.multiply: a,b = stored.args if a.__class__ is Const: expr.prim = prims.multiply neg_a = Const(value = -a.value, type = a.type) expr.args = [neg_a, b] return expr elif b.__class__ is Const: expr.prim = prims.multiply neg_b = Const(value = -b.value, type = b.type) expr.args = [a, neg_b] return expr elif prim == prims.multiply: if is_one(x): return y elif is_one(y): return x elif is_zero(x): return x elif is_zero(y): return y elif prim == prims.divide and is_one(y): return x elif prim == prims.power: if is_one(y): return self.cast(x, expr.type) elif is_zero(y): return one(expr.type) elif y.__class__ is Const: if y.value == 2: return self.cast(self.mul(x, x, "sqr"), expr.type) elif y.value == 1: return self.cast(x, expr.type) elif y.value == 0: return self.cast(one_i64, expr.type) elif y.value == 0.5 and isinstance(expr.type, FloatT): expr.prim = prims.sqrt expr.args = (self.cast(x, expr.type),) return expr elif prim == prims.logical_and: if is_true(x): return y elif is_true(y): return x elif is_false(x) or is_false(y): return false elif prim == prims.logical_or: if is_true(x) or is_true(y): return true elif is_false(x): return y elif is_false(y): return x expr.args = args return expr def transform_Map(self, expr): expr.args = self.transform_simple_exprs(expr.args) expr.fn = self.transform_expr(expr.fn) expr.axis = self.transform_if_expr(expr.axis) max_rank = max(self.rank(arg) for arg in expr.args) # if an axis is the Python value None, turn it into the IR expression for None if max_rank == 1 and self.is_none(expr.axis): expr.axis = zero_i64 elif expr.axis is None: expr.axis = none return expr def transform_OuterMap(self, expr): expr.args = self.transform_simple_exprs(expr.args) expr.fn = self.transform_expr(expr.fn) expr.axis = self.transform_if_expr(expr.axis) max_rank = max(self.rank(arg) for arg in expr.args) # if an axis is the Python value None, turn it into the IR expression for None if max_rank == 1 and self.is_none(expr.axis): expr.axis = zero_i64 elif expr.axis is None: expr.axis = none return expr def transform_shape(self, expr): if isinstance(expr, Tuple): expr.elts = tuple(self.transform_simple_exprs(expr.elts)) return expr else: return self.transform_simple_expr(expr) def transform_ParFor(self, stmt): stmt.bounds = self.transform_shape(stmt.bounds) stmt.fn = self.transform_expr(stmt.fn) return stmt def transform_Reduce(self, expr): expr.axis = self.transform_if_expr(expr.axis) expr.fn = self.transform_expr(expr.fn) expr.combine = self.transform_expr(expr.combine) expr.init = self.transform_if_simple_expr(expr.init) expr.args = self.transform_simple_exprs(expr.args) # if an axis is the Python value None, turn it into the IR expression for None max_rank = max(self.rank(arg) for arg in expr.args) if max_rank == 1 and self.is_none(expr.axis): expr.axis = zero_i64 elif expr.axis is None: expr.axis = none return expr def transform_Scan(self, expr): expr.axis = self.transform_if_expr(expr.axis) expr.fn = self.transform_expr(expr.fn) expr.combine = self.transform_expr(expr.combine) expr.emit = self.transform_expr(expr.emit) expr.init = self.transform_if_simple_expr(expr.init) expr.args = self.transform_simple_exprs(expr.args) max_rank = max(self.rank(arg) for arg in expr.args) if max_rank == 1 and self.is_none(expr.axis): expr.axis = zero_i64 elif expr.axis is None: expr.axis = none return expr def transform_IndexMap(self, expr): expr.fn = self.transform_expr(expr.fn) expr.shape = self.transform_shape(expr.shape) return expr def transform_IndexReduce(self, expr): expr.fn = self.transform_if_expr(expr.fn) expr.combine = self.transform_expr(expr.combine) expr.init = self.transform_if_simple_expr(expr.init) expr.shape = self.transform_shape(expr.shape) return expr def transform_IndexScan(self, expr): expr.fn = self.transform_if_expr(expr.fn) expr.combine = self.transform_expr(expr.combine) expr.emit = self.transform_if_expr(expr.emit) expr.init = self.transform_if_simple_expr(expr.init) expr.shape = self.transform_shape(expr.shape) return expr def transform_ConstArray(self, expr): expr.shape = self.transform_shape(expr.shape) expr.value = self.transform_simple_expr(expr.value) return expr def transform_ConstArrayLike(self, expr): expr.array = self.transform_simple_expr(expr.array) expr.value = self.transform_simple_expr(expr.value) def temp_in_block(self, expr, block, name = None): """ If we need a temporary variable not in the current top scope but in a particular block, then use this function. (this function also modifies the bindings dictionary) """ if name is None: name = "temp" var = self.fresh_var(expr.type, name) block.append(Assign(var, expr)) self.bindings[var.name] = expr return var def set_binding(self, name, value): assert value.__class__ is not Var or \ value.name != name, \ "Can't set name %s bound to itself" % name self.bindings[name] = value def bind_var(self, name, rhs): if rhs.__class__ is Var: old_val = self.bindings.get(rhs.name) if old_val and self.is_simple(old_val): self.set_binding(name, old_val) else: self.set_binding(name, rhs) else: self.set_binding(name, rhs) def bind(self, lhs, rhs): lhs_class = lhs.__class__ if lhs_class is Var: self.bind_var(lhs.name, rhs) elif lhs_class is Tuple and rhs.__class__ is Tuple: assert len(lhs.elts) == len(rhs.elts) for lhs_elt, rhs_elt in zip(lhs.elts, rhs.elts): self.bind(lhs_elt, rhs_elt) def transform_lhs_Index(self, lhs): lhs.index = self.transform_index_expr(lhs.index) if lhs.value.__class__ is Var: stored = self.bindings.get(lhs.value.name) if stored and stored.__class__ is Var: lhs.value = stored else: lhs.value = self.assign_name(lhs.value, "array") return lhs def transform_lhs_Attribute(self, lhs): # lhs.value = self.transform_expr(lhs.value) return lhs def transform_ExprStmt(self, stmt): """Don't run an expression unless it possibly has a side effect""" v = self.transform_expr(stmt.value) if self.immutable(v): return None else: stmt.value = v return stmt def transform_Assign(self, stmt): lhs = stmt.lhs rhs = self.transform_expr(stmt.rhs) lhs_class = lhs.__class__ rhs_class = rhs.__class__ if lhs_class is Var: if lhs.type.__class__ is NoneT and self.use_counts.get(lhs.name,0) == 0: return self.transform_stmt(ExprStmt(rhs)) elif self.immutable(rhs): self.bind_var(lhs.name, rhs) if rhs_class is not Var and rhs_class is not Const: self.available_expressions.setdefault(rhs, lhs) elif lhs_class is Tuple: self.bind(lhs, rhs) elif lhs_class is Index: if rhs_class is Index and \ lhs.value == rhs.value and \ lhs.index == rhs.index: # kill effect-free writes like x[i] = x[i] return None elif rhs_class is Var and \ lhs.value.__class__ is Var and \ lhs.value.name == rhs.name and \ lhs.index.type.__class__ is TupleT and \ all(elt_t == slice_none_t for elt_t in lhs.index.type.elt_types): # also kill x[:] = x return None else: lhs = self.transform_lhs_Index(lhs) # when assigning x[j] = [1,2,3] # just rewrite it as a sequence of element assignments # to avoid if lhs.type.__class__ is ArrayT and \ lhs.type.rank == 1 and \ rhs.__class__ is Array: lhs_slice = self.assign_name(lhs, "lhs_slice") for (elt_idx, elt) in enumerate(rhs.elts): lhs_idx = self.index(lhs_slice, const_int(elt_idx), temp = False) self.assign(lhs_idx, elt) return None elif not self.is_simple(rhs): rhs = self.assign_name(rhs) else: assert lhs_class is Attribute assert False, "Considering making attributes immutable" lhs = self.transform_lhs_Attribute(lhs) if rhs_class is Var and \ rhs.name in self.bindings and \ self.use_counts.get(rhs.name, 1) == 1: self.use_counts[rhs.name] = 0 rhs = self.bindings[rhs.name] stmt.lhs = lhs stmt.rhs = rhs return stmt def transform_block(self, stmts, keep_bindings = False): self.available_expressions.push() self.bindings.push() new_stmts = Transform.transform_block(self, stmts) self.available_expressions.pop() if not keep_bindings: self.bindings.pop() return new_stmts def enter_loop(self, phi_nodes): result = {} for (k, (left,right)) in phi_nodes.iteritems(): new_left = self.transform_expr(left) if new_left == right: self.set_binding(k, new_left) else: result[k] = (new_left, right) return result def transform_merge(self, phi_nodes, left_block, right_block): result = {} for (k, (left, right)) in phi_nodes.iteritems(): new_left = self.transform_expr(left) new_right = self.transform_expr(right) if not isinstance(new_left, (Const, Var)): new_left = self.temp_in_block(new_left, left_block) if not isinstance(new_right, (Const, Var)): new_right = self.temp_in_block(new_right, right_block) if new_left == new_right: # if both control flows yield the same value then # we don't actually need the phi-bound variable, we can just # replace the left value everywhere self.assign(Var(name= k, type = new_left.type), new_left) self.set_binding(k, new_left) else: result[k] = new_left, new_right return result def transform_If(self, stmt): stmt.true = self.transform_block(stmt.true, keep_bindings = True) stmt.false = self.transform_block(stmt.false, keep_bindings=True) stmt.merge = self.transform_merge(stmt.merge, left_block = stmt.true, right_block = stmt.false) self.bindings.pop() self.bindings.pop() stmt.cond = self.transform_simple_expr(stmt.cond, "cond") if len(stmt.true) == 0 and len(stmt.false) == 0 and len(stmt.merge) <= 2: for (lhs_name, (true_expr, false_expr)) in stmt.merge.items(): lhs_type = self.lookup_type(lhs_name) lhs_var = Var(name = lhs_name, type = lhs_type) assert true_expr.type == false_expr.type, \ "Unexpcted type mismatch: %s != %s" % (true_expr.type, false_expr.type) rhs = Select(stmt.cond, true_expr, false_expr, type = true_expr.type) self.bind_var(lhs_name, rhs) self.assign(lhs_var, rhs) return None return stmt def transform_loop_condition(self, expr, outer_block, loop_body, merge): """Normalize loop conditions so they are just simple variables""" if self.is_simple(expr): return self.transform_expr(expr) else: loop_carried_vars = [name for name in collect_var_names(expr) if name in merge] if len(loop_carried_vars) == 0: return expr left_values = [merge[name][0] for name in loop_carried_vars] right_values = [merge[name][1] for name in loop_carried_vars] left_cond = subst.subst_expr(expr, dict(zip(loop_carried_vars, left_values))) if not self.is_simple(left_cond): left_cond = self.temp_in_block(left_cond, outer_block, name = "cond") right_cond = subst.subst_expr(expr, dict(zip(loop_carried_vars, right_values))) if not self.is_simple(right_cond): right_cond = self.temp_in_block(right_cond, loop_body, name = "cond") cond_var = self.fresh_var(left_cond.type, "cond") merge[cond_var.name] = (left_cond, right_cond) return cond_var def transform_While(self, stmt): merge = self.enter_loop(stmt.merge) stmt.body = self.transform_block(stmt.body) stmt.merge = self.transform_merge(merge, left_block = self.blocks.current(), right_block = stmt.body) stmt.cond = \ self.transform_loop_condition(stmt.cond, outer_block = self.blocks.current(), loop_body = stmt.body, merge = stmt.merge) return stmt def transform_ForLoop(self, stmt): merge = self.enter_loop(stmt.merge) stmt.body = self.transform_block(stmt.body) stmt.merge = self.transform_merge(merge, left_block = self.blocks.current(), right_block = stmt.body) stmt.start = self.transform_simple_expr(stmt.start, 'start') stmt.stop = self.transform_simple_expr(stmt.stop, 'stop') if self.is_none(stmt.step): stmt.step = one(stmt.start.type) else: stmt.step = self.transform_simple_expr(stmt.step, 'step') # if a loop is only going to run for one iteration, might as well get rid of # it if stmt.start.__class__ is Const and \ stmt.stop.__class__ is Const and \ stmt.step.__class__ is Const: if stmt.start.value >= stmt.stop.value: for (var_name, (input_value, _)) in stmt.merge.iteritems(): var = Var(var_name, input_value.type) self.blocks.append(Assign(var, input_value)) return None elif stmt.start.value + stmt.step.value >= stmt.stop.value: for (var_name, (input_value, _)) in stmt.merge.iteritems(): var = Var(var_name, input_value.type) self.blocks.append(Assign(var, input_value)) self.assign(stmt.var, stmt.start) self.blocks.top().extend(stmt.body) return None return stmt def transform_Return(self, stmt): new_value = self.transform_expr(stmt.value) if new_value != stmt.value: stmt.value = new_value return stmt