def map_basic_index(self, expr: BasicIndex) -> IndexLambda: vng = UniqueNameGenerator() indices = [] in_ary = vng("in") bindings = {in_ary: self.rec(expr.array)} islice_idx = 0 for idx, axis_len in zip(expr.indices, expr.array.shape): if isinstance(idx, INT_CLASSES): if isinstance(axis_len, INT_CLASSES): indices.append(idx % axis_len) else: bnd_name = vng("in") bindings[bnd_name] = axis_len indices.append(idx % prim.Variable(bnd_name)) elif isinstance(idx, NormalizedSlice): indices.append(idx.start + idx.step * prim.Variable(f"_{islice_idx}")) islice_idx += 1 else: raise NotImplementedError return IndexLambda(expr=prim.Subscript(prim.Variable(in_ary), tuple(indices)), bindings=bindings, shape=expr.shape, dtype=expr.dtype, axes=expr.axes, tags=expr.tags, )
def _expression_mathfunction(expr, ctx): if expr.name.startswith('cyl_bessel_'): # Bessel functions if is_complex(ctx.scalar_type): raise NotImplementedError("Bessel functions for complex numbers: " "missing implementation") nu, arg = expr.children nu_ = expression(nu, ctx) arg_ = expression(arg, ctx) # Modified Bessel functions (C++ only) # # These mappings work for FEniCS only, and fail with Firedrake # since no Boost available. if expr.name in {'cyl_bessel_i', 'cyl_bessel_k'}: name = 'boost::math::' + expr.name return p.Variable(name)(nu_, arg_) else: # cyl_bessel_{jy} -> {jy} name = expr.name[-1:] if nu == gem.Zero(): return p.Variable(f"{name}0")(arg_) elif nu == gem.one: return p.Variable(f"{name}1")(arg_) else: return p.Variable(f"{name}n")(nu_, arg_) else: if expr.name == "ln": name = "log" else: name = expr.name # Not all mathfunctions apply to complex numbers, but this # will be picked up in loopy. This way we allow erf(real(...)) # in complex mode (say). return p.Variable(name)(*(expression(c, ctx) for c in expr.children))
def _test_to_pymbolic(mapper, sym, use_symengine): x, y = sym.symbols("x,y") assert mapper(sym.Rational(3, 4)) == prim.Quotient(3, 4) assert mapper(sym.Integer(6)) == 6 if not use_symengine: assert mapper(sym.Subs(x**2, (x,), (y,))) == \ prim.Substitution(x_**2, ("x",), (y_,)) deriv = sym.Derivative(x**2, x) assert mapper(deriv) == prim.Derivative(x_**2, ("x", )) else: assert mapper(sym.Subs(x**2, (x,), (y,))) == \ y_**2 deriv = sym.Derivative(x**2, x) assert mapper(deriv) == 2 * x_ # functions assert mapper(sym.Function("f")(x)) == prim.Variable("f")(x_) assert mapper(sym.exp(x)) == prim.Variable("exp")(x_) # indexed accesses if not use_symengine: i, j = sym.symbols("i,j") assert mapper(sym.Indexed(x, i, j)) == x_[i_, j_] # constants import math # FIXME: Why isn't this exact? assert abs(mapper(sym.pi) - math.pi) < 1e-14 assert abs(mapper(sym.E) - math.e) < 1e-14 assert mapper(sym.I) == 1j
def facet_integral_predicates(self, mesh, integral_type, kinfo): self.bag.needs_cell_facets = True # Number of recerence cell facets if mesh.cell_set._extruded: self.num_facets = mesh._base_mesh.ufl_cell().num_facets() else: self.num_facets = mesh.ufl_cell().num_facets() # Index for loop over cell faces of reference cell fidx = self.bag.index_creator((self.num_facets,)) # Cell is interior or exterior select = 1 if integral_type.startswith("interior_facet") else 0 i = self.bag.index_creator((1,)) predicates = [pym.Comparison(pym.Subscript(pym.Variable(self.cell_facets_arg), (fidx[0], 0)), "==", select)] # TODO subdomain boundary integrals, this does the wrong thing for integrals like f*ds + g*ds(1) # "otherwise" is treated incorrectly as "everywhere" # However, this replicates an existing slate bug. if kinfo.subdomain_id != "otherwise": predicates.append(pym.Comparison(pym.Subscript(pym.Variable(self.cell_facets_arg), (fidx[0], 1)), "==", kinfo.subdomain_id)) # Additional facet array argument to be fed into tsfc loopy kernel subscript = pym.Subscript(pym.Variable(self.local_facet_array_arg), (pym.Sum((i[0], fidx[0])))) facet_arg = SubArrayRef(i, subscript) return predicates, fidx, facet_arg
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 test_parser(): from pymbolic import parse parse("(2*a[1]*b[1]+2*a[0]*b[0])*(hankel_1(-1,sqrt(a[1]**2+a[0]**2)*k) " "-hankel_1(1,sqrt(a[1]**2+a[0]**2)*k))*k /(4*sqrt(a[1]**2+a[0]**2)) " "+hankel_1(0,sqrt(a[1]**2+a[0]**2)*k)") print repr(parse("d4knl0")) print repr(parse("0.")) print repr(parse("0.e1")) print repr(parse("0.e1")) print repr(parse("a >= 1")) print repr(parse("a <= 1")) print repr(parse(":")) print repr(parse("1:")) print repr(parse(":2")) print repr(parse("1:2")) print repr(parse("::")) print repr(parse("1::")) print repr(parse(":1:")) print repr(parse("::1")) print repr(parse("3::1")) print repr(parse(":5:1")) print repr(parse("3:5:1")) print parse("3::1") assert parse("e1") == prim.Variable("e1") assert parse("d1") == prim.Variable("d1")
def map_power(self, expr, type_context): tgt_dtype = self.infer_type(expr) base_dtype = self.infer_type(expr.base) exponent_dtype = self.infer_type(expr.exponent) if not self.allow_complex or (not tgt_dtype.is_complex()): return super().map_power(expr, type_context) if expr.exponent in [2, 3, 4]: value = expr.base for _i in range(expr.exponent - 1): value = value * expr.base return self.rec(value, type_context) else: b_complex = base_dtype.is_complex() e_complex = exponent_dtype.is_complex() if b_complex and not e_complex: return p.Variable( "%s_powr" % self.complex_type_name(tgt_dtype))(self.rec( expr.base, type_context, tgt_dtype), self.rec(expr.exponent, type_context)) else: return p.Variable( "%s_pow" % self.complex_type_name(tgt_dtype))( self.rec(expr.base, type_context, tgt_dtype), self.rec(expr.exponent, type_context, tgt_dtype))
def initialise_terminals(self, var2terminal, coefficients): """ Initilisation of the variables in which coefficients and the Tensors coming from TSFC are saved. :arg var2terminal: dictionary that maps Slate Tensors to gem Variables """ tensor2temp = OrderedDict() inits = [] for gem_tensor, slate_tensor in var2terminal.items(): assert slate_tensor.terminal, "Only terminal tensors need to be initialised in Slate kernels." (_, dtype), = assign_dtypes([gem_tensor], self.tsfc_parameters["scalar_type"]) loopy_tensor = loopy.TemporaryVariable( gem_tensor.name, dtype=dtype, shape=gem_tensor.shape, address_space=loopy.AddressSpace.LOCAL) tensor2temp[slate_tensor] = loopy_tensor if not slate_tensor.assembled: indices = self.bag.index_creator(self.shape(slate_tensor)) inames = {var.name for var in indices} var = pym.Subscript(pym.Variable(loopy_tensor.name), indices) inits.append( loopy.Assignment(var, "0.", id="init%d" % len(inits), within_inames=frozenset(inames))) else: f = slate_tensor.form if isinstance( slate_tensor.form, tuple) else (slate_tensor.form, ) coeff = tuple(coefficients[c] for c in f) offset = 0 ismixed = tuple( (type(c.ufl_element()) == MixedElement) for c in f) names = [] for (im, c) in zip(ismixed, coeff): names += [name for (name, ext) in c.values()] if im else [c[0]] # Mixed coefficients come as seperate parameter (one per space) for i, shp in enumerate(*slate_tensor.shapes.values()): indices = self.bag.index_creator((shp, )) inames = {var.name for var in indices} offset_index = (pym.Sum((offset, indices[0])), ) name = names[i] if ismixed else names var = pym.Subscript(pym.Variable(loopy_tensor.name), offset_index) c = pym.Subscript(pym.Variable(name), indices) inits.append( loopy.Assignment(var, c, id="init%d" % len(inits), within_inames=frozenset(inames))) offset += shp return inits, tensor2temp
def map_non_contiguous_advanced_index(self, expr: AdvancedIndexInNoncontiguousAxes ) -> IndexLambda: from pytato.utils import (get_shape_after_broadcasting, get_indexing_expression) i_adv_indices = tuple(i for i, idx_expr in enumerate(expr.indices) if isinstance(idx_expr, (Array, INT_CLASSES))) adv_idx_shape = get_shape_after_broadcasting([expr.indices[i_idx] for i_idx in i_adv_indices]) vng = UniqueNameGenerator() indices = [] in_ary = vng("in") bindings = {in_ary: self.rec(expr.array)} islice_idx = len(adv_idx_shape) for idx, axis_len in zip(expr.indices, expr.array.shape): if isinstance(idx, INT_CLASSES): if isinstance(axis_len, INT_CLASSES): indices.append(idx % axis_len) else: bnd_name = vng("in") bindings[bnd_name] = self.rec(axis_len) indices.append(idx % prim.Variable(bnd_name)) elif isinstance(idx, NormalizedSlice): indices.append(idx.start + idx.step * prim.Variable(f"_{islice_idx}")) islice_idx += 1 elif isinstance(idx, Array): if isinstance(axis_len, INT_CLASSES): bnd_name = vng("in") bindings[bnd_name] = self.rec(idx) indirect_idx_expr = prim.Subscript(prim.Variable(bnd_name), get_indexing_expression( idx.shape, adv_idx_shape)) if not idx.tags_of_type(AssumeNonNegative): indirect_idx_expr = indirect_idx_expr % axis_len indices.append(indirect_idx_expr) else: raise NotImplementedError("Advanced indexing over" " parametric axis lengths.") else: raise NotImplementedError(f"Indices of type {type(idx)}.") return IndexLambda(expr=prim.Subscript(prim.Variable(in_ary), tuple(indices)), bindings=bindings, shape=expr.shape, dtype=expr.dtype, axes=expr.axes, tags=expr.tags, )
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 map_constant(self, expr, type_context): from loopy.symbolic import Literal if isinstance(expr, (complex, np.complexfloating)): real = self.rec(expr.real) imag = self.rec(expr.imag) iota = p.Variable("I" if "I" not in self.kernel.all_variable_names( ) else "_Complex_I") return real + imag * iota elif np.isnan(expr): return p.Variable("NAN") elif np.isneginf(expr): return -p.Variable("INFINITY") elif np.isinf(expr): return p.Variable("INFINITY") elif isinstance(expr, np.generic): # Explicitly typed: Generated code must reflect type exactly. # FIXME: This assumes a 32-bit architecture. if isinstance(expr, np.float32): return Literal(repr(expr) + "f") elif isinstance(expr, np.float64): return Literal(repr(expr)) # Disabled for now, possibly should be a subtarget. # elif isinstance(expr, np.float128): # return Literal(repr(expr)+"l") elif isinstance(expr, np.integer): suffix = "" iinfo = np.iinfo(expr) if iinfo.min == 0: suffix += "u" if iinfo.max > (2**31 - 1): suffix += "l" return Literal(repr(expr) + suffix) elif isinstance(expr, np.bool_): return Literal("true") if expr else Literal("false") else: raise LoopyError("do not know how to generate code for " "constant of numpy type '%s'" % type(expr).__name__) elif np.isfinite(expr): if type_context == "f": return Literal(repr(np.float32(expr)) + "f") elif type_context == "d": return Literal(repr(float(expr))) elif type_context in ["i", "b"]: return int(expr) else: if is_integer(expr): return int(expr) raise RuntimeError("don't know how to generate code " "for constant '%s'" % expr) else: raise LoopyError("don't know how to generate code " "for constant '%s'" % expr)
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 parse_terminal(self, pstate): scope = self.tree_walker.scope_stack[-1] from pymbolic.parser import (_identifier, _openpar, _closepar, _float) next_tag = pstate.next_tag() if next_tag is _float: value = pstate.next_str_and_advance().lower() if "d" in value: dtype = np.float64 else: dtype = np.float32 value = value.replace("d", "e") if value.startswith("."): prev_value = value value = "0" + value print(value, prev_value) elif value.startswith("-."): prev_value = value value = "-0" + value[1:] print(value, prev_value) return TypedLiteral(value, dtype) elif next_tag is _identifier: name = pstate.next_str_and_advance() if pstate.is_at_end() or pstate.next_tag() is not _openpar: # not a subscript scope.use_name(name) return p.Variable(name) left_exp = p.Variable(name) pstate.advance() pstate.expect_not_end() if scope.is_known(name): cls = p.Subscript else: cls = p.Call if pstate.next_tag is _closepar: pstate.advance() left_exp = cls(left_exp, ()) else: args = self.parse_expression(pstate, self._PREC_FUNC_ARGS) if not isinstance(args, tuple): args = (args, ) left_exp = cls(left_exp, args) pstate.expect(_closepar) pstate.advance() return left_exp else: return ExpressionParserBase.parse_terminal(self, pstate)
def _(poly: ct.NasaPoly2, arg_name): log = p.Variable("log") def gen(c, t): assert len(c) == 7 return (c[0] * log(t) + c[1] * t + c[2] / 2 * t**2 + c[3] / 3 * t**3 + c[4] / 4 * t**4 + c[6]) return nasa7_conditional(p.Variable(arg_name), poly, gen)
def to_loopy_expression( self, indices: SymbolicIndex, expr_context: LoopyExpressionContext) -> ScalarExpression: assert len(indices) == self.num_indices expr_context.update_depends_on(self.depends_on) if indices == (): return prim.Variable(self.name) else: return prim.Variable(self.name)[indices]
def wrap_in_typecast(self, actual_type, needed_dtype, s): if (actual_type.is_complex() and needed_dtype.is_complex() and actual_type != needed_dtype): return p.Variable("%s_cast" % self.complex_type_name(needed_dtype))(s) elif not actual_type.is_complex() and needed_dtype.is_complex(): return p.Variable("%s_fromreal" % self.complex_type_name(needed_dtype))(s) else: return super().wrap_in_typecast_lazy(actual_type, needed_dtype, s)
def layer_integral_predicates(self, tensor, integral_type): self.bag.needs_mesh_layers = True layer = pym.Variable(self.layer_arg) # TODO: Variable layers nlayer = pym.Variable(self.layer_count) which = {"interior_facet_horiz_top": pym.Comparison(layer, "<", nlayer), "interior_facet_horiz_bottom": pym.Comparison(layer, ">", 0), "exterior_facet_top": pym.Comparison(layer, "==", nlayer), "exterior_facet_bottom": pym.Comparison(layer, "==", 0)}[integral_type] return [which]
def test_flop_counter(): x = prim.Variable("x") y = prim.Variable("y") z = prim.Variable("z") subexpr = prim.CommonSubexpression(3 * (x**2 + y + z)) expr = 3 * subexpr + subexpr from pymbolic.mapper.flop_counter import FlopCounter, CSEAwareFlopCounter assert FlopCounter()(expr) == 4 * 2 + 2 assert CSEAwareFlopCounter()(expr) == 4 + 2
def test_parser(): from pymbolic import parse parse("(2*a[1]*b[1]+2*a[0]*b[0])*(hankel_1(-1,sqrt(a[1]**2+a[0]**2)*k) " "-hankel_1(1,sqrt(a[1]**2+a[0]**2)*k))*k /(4*sqrt(a[1]**2+a[0]**2)) " "+hankel_1(0,sqrt(a[1]**2+a[0]**2)*k)") print(repr(parse("d4knl0"))) print(repr(parse("0."))) print(repr(parse("0.e1"))) assert parse("0.e1") == 0 assert parse("1e-12") == 1e-12 print(repr(parse("a >= 1"))) print(repr(parse("a <= 1"))) print(repr(parse(":"))) print(repr(parse("1:"))) print(repr(parse(":2"))) print(repr(parse("1:2"))) print(repr(parse("::"))) print(repr(parse("1::"))) print(repr(parse(":1:"))) print(repr(parse("::1"))) print(repr(parse("3::1"))) print(repr(parse(":5:1"))) print(repr(parse("3:5:1"))) print(repr(parse("g[i,k]+2.0*h[i,k]"))) print(repr(parse("g[i,k]+(+2.0)*h[i,k]"))) print(repr(parse("a - b - c"))) print(repr(parse("-a - -b - -c"))) print(repr(parse("- - - a - - - - b - - - - - c"))) print(repr(parse("~(a ^ b)"))) print(repr(parse("(a | b) | ~(~a & ~b)"))) print(repr(parse("3 << 1"))) print(repr(parse("1 >> 3"))) print(parse("3::1")) assert parse("e1") == prim.Variable("e1") assert parse("d1") == prim.Variable("d1") from pymbolic import variables f, x, y, z = variables("f x y z") assert parse("f((x,y),z)") == f((x, y), z) assert parse("f((x,),z)") == f((x, ), z) assert parse("f(x,(y,z),z)") == f(x, (y, z), z) assert parse("f(x,(y,z),z, name=15)") == f(x, (y, z), z, name=15) assert parse("f(x,(y,z),z, name=15, name2=17)") == f(x, (y, z), z, name=15, name2=17)
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_parser(): from pymbolic import var from dagrt.expression import parse parse("(2*a[1]*b[1]+2*a[0]*b[0])*(hankel_1(-1,sqrt(a[1]**2+a[0]**2)*k) " "-hankel_1(1,sqrt(a[1]**2+a[0]**2)*k))*k /(4*sqrt(a[1]**2+a[0]**2)) " "+hankel_1(0,sqrt(a[1]**2+a[0]**2)*k)") print(repr(parse("d4knl0"))) print(repr(parse("0."))) print(repr(parse("0.e1"))) assert parse("0.e1") == 0 assert parse("1e-12") == 1e-12 print(repr(parse("a >= 1"))) print(repr(parse("a <= 1"))) print(repr(parse("g[i,k]+2.0*h[i,k]"))) print(repr(parse("g[i,k]+(+2.0)*h[i,k]"))) print(repr(parse("a - b - c"))) print(repr(parse("-a - -b - -c"))) print(repr(parse("- - - a - - - - b - - - - - c"))) print(repr(parse("~(a ^ b)"))) print(repr(parse("(a | b) | ~(~a & ~b)"))) print(repr(parse("3 << 1"))) print(repr(parse("1 >> 3"))) print(parse("3::1")) import pymbolic.primitives as prim assert parse("e1") == prim.Variable("e1") assert parse("d1") == prim.Variable("d1") from pymbolic import variables f, x, y, z = variables("f x y z") assert parse("f((x,y),z)") == f((x, y), z) assert parse("f((x,),z)") == f((x, ), z) assert parse("f(x,(y,z),z)") == f(x, (y, z), z) assert parse("f(x,(y,z),z, name=15)") == f(x, (y, z), z, name=15) assert parse("f(x,(y,z),z, name=15, name2=17)") == f(x, (y, z), z, name=15, name2=17) assert parse("<func>yoink") == var("<func>yoink") assert parse("-< func > yoink") == -var("<func>yoink") print(repr(parse("<func>yoink < <p>x"))) print(repr(parse("<func>yoink < - <p>x")))
def loopy_inst_mathfunc(expr, context): children = [loopy_instructions(o, context) for o in expr.ufl_operands] if expr._name == "ln": name = "log" else: name = expr._name return p.Variable(name)(*children)
def binary_tree_mul(start, end): if start + 1 == end: return c_applied[start] mid = (start + end) // 2 lsum = binary_tree_mul(start, mid) rsum = binary_tree_mul(mid, end) return p.Variable("%s_mul" % tgt_name)(lsum, rsum)
def _indices_for_reshape(self, expr: Reshape) -> SymbolicIndex: if expr.array.shape == (): # RHS must be a scalar i.e. RHS' indices are empty assert expr.size == 1 return () newstrides = [1] # reshaped array strides for new_axis_len in reversed(expr.shape[1:]): assert isinstance(new_axis_len, INT_CLASSES) newstrides.insert(0, newstrides[0]*new_axis_len) flattened_idx = sum(prim.Variable(f"_{i}")*stride for i, stride in enumerate(newstrides)) oldstrides: List[IntegralT] = [1] # input array strides for axis_len in reversed(expr.array.shape[1:]): assert isinstance(axis_len, INT_CLASSES) oldstrides.insert(0, oldstrides[0]*axis_len) assert isinstance(expr.array.shape[-1], INT_CLASSES) oldsizetills = [expr.array.shape[-1]] # input array size till for axes idx for old_axis_len in reversed(expr.array.shape[:-1]): assert isinstance(old_axis_len, INT_CLASSES) oldsizetills.insert(0, oldsizetills[0]*old_axis_len) return tuple(((flattened_idx % sizetill) // stride) for stride, sizetill in zip(oldstrides, oldsizetills))
def map_Pow(self, expr): # noqa # SymEngine likes to use as e**a to express exp(a); we undo that here. base, exp = expr.args if base == symengine.E: return prim.Variable("exp")(self.rec(exp)) else: return prim.Power(self.rec(base), self.rec(exp))
def expression_kernel(expr, args): r"""Produce a :class:`pyop2.Kernel` from the processed UFL expression expr and the corresponding args.""" # Empty slot indicating assignment to indexed LHS, so don't do anything if type(expr) is Zero: return fs = args[0].function.function_space() import islpy as isl inames = isl.make_zero_and_vars(["d"]) domain = (inames[0].le_set( inames["d"])) & (inames["d"].lt_set(inames[0] + fs.dof_dset.cdim)) context = Bag() context.within_inames = frozenset(["d"]) context.indices = (p.Variable("d"), ) insn = loopy_instructions(expr, context) data = [arg.arg for arg in args] knl = loopy.make_function([domain], [insn], data, name="expression", silenced_warnings=["summing_if_branches_ops"]) return op2.Kernel(knl, "expression")
def make_random_expression(var_values, size): from random import randrange import pymbolic.primitives as p v = randrange(1500) size[0] += 1 if v < 500 and size[0] < 40: term_count = randrange(2, 5) if randrange(2) < 1: cls = p.Sum else: cls = p.Product return cls( tuple( make_random_expression(var_values, size) for i in range(term_count))) elif v < 750: return make_random_value() elif v < 1000: var_name = "var_%d" % len(var_values) assert var_name not in var_values var_values[var_name] = make_random_value() return p.Variable(var_name) elif v < 1250: # Cannot use '-' because that destroys numpy constants. return p.Sum((make_random_expression(var_values, size), -make_random_expression(var_values, size))) elif v < 1500: # Cannot use '/' because that destroys numpy constants. return p.Quotient(make_random_expression(var_values, size), make_random_expression(var_values, size))
def statement_for(tree, ctx): extent = tree.index.extent assert extent idx = ctx.name_gen(ctx.index_names[tree.index]) ctx.index_extent[idx] = extent with active_indices({tree.index: p.Variable(idx)}, ctx) as ctx_active: return statement(tree.children[0], ctx_active)
def pymbolic_multiindex(self, shape): indices = [] for extent in shape: name = self.name_gen(self.index_names[extent]) self.index_extent[name] = extent indices.append(p.Variable(name)) return tuple(indices)
def _(poly: ct.NasaPoly2, arg_name): def gen(c, t): assert len(c) == 7 return (c[0] + c[1] / 2 * t + c[2] / 3 * t**2 + c[3] / 4 * t**3 + c[4] / 5 * t**4 + c[5] / t) return nasa7_conditional(p.Variable(arg_name), poly, gen)