def statement_evaluate(leaf, ctx): expr = leaf.expression if isinstance(expr, gem.ListTensor): ops = [] var, index = ctx.pymbolic_variable_and_destruct(expr) for multiindex, value in numpy.ndenumerate(expr.array): ops.append( lp.Assignment(p.Subscript(var, index + multiindex), expression(value, ctx), within_inames=ctx.active_inames())) return ops elif isinstance(expr, gem.Constant): return [] elif isinstance(expr, gem.ComponentTensor): idx = ctx.gem_to_pym_multiindex(expr.multiindex) var, sub_idx = ctx.pymbolic_variable_and_destruct(expr) lhs = p.Subscript(var, idx + sub_idx) with active_indices(dict(zip(expr.multiindex, idx)), ctx) as ctx_active: return [ lp.Assignment(lhs, expression(expr.children[0], ctx_active), within_inames=ctx_active.active_inames()) ] elif isinstance(expr, gem.Inverse): idx = ctx.pymbolic_multiindex(expr.shape) var = ctx.pymbolic_variable(expr) lhs = (SubArrayRef(idx, p.Subscript(var, idx)), ) idx_reads = ctx.pymbolic_multiindex(expr.children[0].shape) var_reads = ctx.pymbolic_variable(expr.children[0]) reads = (SubArrayRef(idx_reads, p.Subscript(var_reads, idx_reads)), ) rhs = p.Call(p.Variable("inverse"), reads) return [ lp.CallInstruction(lhs, rhs, within_inames=ctx.active_inames()) ] elif isinstance(expr, gem.Solve): idx = ctx.pymbolic_multiindex(expr.shape) var = ctx.pymbolic_variable(expr) lhs = (SubArrayRef(idx, p.Subscript(var, idx)), ) reads = [] for child in expr.children: idx_reads = ctx.pymbolic_multiindex(child.shape) var_reads = ctx.pymbolic_variable(child) reads.append( SubArrayRef(idx_reads, p.Subscript(var_reads, idx_reads))) rhs = p.Call(p.Variable("solve"), tuple(reads)) return [ lp.CallInstruction(lhs, rhs, within_inames=ctx.active_inames()) ] else: return [ lp.Assignment(ctx.pymbolic_variable(expr), expression(expr, ctx, top=True), within_inames=ctx.active_inames()) ]
def slate_call(self, kernel): # Slate kernel call call = pym.Call(pym.Variable(kernel.name), tuple()) output_var = pym.Variable(kernel.args[0].name) slate_kernel_call_output = self.generate_lhs(self.expression, output_var) insn = loopy.CallInstruction((slate_kernel_call_output,), call, id="slate_kernel_call") return [insn]
def generate_tsfc_calls(self, terminal, loopy_tensor): """A setup method to initialize all the local assembly kernels generated by TSFC. This function also collects any information regarding orientations and extra include directories. """ cxt_kernels = self.tsfc_cxt_kernels(terminal) for cxt_kernel in cxt_kernels: for tsfc_kernel in cxt_kernel.tsfc_kernels: integral_type = cxt_kernel.original_integral_type slate_tensor = cxt_kernel.tensor mesh = slate_tensor.ufl_domain() kinfo = tsfc_kernel.kinfo reads = [] inames_dep = [] if integral_type not in self.supported_integral_types: raise ValueError("Integral type '%s' not recognized" % integral_type) # Prepare lhs and args for call to tsfc kernel output_var = pym.Variable(loopy_tensor.name) reads.append(output_var) output = self.generate_lhs(slate_tensor, output_var) kernel_data = self.collect_tsfc_kernel_data( mesh, cxt_kernel.coefficients, self.bag.coefficients, kinfo) reads.extend(self.loopify_tsfc_kernel_data(kernel_data)) # Generate predicates for different integral types if self.is_integral_type(integral_type, "cell_integral"): predicates = None if kinfo.subdomain_id != "otherwise": raise NotImplementedError( "No subdomain markers for cells yet") elif self.is_integral_type(integral_type, "facet_integral"): predicates, fidx, facet_arg = self.facet_integral_predicates( mesh, integral_type, kinfo) reads.append(facet_arg) inames_dep.append(fidx[0].name) elif self.is_integral_type(integral_type, "layer_integral"): predicates = self.layer_integral_predicates( slate_tensor, integral_type) else: raise ValueError( "Unhandled integral type {}".format(integral_type)) # TSFC kernel call key = self.bag.call_name_generator(integral_type) call = pym.Call(pym.Variable(kinfo.kernel.name), tuple(reads)) insn = loopy.CallInstruction( (output, ), call, within_inames=frozenset(inames_dep), predicates=predicates, id=key) yield insn, kinfo.kernel.code
def slate_call(self, kernel, temporaries): # Slate kernel call reads = [] for t in temporaries: shape = t.shape name = t.name idx = self.bag.index_creator(shape) reads.append(SubArrayRef(idx, pym.Subscript(pym.Variable(name), idx))) call = pym.Call(pym.Variable(kernel.name), tuple(reads)) output_var = pym.Variable(kernel.args[0].name) slate_kernel_call_output = self.generate_lhs(self.expression, output_var) insn = loopy.CallInstruction((slate_kernel_call_output,), call, id="slate_kernel_call") return insn
def test_call_with_no_returned_value(ctx_factory): import pymbolic.primitives as p ctx = ctx_factory() queue = cl.CommandQueue(ctx) knl = lp.make_kernel("{:}", [lp.CallInstruction((), p.Call(p.Variable("f"), ()))]) from library_for_test import no_ret_f_mangler, no_ret_f_preamble_gen knl = lp.register_function_manglers(knl, [no_ret_f_mangler]) knl = lp.register_preamble_generators(knl, [no_ret_f_preamble_gen]) evt, _ = knl(queue)
def statement_functioncall(expr, context): parameters = context.parameters free_indices = set(i.name for i in expr.free_indices) writes = [] reads = [] for access, child in zip(expr.access, expr.children): var = expression(child, parameters) if isinstance(var, pym.Subscript): # tensor argument indices = [] sweeping_indices = [] for index in var.index_tuple: indices.append(index) if isinstance(index, pym.Variable) and index.name in free_indices: sweeping_indices.append(index) arg = SubArrayRef(tuple(sweeping_indices), var) else: # scalar argument or constant arg = var if access is READ or (isinstance(child, Argument) and isinstance(child.dtype, OpaqueType)): reads.append(arg) else: writes.append(arg) within_inames = context.within_inames[expr] predicates = frozenset(context.conditions) id, depends_on = context.instruction_dependencies[expr] call = pym.Call(pym.Variable(expr.name), tuple(reads)) return loopy.CallInstruction(tuple(writes), call, within_inames=within_inames, predicates=predicates, id=id, depends_on=depends_on, depends_on_is_final=True)