def make_yask_ast(expr, yc_soln, mapper=None): def nary2binary(args, op): r = make_yask_ast(args[0], yc_soln, mapper) return r if len(args) == 1 else op(r, nary2binary(args[1:], op)) if mapper is None: mapper = {} if expr.is_Integer: return nfac.new_const_number_node(int(expr)) elif expr.is_Float: return nfac.new_const_number_node(float(expr)) elif expr.is_Rational: a, b = expr.as_numer_denom() return nfac.new_const_number_node(float(a) / float(b)) elif expr.is_Symbol: function = expr.function if function.is_Constant: # Create a YASK var if it's the first time we encounter the embedded Function if function not in mapper: mapper[function] = yc_soln.new_var(function.name, []) # Allow number of time-steps to be set in YASK kernel. mapper[function].set_dynamic_step_alloc(True) return mapper[function].new_var_point([]) elif function.is_Dimension: if expr.is_Time: return nfac.new_step_index(expr.name) elif expr.is_Space: # `expr.root` instead of `expr` because YASK wants the SubDimension # information to be provided as if-conditions, and this is handled # a-posteriori directly by `yaskit` return nfac.new_domain_index(expr.root.name) else: return nfac.new_misc_index(expr.name) else: # E.g., A DSE-generated temporary, which must have already been # encountered as a LHS of a previous expression assert function in mapper return mapper[function] elif expr.is_Indexed: function = expr.function # Create a YASK var if it's the first time we encounter the embedded Function if function not in mapper: dimensions = [ make_yask_ast(i.root, yc_soln, mapper) for i in function.indices ] mapper[function] = yc_soln.new_var(function.name, dimensions) # Allow number of time-steps to be set in YASK kernel. mapper[function].set_dynamic_step_alloc(True) # We also get to know some relevant Dimension-related symbols # For example, the min point of the `x` Dimension, `x_m`, should # be mapped to YASK's `FIRST(x)` for d in function.indices: node = nfac.new_domain_index(d.name) mapper[d.symbolic_min] = nfac.new_first_domain_index(node) mapper[d.symbolic_max] = nfac.new_last_domain_index(node) indices = [make_yask_ast(i, yc_soln, mapper) for i in expr.indices] return mapper[function].new_var_point(indices) elif expr.is_Add: return nary2binary(expr.args, nfac.new_add_node) elif expr.is_Mul: return nary2binary(expr.args, nfac.new_multiply_node) elif expr.is_Pow: base, exp = expr.as_base_exp() if not exp.is_integer: raise NotImplementedError("Non-integer powers unsupported in " "Devito-YASK translation") if int(exp) < 0: num, den = expr.as_numer_denom() return nfac.new_divide_node(make_yask_ast(num, yc_soln, mapper), make_yask_ast(den, yc_soln, mapper)) elif int(exp) >= 1: return nary2binary([base] * exp, nfac.new_multiply_node) else: warning("0-power found in Devito-YASK translation? setting to 1") return nfac.new_const_number_node(1) elif isinstance(expr, IntDiv): return nfac.new_divide_node(make_yask_ast(expr.lhs, yc_soln, mapper), make_yask_ast(expr.rhs, yc_soln, mapper)) elif expr.is_Equality: if expr.lhs.is_Symbol: function = expr.lhs.function # The IETs are always in SSA form, so the only situation in # which `function` may already appear in `mapper` is when we've # already processed it as part of a different set of # boundary conditions. For example consider `expr = a[x]*2`: # first time, expr executed iff `x == FIRST_INDEX(x) + 7` # second time, expr executed iff `x == FIRST_INDEX(x) + 6` if function not in mapper: mapper[function] = make_yask_ast(expr.rhs, yc_soln, mapper) else: return nfac.new_equation_node( *[make_yask_ast(i, yc_soln, mapper) for i in expr.args]) else: raise NotImplementedError("Missing handler in Devito-YASK translation")
def yaskit(trees, yc_soln): """ Populate a YASK compiler solution with the Expressions found in an IET. The necessary YASK vars are instantiated. Parameters ---------- trees : list of IterationTree One or more Iteration trees within which the convertible Expressions are searched. yc_soln The YASK compiler solution to be populated. """ # It's up to Devito to organize the equations into a flow graph yc_soln.set_dependency_checker_enabled(False) mapper = {} processed = [] for tree in trees: # All expressions within `tree` expressions = [i.expr for i in FindNodes(Expression).visit(tree.inner)] # Attach conditional expression for sub-domains conditions = [(i, []) for i in expressions] for i in tree: if not i.dim.is_Sub: continue # Can we express both Iteration extremes as # `FIRST(i.dim) + integer` OR `LAST(i.dim) + integer` ? # If not, one of the following lines will throw a TypeError exception lvalue, lextreme, _ = i.dim._offset_left rvalue, rextreme, _ = i.dim._offset_right if i.is_Parallel: # At this point, no issues are expected -- we should just be able to # build the YASK conditions under which to execute this parallel Iteration ydim = nfac.new_domain_index(i.dim.parent.name) # Handle lower extreme if lextreme == i.dim.parent.symbolic_min: node = nfac.new_first_domain_index(ydim) else: node = nfac.new_last_domain_index(ydim) expr = nfac.new_add_node(node, nfac.new_const_number_node(lvalue)) for _, v in conditions: v.append(nfac.new_not_less_than_node(ydim, expr)) # Handle upper extreme if rextreme == i.dim.parent.symbolic_min: node = nfac.new_first_domain_index(ydim) else: node = nfac.new_last_domain_index(ydim) expr = nfac.new_add_node(node, nfac.new_const_number_node(rvalue)) for _, v in conditions: v.append(nfac.new_not_greater_than_node(ydim, expr)) elif i.is_Sequential: # For sequential Iterations, the extent *must* be statically known, # otherwise we don't know how to handle this try: int(i.dim._thickness_map.get(i.size())) except TypeError: raise NotImplementedError( "Found sequential Iteration with " "statically unknown extent") assert lextreme == rextreme # A corollary of getting up to this point n = lextreme ydim = nfac.new_domain_index(i.dim.parent.name) if n == i.dim.parent.symbolic_min: node = nfac.new_first_domain_index(ydim) else: node = nfac.new_last_domain_index(ydim) if i.direction is Backward: _range = range(rvalue, lvalue - 1, -1) else: _range = range(lvalue, rvalue + 1) unwound = [] for e, v in conditions: for r in _range: expr = nfac.new_add_node(node, nfac.new_const_number_node(r)) unwound.append( (e, v + [nfac.new_equals_node(ydim, expr)])) conditions = unwound # Build the YASK equations as well as all necessary vars for k, v in conditions: yask_expr = make_yask_ast(k, yc_soln, mapper) if yask_expr is not None: processed.append(yask_expr) # Is there a sub-domain to attach ? if v: condition = v.pop(0) for i in v: condition = nfac.new_and_node(condition, i) yask_expr.set_cond(condition) # Add flow dependences to the offloaded equations # TODO: This can be improved by spotting supergroups ? for to, frm in zip(processed, processed[1:]): yc_soln.add_flow_dependency(frm, to) # Have we built any local vars (i.e., DSE-produced tensor temporaries)? # If so, eventually these will be mapped to YASK scratch vars local_vars = [i for i in mapper if i.is_Array] return local_vars
def make_yask_ast(expr, yc_soln, mapper=None): def nary2binary(args, op): r = make_yask_ast(args[0], yc_soln, mapper) return r if len(args) == 1 else op(r, nary2binary(args[1:], op)) if mapper is None: mapper = {} if expr.is_Integer: return nfac.new_const_number_node(int(expr)) elif expr.is_Float: return nfac.new_const_number_node(float(expr)) elif expr.is_Rational: a, b = expr.as_numer_denom() return nfac.new_const_number_node(float(a)/float(b)) elif expr.is_Symbol: function = expr.function if function.is_Constant: # Create a YASK grid if it's the first time we encounter the embedded Function if function not in mapper: mapper[function] = yc_soln.new_grid(function.name, []) # Allow number of time-steps to be set in YASK kernel. mapper[function].set_dynamic_step_alloc(True) return mapper[function].new_grid_point([]) elif function.is_Dimension: if expr.is_Time: return nfac.new_step_index(expr.name) elif expr.is_Space: # `expr.root` instead of `expr` because YASK wants the SubDimension # information to be provided as if-conditions, and this is handled # a-posteriori directly by `yaskit` return nfac.new_domain_index(expr.root.name) else: return nfac.new_misc_index(expr.name) else: # E.g., A DSE-generated temporary, which must have already been # encountered as a LHS of a previous expression assert function in mapper return mapper[function] elif expr.is_Indexed: function = expr.function # Create a YASK grid if it's the first time we encounter the embedded Function if function not in mapper: dimensions = [make_yask_ast(i.root, yc_soln, mapper) for i in function.indices] mapper[function] = yc_soln.new_grid(function.name, dimensions) # Allow number of time-steps to be set in YASK kernel. mapper[function].set_dynamic_step_alloc(True) # We also get to know some relevant Dimension-related symbols # For example, the min point of the `x` Dimension, `x_m`, should # be mapped to YASK's `FIRST(x)` for d in function.indices: node = nfac.new_domain_index(d.name) mapper[d.symbolic_min] = nfac.new_first_domain_index(node) mapper[d.symbolic_max] = nfac.new_last_domain_index(node) indices = [make_yask_ast(i, yc_soln, mapper) for i in expr.indices] return mapper[function].new_grid_point(indices) elif expr.is_Add: return nary2binary(expr.args, nfac.new_add_node) elif expr.is_Mul: return nary2binary(expr.args, nfac.new_multiply_node) elif expr.is_Pow: base, exp = expr.as_base_exp() if not exp.is_integer: raise NotImplementedError("Non-integer powers unsupported in " "Devito-YASK translation") if int(exp) < 0: num, den = expr.as_numer_denom() return nfac.new_divide_node(make_yask_ast(num, yc_soln, mapper), make_yask_ast(den, yc_soln, mapper)) elif int(exp) >= 1: return nary2binary([base] * exp, nfac.new_multiply_node) else: warning("0-power found in Devito-YASK translation? setting to 1") return nfac.new_const_number_node(1) elif isinstance(expr, IntDiv): return nfac.new_divide_node(make_yask_ast(expr.lhs, yc_soln, mapper), make_yask_ast(expr.rhs, yc_soln, mapper)) elif expr.is_Equality: if expr.lhs.is_Symbol: function = expr.lhs.function # The IETs are always in SSA form, so the only situation in # which `function` may already appear in `mapper` is when we've # already processed it as part of a different set of # boundary conditions. For example consider `expr = a[x]*2`: # first time, expr executed iff `x == FIRST_INDEX(x) + 7` # second time, expr executed iff `x == FIRST_INDEX(x) + 6` if function not in mapper: mapper[function] = make_yask_ast(expr.rhs, yc_soln, mapper) else: return nfac.new_equation_node(*[make_yask_ast(i, yc_soln, mapper) for i in expr.args]) else: raise NotImplementedError("Missing handler in Devito-YASK translation")
def yaskit(trees, yc_soln): """ Populate a YASK compiler solution with the :class:`Expression`s found in an IET. The necessary YASK grids are instantiated. :param trees: A sequence of offloadable :class:`IterationTree`s, in which the Expressions are searched. :param yc_soln: The YASK compiler solution to be populated. """ # Track all created YASK grids mapper = {} # It's up to Devito to organize the equations into a flow graph yc_soln.set_dependency_checker_enabled(False) processed = [] for tree in trees: # All expressions within `tree` expressions = [i.expr for i in FindNodes(Expression).visit(tree.inner)] # Attach conditional expression for sub-domains conditions = [(i, []) for i in expressions] for i in tree: if not i.dim.is_Sub: continue # Can we express both Iteration extremes as # `FIRST(i.dim) + integer` OR `LAST(i.dim) + integer` ? # If not, one of the following lines will throw a TypeError exception lower_ofs, lower_sym = i.dim.offset_left() upper_ofs, upper_sym = i.dim.offset_right() if i.is_Parallel: # At this point, no issues are expected -- we should just be able to # build the YASK conditions under which to execute this parallel Iteration ydim = nfac.new_domain_index(i.dim.parent.name) # Handle lower extreme if lower_sym == i.dim.parent.symbolic_start: node = nfac.new_first_domain_index(ydim) else: node = nfac.new_last_domain_index(ydim) expr = nfac.new_add_node(node, nfac.new_const_number_node(lower_ofs)) for _, v in conditions: v.append(nfac.new_not_less_than_node(ydim, expr)) # Handle upper extreme if upper_sym == i.dim.parent.symbolic_start: node = nfac.new_first_domain_index(ydim) else: node = nfac.new_last_domain_index(ydim) expr = nfac.new_add_node(node, nfac.new_const_number_node(upper_ofs)) for _, v in conditions: v.append(nfac.new_not_greater_than_node(ydim, expr)) elif i.is_Sequential: # For sequential Iterations, the extent *must* be statically known, # otherwise we don't know how to handle this try: int(i.extent()) except TypeError: raise NotImplementedError( "Found sequential Iteration with " "statically unknown extent") assert lower_sym == upper_sym # A corollary of getting up to this point n = lower_sym ydim = nfac.new_domain_index(i.dim.parent.name) if n == i.dim.parent.symbolic_start: node = nfac.new_first_domain_index(ydim) else: node = nfac.new_last_domain_index(ydim) if i.direction is Backward: _range = range(upper_ofs, lower_ofs - 1, -1) else: _range = range(lower_ofs, upper_ofs + 1) unwound = [] for e, v in conditions: for r in _range: expr = nfac.new_add_node(node, nfac.new_const_number_node(r)) unwound.append( (e, v + [nfac.new_equals_node(ydim, expr)])) conditions = unwound # Build the YASK equations as well as all necessary grids for k, v in conditions: yask_expr = make_yask_ast(k, yc_soln, mapper) if yask_expr is not None: processed.append(yask_expr) # Is there a sub-domain to attach ? if v: condition = v.pop(0) for i in v: condition = nfac.new_and_node(condition, i) yask_expr.set_cond(condition) # Add flow dependences to the offloaded equations # TODO: This can be improved by spotting supergroups ? for to, frm in zip(processed, processed[1:]): yc_soln.add_flow_dependency(frm, to) return mapper
def yaskit(trees, yc_soln): """ Populate a YASK compiler solution with the :class:`Expression`s found in an IET. The necessary YASK grids are instantiated. Parameters ---------- trees : list of IterationTree One or more Iteration trees within which the convertible Expressions are searched. yc_soln The YASK compiler solution to be populated. """ # It's up to Devito to organize the equations into a flow graph yc_soln.set_dependency_checker_enabled(False) mapper = {} processed = [] for tree in trees: # All expressions within `tree` expressions = [i.expr for i in FindNodes(Expression).visit(tree.inner)] # Attach conditional expression for sub-domains conditions = [(i, []) for i in expressions] for i in tree: if not i.dim.is_Sub: continue # Can we express both Iteration extremes as # `FIRST(i.dim) + integer` OR `LAST(i.dim) + integer` ? # If not, one of the following lines will throw a TypeError exception lower_ofs, lower_sym = i.dim._offset_left() upper_ofs, upper_sym = i.dim._offset_right() if i.is_Parallel: # At this point, no issues are expected -- we should just be able to # build the YASK conditions under which to execute this parallel Iteration ydim = nfac.new_domain_index(i.dim.parent.name) # Handle lower extreme if lower_sym == i.dim.parent.symbolic_min: node = nfac.new_first_domain_index(ydim) else: node = nfac.new_last_domain_index(ydim) expr = nfac.new_add_node(node, nfac.new_const_number_node(lower_ofs)) for _, v in conditions: v.append(nfac.new_not_less_than_node(ydim, expr)) # Handle upper extreme if upper_sym == i.dim.parent.symbolic_min: node = nfac.new_first_domain_index(ydim) else: node = nfac.new_last_domain_index(ydim) expr = nfac.new_add_node(node, nfac.new_const_number_node(upper_ofs)) for _, v in conditions: v.append(nfac.new_not_greater_than_node(ydim, expr)) elif i.is_Sequential: # For sequential Iterations, the extent *must* be statically known, # otherwise we don't know how to handle this try: int(i.dim._thickness_map.get(i.size())) except TypeError: raise NotImplementedError("Found sequential Iteration with " "statically unknown extent") assert lower_sym == upper_sym # A corollary of getting up to this point n = lower_sym ydim = nfac.new_domain_index(i.dim.parent.name) if n == i.dim.parent.symbolic_min: node = nfac.new_first_domain_index(ydim) else: node = nfac.new_last_domain_index(ydim) if i.direction is Backward: _range = range(upper_ofs, lower_ofs - 1, -1) else: _range = range(lower_ofs, upper_ofs + 1) unwound = [] for e, v in conditions: for r in _range: expr = nfac.new_add_node(node, nfac.new_const_number_node(r)) unwound.append((e, v + [nfac.new_equals_node(ydim, expr)])) conditions = unwound # Build the YASK equations as well as all necessary grids for k, v in conditions: yask_expr = make_yask_ast(k, yc_soln, mapper) if yask_expr is not None: processed.append(yask_expr) # Is there a sub-domain to attach ? if v: condition = v.pop(0) for i in v: condition = nfac.new_and_node(condition, i) yask_expr.set_cond(condition) # Add flow dependences to the offloaded equations # TODO: This can be improved by spotting supergroups ? for to, frm in zip(processed, processed[1:]): yc_soln.add_flow_dependency(frm, to) # Have we built any local grids (i.e., DSE-produced tensor temporaries)? # If so, eventually these will be mapped to YASK scratch grids local_grids = [i for i in mapper if i.is_Array] return local_grids