def _slate2gem_mul(expr, self): A, B = map(self, expr.children) *i, k = tuple(make_indices(len(A.shape))) _, *j = tuple(make_indices(len(B.shape))) ABikj = Product(Indexed(A, tuple(i + [k])), Indexed(B, tuple([k] + j))) return ComponentTensor(IndexSum(ABikj, (k, )), tuple(i + j))
def gem_to_loopy(gem_expr, var2terminal, scalar_type): """ Method encapsulating stage 2. Converts the gem expression dag into imperoc first, and then further into loopy. :return slate_loopy: 2-tuple of loopy kernel for slate operations and loopy GlobalArg for the output variable. """ # Creation of return variables for outer loopy shape = gem_expr.shape if len(gem_expr.shape) != 0 else (1, ) idx = make_indices(len(shape)) indexed_gem_expr = gem.Indexed(gem_expr, idx) args = ([ loopy.GlobalArg("output", shape=shape, dtype=scalar_type, is_output=True, is_input=True) ] + [ loopy.GlobalArg(var.name, shape=var.shape, dtype=scalar_type) for var in var2terminal.keys() ]) ret_vars = [gem.Indexed(gem.Variable("output", shape), idx)] preprocessed_gem_expr = impero_utils.preprocess_gem([indexed_gem_expr]) # glue assignments to return variable assignments = list(zip(ret_vars, preprocessed_gem_expr)) # Part A: slate to impero_c impero_c = impero_utils.compile_gem(assignments, (), remove_zeros=False) # Part B: impero_c to loopy return generate_loopy(impero_c, args, scalar_type, "slate_loopy", []), args[0].copy()
def gem_to_loopy(gem_expr): """ Method encapsulating stage 2. Converts the gem expression dag into imperoc first, and then further into loopy. :return slate_loopy: loopy kernel for slate operations. """ # Creation of return variables for outer loopy shape = gem_expr.shape if len(gem_expr.shape) != 0 else (1,) idx = make_indices(len(shape)) indexed_gem_expr = gem.Indexed(gem_expr, idx) arg = [loopy.GlobalArg("output", shape=shape)] ret_vars = [gem.Indexed(gem.Variable("output", shape), idx)] preprocessed_gem_expr = impero_utils.preprocess_gem([indexed_gem_expr]) # glue assignments to return variable assignments = list(zip(ret_vars, preprocessed_gem_expr)) # Part A: slate to impero_c impero_c = impero_utils.compile_gem(assignments, (), remove_zeros=False) # Part B: impero_c to loopy return generate_loopy(impero_c, arg, parameters["form_compiler"]["scalar_type"], "slate_loopy", [])
def _slate2gem_add(expr, self): A, B = map(self, expr.children) indices = tuple(make_indices(len(A.shape))) return ComponentTensor(Sum(Indexed(A, indices), Indexed(B, indices)), indices)
def _slate2gem_negative(expr, self): child, = map(self, expr.children) indices = tuple(make_indices(len(child.shape))) return ComponentTensor(Product(Literal(-1), Indexed(child, indices)), indices)
def _slate2gem_transpose(expr, self): child, = map(self, expr.children) indices = tuple(make_indices(len(child.shape))) return ComponentTensor(Indexed(child, indices), tuple(indices[::-1]))
def _slate2gem_reciprocal(expr, self): child, = map(self, expr.children) indices = tuple(make_indices(len(child.shape))) return ComponentTensor(Division(Literal(1.), Indexed(child, indices)), indices)