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)
