def __init__(self, pdf_field, stencil, prev_timestep=Timestep.BOTH, streaming_pattern='pull', index_dtype=np.int64, offsets_dtype=np.int64): if prev_timestep == Timestep.BOTH and is_inplace(streaming_pattern): raise ValueError('Cannot create index arrays for both kinds of timesteps for inplace streaming pattern ' + streaming_pattern) prev_accessor = get_accessor(streaming_pattern, prev_timestep) next_accessor = get_accessor(streaming_pattern, prev_timestep.next()) outward_accesses = prev_accessor.write(pdf_field, stencil) inward_accesses = next_accessor.read(pdf_field, stencil) self._accesses = {'out': outward_accesses, 'in': inward_accesses} self._pdf_field = pdf_field self._stencil = stencil self._dim = stencil.D self._q = stencil.Q self._coordinate_names = ['x', 'y', 'z'][:self._dim] self._index_dtype = create_type(index_dtype) self._offsets_dtype = create_type(offsets_dtype) self._required_index_arrays = set() self._required_offset_arrays = set() self._trivial_index_translations, self._trivial_offset_translations = self._collect_trivial_translations()
def test_prod_var_limit(): k = pystencils.TypedSymbol('k', create_type('int64')) limit = pystencils.TypedSymbol('limit', create_type('int64')) sum = sympy.Sum(k, (k, 1, limit)) expanded_sum = sum.replace(limit, 100).doit() print(sum) print(expanded_sum) x = pystencils.fields('x: int64[1d]') assignments = pystencils.AssignmentCollection({x.center(): sum}) ast = pystencils.create_kernel(assignments) code = str(pystencils.show_code(ast)) kernel = ast.compile() print(code) array = np.zeros((10, ), np.int64) kernel(x=array, limit=100) assert np.allclose(array, int(expanded_sum) * np.ones_like(array))
def _print_Function(self, expr): if isinstance(expr, vector_memory_access): arg, data_type, aligned, _, mask, stride = expr.args if stride != 1: return self.instruction_set['loadS'].format(f"& {self._print(arg)}", stride, **self._kwargs) instruction = self.instruction_set['loadA'] if aligned else self.instruction_set['loadU'] return instruction.format(f"& {self._print(arg)}", **self._kwargs) elif isinstance(expr, cast_func): arg, data_type = expr.args if type(data_type) is VectorType: # vector_memory_access is a cast_func itself so it should't be directly inside a cast_func assert not isinstance(arg, vector_memory_access) if isinstance(arg, sp.Tuple): is_boolean = get_type_of_expression(arg[0]) == create_type("bool") is_integer = get_type_of_expression(arg[0]) == create_type("int") printed_args = [self._print(a) for a in arg] instruction = 'makeVecBool' if is_boolean else 'makeVecInt' if is_integer else 'makeVec' if instruction == 'makeVecInt' and 'makeVecIndex' in self.instruction_set: increments = np.array(arg)[1:] - np.array(arg)[:-1] if len(set(increments)) == 1: return self.instruction_set['makeVecIndex'].format(printed_args[0], increments[0], **self._kwargs) return self.instruction_set[instruction].format(*printed_args, **self._kwargs) else: is_boolean = get_type_of_expression(arg) == create_type("bool") is_integer = get_type_of_expression(arg) == create_type("int") or \ (isinstance(arg, TypedSymbol) and not isinstance(arg.dtype, VectorType) and arg.dtype.is_int()) instruction = 'makeVecConstBool' if is_boolean else \ 'makeVecConstInt' if is_integer else 'makeVecConst' return self.instruction_set[instruction].format(self._print(arg), **self._kwargs) elif expr.func == fast_division: result = self._scalarFallback('_print_Function', expr) if not result: result = self.instruction_set['/'].format(self._print(expr.args[0]), self._print(expr.args[1]), **self._kwargs) return result elif expr.func == fast_sqrt: return f"({self._print(sp.sqrt(expr.args[0]))})" elif expr.func == fast_inv_sqrt: result = self._scalarFallback('_print_Function', expr) if not result: if 'rsqrt' in self.instruction_set: return self.instruction_set['rsqrt'].format(self._print(expr.args[0]), **self._kwargs) else: return f"({self._print(1 / sp.sqrt(expr.args[0]))})" elif isinstance(expr, vec_any) or isinstance(expr, vec_all): instr = 'any' if isinstance(expr, vec_any) else 'all' expr_type = get_type_of_expression(expr.args[0]) if type(expr_type) is not VectorType: return self._print(expr.args[0]) else: if isinstance(expr.args[0], sp.Rel): op = expr.args[0].rel_op if (instr, op) in self.instruction_set: return self.instruction_set[(instr, op)].format(*[self._print(a) for a in expr.args[0].args], **self._kwargs) return self.instruction_set[instr].format(self._print(expr.args[0]), **self._kwargs) return super(VectorizedCustomSympyPrinter, self)._print_Function(expr)
def test_collation(): double_type = create_type("double") float_type = create_type("float32") double4_type = VectorType(double_type, 4) float4_type = VectorType(float_type, 4) assert collate_types([double_type, float_type]) == double_type assert collate_types([double4_type, float_type]) == double4_type assert collate_types([double4_type, float4_type]) == double4_type
def _print_Number(self, n): if get_type_of_expression(n) == create_type("int"): return ir.Constant(self.integer, int(n)) elif get_type_of_expression(n) == create_type("double"): return ir.Constant(self.fp_type, float(n)) else: raise NotImplementedError("Numbers can only have int and double", n)
def test_vector_type(): double_type = create_type("double") float_type = create_type("float32") double4_type = VectorType(double_type, 4) float4_type = VectorType(float_type, 4) assert double4_type.item_size == 4 assert float4_type.item_size == 4 assert not double4_type == 4
def test_assumptions(): x = ps.fields('x: float32[3d]') assert x.shape[0].is_nonnegative assert (2 * x.shape[0]).is_nonnegative assert (2 * x.shape[0]).is_integer assert (TypedSymbol('a', create_type('uint64'))).is_nonnegative assert (TypedSymbol('a', create_type('uint64'))).is_positive is None assert (TypedSymbol('a', create_type('uint64')) + 1).is_positive assert (x.shape[0] + 1).is_real
def test_pointer_type(): double_type = create_type("double") float_type = create_type("float32") double4_type = PointerType(double_type, restrict=True) float4_type = PointerType(float_type, restrict=False) assert double4_type.item_size == 1 assert float4_type.item_size == 1 assert not double4_type == 4 assert not double4_type.alias assert float4_type.alias
def test_address_of_with_cse(): x, y = pystencils.fields('x,y: int64[2d]') s = pystencils.TypedSymbol('s', PointerType(create_type('int64'))) assignments = pystencils.AssignmentCollection({ y[0, 0]: cast_func(address_of(x[0, 0]), create_type('int64')) + s, x[0, 0]: cast_func(address_of(x[0, 0]), create_type('int64')) + 1 }, {}) ast = pystencils.create_kernel(assignments) pystencils.show_code(ast) assignments_cse = sympy_cse(assignments) ast = pystencils.create_kernel(assignments_cse) pystencils.show_code(ast)
def __init__(self): super(CustomSympyPrinter, self).__init__() self._float_type = create_type("float32") if 'Min' in self.known_functions: del self.known_functions['Min'] if 'Max' in self.known_functions: del self.known_functions['Max']
def test_floor_ceil_float_no_optimization(): x, y = pystencils.fields('x,y: float32[2d]') a, b, c = sp.symbols('a, b, c') int_symbol = sp.Symbol('int_symbol', integer=True) typed_symbol = pystencils.TypedSymbol('typed_symbol', create_type('float32')) assignments = pystencils.AssignmentCollection({ a: sp.floor(1), b: sp.ceiling(typed_symbol), c: sp.floor(int_symbol), y.center(): sp.ceiling(x.center()) + sp.floor(x.center()) }) assert not typed_symbol.is_integer print(sp.simplify(sp.ceiling(typed_symbol))) print(assignments) wild_floor = sp.floor(sp.Wild('w1')) assert not sp.floor(int_symbol).match(wild_floor) assert sp.floor(a).match(wild_floor) assert assignments.find(wild_floor)
def test_dynamic_matrix_location_dependent(): try: from pystencils.data_types import TypedMatrixSymbol except ImportError: import pytest pytest.skip() x, y = pystencils.fields('x, y: float32[3d]') A = TypedMatrixSymbol('A', 3, 1, create_type('double'), CustomCppType('Vector3<double>')) my_fun_call = DynamicFunction( TypedSymbol('my_fun', 'std::function<Vector3<double>(int, int, int)>'), A.dtype, *pystencils.x_vector(3)) assignments = pystencils.AssignmentCollection({ A: my_fun_call, y.center: A[0] + A[1] + A[2] }) ast = pystencils.create_kernel(assignments) pystencils.show_code(ast, custom_backend=FrameworkIntegrationPrinter()) my_fun_call = DynamicFunction( TypedSymbol('my_fun', TemplateType('Functor_T')), A.dtype, *pystencils.x_vector(3)) assignments = pystencils.AssignmentCollection({ A: my_fun_call, y.center: A[0] + A[1] + A[2] }) ast = pystencils.create_kernel(assignments) pystencils.show_code(ast, custom_backend=FrameworkIntegrationPrinter())
def test_complex_execution(dtype, target, with_complex_argument): complex_dtype = f'complex{64 if dtype ==np.float32 else 128}' x, y = pystencils.fields(f'x, y: {complex_dtype}[2d]') x_arr = np.zeros((20, 30), complex_dtype) y_arr = np.zeros((20, 30), complex_dtype) if with_complex_argument: a = pystencils.TypedSymbol('a', create_type(complex_dtype)) else: a = (2j + 1) assignments = AssignmentCollection({y.center: x.center + a}) if target == pystencils.Target.GPU: pytest.importorskip('pycuda') from pycuda.gpuarray import zeros x_arr = zeros((20, 30), complex_dtype) y_arr = zeros((20, 30), complex_dtype) kernel = pystencils.create_kernel(assignments, target=target, data_type=dtype).compile() if with_complex_argument: kernel(x=x_arr, y=y_arr, a=2j + 1) else: kernel(x=x_arr, y=y_arr) if target == pystencils.Target.GPU: y_arr = y_arr.get() assert np.allclose(y_arr, 2j + 1)
def _comparison(self, cmpop, expr): if collate_types([get_type_of_expression(arg) for arg in expr.args]) == create_type('double'): comparison = self.builder.fcmp_unordered else: comparison = self.builder.icmp_signed return comparison(cmpop, self._print(expr.lhs), self._print(expr.rhs))
def upsample(input: {'field_type': pystencils.field.FieldType.CUSTOM}, result, factor): ndim = input.spatial_dimensions here = pystencils.x_vector(ndim) assignments = AssignmentCollection({ result.center: pystencils.astnodes.ConditionalFieldAccess( input.absolute_access( tuple( cast_func(sympy.S(1) / factor * h, create_type('int64')) for h in here), ()), sympy.Or(*[s % cast_func(factor, 'int64') > 0 for s in here])) }) def create_autodiff(self, constant_fields=None, **kwargs): backward_assignments = downsample(AdjointField(result), AdjointField(input), factor) self._autodiff = pystencils.autodiff.AutoDiffOp( assignments, "", backward_assignments=backward_assignments, **kwargs) assignments._create_autodiff = types.MethodType(create_autodiff, assignments) return assignments
def _print_Piecewise(self, expr): result = self._scalarFallback('_print_Piecewise', expr) if result: return result if expr.args[-1].cond.args[0] is not sp.sympify(True): # We need the last conditional to be a True, otherwise the resulting # function may not return a result. raise ValueError("All Piecewise expressions must contain an " "(expr, True) statement to be used as a default " "condition. Without one, the generated " "expression may not evaluate to anything under " "some condition.") result = self._print(expr.args[-1][0]) for true_expr, condition in reversed(expr.args[:-1]): if isinstance(condition, cast_func) and get_type_of_expression( condition.args[0]) == create_type("bool"): if not KERNCRAFT_NO_TERNARY_MODE: result = "(({}) ? ({}) : ({}))".format( self._print(condition.args[0]), self._print(true_expr), result) else: print("Warning - skipping ternary op") else: # noinspection SpellCheckingInspection result = self.instruction_set['blendv'].format( result, self._print(true_expr), self._print(condition)) return result
def struct_from_numpy_dtype(struct_name, numpy_dtype): result = "struct %s { \n" % (struct_name, ) equality_compare = [] constructor_params = [] constructor_initializer_list = [] for name, (sub_type, offset) in numpy_dtype.fields.items(): pystencils_type = create_type(sub_type) result += " %s %s;\n" % (pystencils_type, name) if name in boundary_index_array_coordinate_names or name == direction_member_name: constructor_params.append("%s %s_" % (pystencils_type, name)) constructor_initializer_list.append("%s(%s_)" % (name, name)) else: constructor_initializer_list.append("%s()" % name) if pystencils_type.is_float(): equality_compare.append("floatIsEqual(%s, o.%s)" % (name, name)) else: equality_compare.append("%s == o.%s" % (name, name)) result += " %s(%s) : %s {}\n" % \ (struct_name, ", ".join(constructor_params), ", ".join(constructor_initializer_list)) result += " bool operator==(const %s & o) const {\n return %s;\n }\n" % \ (struct_name, " && ".join(equality_compare)) result += "};\n" return result
def test_sum_use_float(): sum = sympy.Sum(k, (k, 1, 100)) expanded_sum = sum.doit() print(sum) print(expanded_sum) x = pystencils.fields('x: float32[1d]') assignments = pystencils.AssignmentCollection({x.center(): sum}) ast = pystencils.create_kernel(assignments, data_type=create_type('float32')) code = str(pystencils.show_code(ast)) kernel = ast.compile() print(code) print(pystencils.show_code(ast)) assert 'float sum' in code array = np.zeros((10, ), np.float32) kernel(x=array) assert np.allclose(array, int(expanded_sum) * np.ones_like(array))
def test_sum_use_float(default_assignment_simplifications): sum = sympy.Sum(sp.abc.k, (sp.abc.k, 1, 100)) expanded_sum = sum.doit() print(sum) print(expanded_sum) x = ps.fields('x: float32[1d]') assignments = ps.AssignmentCollection({x.center(): sum}) config = ps.CreateKernelConfig( default_assignment_simplifications=default_assignment_simplifications, data_type=create_type('float32')) ast = ps.create_kernel(assignments, config=config) code = ps.get_code_str(ast) kernel = ast.compile() print(code) if default_assignment_simplifications is False: assert 'float sum' in code array = np.zeros((10, ), np.float32) kernel(x=array) assert np.allclose(array, int(expanded_sum) * np.ones_like(array))
def test_product(default_assignment_simplifications): k = ps.TypedSymbol('k', create_type('int64')) sum = sympy.Product(k, (k, 1, 10)) expanded_sum = sum.doit() print(sum) print(expanded_sum) x = ps.fields('x: int64[1d]') assignments = ps.AssignmentCollection({x.center(): sum}) config = ps.CreateKernelConfig( default_assignment_simplifications=default_assignment_simplifications) ast = ps.create_kernel(assignments, config=config) code = ps.get_code_str(ast) kernel = ast.compile() print(code) if default_assignment_simplifications is False: assert 'int64_t product' in code array = np.zeros((10, ), np.int64) kernel(x=array) assert np.allclose(array, int(expanded_sum) * np.ones_like(array))
def additional_data(self): """ In case of the UBB boundary additional data is a velocity vector. This vector is added to each cell to realize velocity profiles for the inlet.""" if self.velocity_is_callable: return [(f'vel_{i}', create_type(self.data_type)) for i in range(self.dim)] else: return []
def is_loop_counter_symbol(symbol): prefix = LoopOverCoordinate.LOOP_COUNTER_NAME_PREFIX if not symbol.name.startswith(prefix): return None if symbol.dtype != create_type('int'): return None coordinate = int(symbol.name[len(prefix) + 1:]) return coordinate
def test_loop_information(): f, g = ps.fields("f, g: double[2D]") update_rule = ps.Assignment(g[0, 0], f[0, 0]) ast = ps.create_kernel(update_rule) inner_loops = [l for l in filtered_tree_iteration(ast, LoopOverCoordinate, stop_type=SympyAssignment) if l.is_innermost_loop] loop_order = [] for i in get_loop_hierarchy(inner_loops[0].args[0]): loop_order.append(i) assert loop_order == [0, 1] loop_symbols = get_loop_counter_symbol_hierarchy(inner_loops[0].args[0]) assert loop_symbols == [TypedSymbol("ctr_1", create_type("int"), nonnegative=True), TypedSymbol("ctr_0", create_type("int"), nonnegative=True)]
def test_type_interference(): x = pystencils.fields('x: float32[3d]') assignments = pystencils.AssignmentCollection({ a: cast_func(10, create_type('float64')), b: cast_func(10, create_type('uint16')), e: 11, c: b, f: c + b, d: c + b + x.center + e, x.center: c + b + x.center }) ast = pystencils.create_kernel(assignments) code = str(pystencils.get_code_str(ast)) assert 'double a' in code assert 'uint16_t b' in code assert 'uint16_t f' in code assert 'int64_t e' in code
def type_symbol(term): if isinstance(term, Field.Access) or isinstance(term, TypedSymbol): return term elif isinstance(term, sp.Symbol): if not hasattr(type_info, '__getitem__'): return TypedSymbol(term.name, create_type(type_info)) else: return TypedSymbol(term.name, type_info[term.name]) else: raise ValueError("Term has to be field access or symbol")
def _print_Mul(self, expr): nodes = [self._print(a) for a in expr.args] e = nodes[0] if get_type_of_expression(expr) == create_type('double'): mul = self.builder.fmul else: # int TODO unsigned/signed mul = self.builder.mul for node in nodes[1:]: e = mul(e, node) return e
def _print_Add(self, expr): nodes = [self._print(a) for a in expr.args] e = nodes[0] if get_type_of_expression(expr) == create_type('double'): add = self.builder.fadd else: # int TODO unsigned/signed add = self.builder.add for node in nodes[1:]: e = add(e, node) return e
def test_wild_typed_symbol(): x = pystencils.fields('x: float32[3d]') typed_symbol = pystencils.data_types.TypedSymbol('a', create_type('float64')) assert x.center().match(sp.Wild('w1')) assert typed_symbol.match(sp.Wild('w1')) wild_ceiling = sp.ceiling(sp.Wild('w1')) assert sp.ceiling(x.center()).match(wild_ceiling) assert sp.ceiling(typed_symbol).match(wild_ceiling)
def additional_data(self): """Used internally only. For the FreeSlip boundary the information of the normal direction for each pdf direction is needed. This information is stored in the index vector.""" if self.normal_direction: return [] else: data_type = create_type('int32') wnz = [] if self.dim == 2 else [('wnz', data_type)] data = [('wnx', data_type), ('wny', data_type)] + wnz return data + [('ref_dir', data_type)]
def __init__(self, stencil, mirror_axis, dtype=np.int64): offsets_dtype = create_type(dtype) mirrored_stencil_symbol = MirroredStencilDirections._mirrored_symbol(mirror_axis) mirrored_directions = [stencil.index(MirroredStencilDirections.mirror_stencil(direction, mirror_axis)) for direction in stencil] code = "\n" code += _array_pattern(offsets_dtype, mirrored_stencil_symbol.name, mirrored_directions) super(MirroredStencilDirections, self).__init__(code, symbols_read=set(), symbols_defined={mirrored_stencil_symbol})