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})
