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 test_projection():
volume = pystencils.fields('volume: float32[100,200,300]')
projections = pystencils.fields('projections: float32[2D]')
projection_matrix = sympy.Matrix(
[[-289.0098977737411, -1205.2274801832275, 0.0, 186000.0],
[-239.9634468375339, -4.188577544948043, 1200.0, 144000.0],
[-0.9998476951563913, -0.01745240643728351, 0.0, 600.0]])
assignments = forward_projection(
volume, projections, projection_matrix) # .compile(target='gpu')
print(type(assignments))
print(pickle.dumps(assignments))
# a = sympy.Symbol('a')
# projection_matrix = sympy.Matrix([[-289.0098977737411, -1205.2274801832275, 0.0, 186000.0],
# [a, - 4.188577544948043, 1200.0, 144000.0],
# [-0.9998476951563913, -0.01745240643728351, 0.0, 600.0]])
# kernel = forward_projection(volume, projections, projection_matrix).compile(target='gpu')
# print(kernel.code)
a = sympy.symbols('a:12')
a = [pystencils.TypedSymbol(s.name, 'float32') for s in a]
A = sympy.Matrix(3, 4, lambda i, j: a[i * 4 + j])
# A = sympy.MatrixSymbol('A', 3, 4)
projection_matrix = A
forward_projection(volume, projections,
projection_matrix).compile(target='gpu')
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 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 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_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 test_address_of():
x, y = pystencils.fields('x,y: int64[2d]')
s = pystencils.TypedSymbol('s', PointerType('int64'))
assignments = pystencils.AssignmentCollection(
{
s: address_of(x[0, 0]),
y[0, 0]: cast_func(s, 'int64')
}, {})
ast = pystencils.create_kernel(assignments)
code = pystencils.show_code(ast)
print(code)
assignments = pystencils.AssignmentCollection(
{y[0, 0]: cast_func(address_of(x[0, 0]), 'int64')}, {})
ast = pystencils.create_kernel(assignments)
code = pystencils.show_code(ast)
print(code)
def test_address_of():
x, y = pystencils.fields('x,y: int64[2d]')
s = pystencils.TypedSymbol('s', PointerType(create_type('int64')))
assert address_of(x[0, 0]).canonical() == x[0, 0]
assert address_of(x[0, 0]).dtype == PointerType(x[0, 0].dtype, restrict=True)
assert address_of(sp.Symbol("a")).dtype == PointerType('void', restrict=True)
assignments = pystencils.AssignmentCollection({
s: address_of(x[0, 0]),
y[0, 0]: cast_func(s, create_type('int64'))
}, {})
ast = pystencils.create_kernel(assignments)
pystencils.show_code(ast)
assignments = pystencils.AssignmentCollection({
y[0, 0]: cast_func(address_of(x[0, 0]), create_type('int64'))
}, {})
ast = pystencils.create_kernel(assignments)
pystencils.show_code(ast)
def __init__(self, body, what, from_iterable, predicate, counter_symbol,
hit_counter_symbol, max_selected, compilation_target):
self.num_iterations = len(from_iterable)
self.elements = what
self.iterable = from_iterable
self._body = body
self.counter_symbol = counter_symbol
self.hit_counter_symbol = hit_counter_symbol
self.predicate = predicate
self.max_selected = max_selected
self.predicate_symbol = pystencils.TypedSymbol('predicate',
create_type('bool'))
predicate_ast = pystencils.create_kernel(pystencils.Assignment(
self.predicate_symbol, predicate),
target=compilation_target)
predicate_assignments = []
stack = [predicate_ast]
while stack:
top = stack.pop()
if isinstance(top, SympyAssignment):
predicate_assignments.insert(0, top)
stack.extend(top.args)
self.predicate_assignments = predicate_assignments
def __next__(self):
name = f"{self._symbol}_{self._ctr}"
self._ctr += 1
if self._dtype is not None:
return pystencils.TypedSymbol(name, self._dtype)
return sp.Symbol(name)