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_assignment_collection_dict_conversion():
x, y = pystencils.fields('x,y: [2D]')
collection_normal = pystencils.AssignmentCollection(
[pystencils.Assignment(x.center(), y[1, 0] + y[0, 0])], [])
collection_dict = pystencils.AssignmentCollection(
{x.center(): y[1, 0] + y[0, 0]}, {})
assert str(collection_normal) == str(collection_dict)
assert collection_dict.main_assignments_dict == {
x.center(): y[1, 0] + y[0, 0]
}
assert collection_dict.subexpressions_dict == {}
collection_normal = pystencils.AssignmentCollection([
pystencils.Assignment(y[1, 0], x.center()),
pystencils.Assignment(y[0, 0], x.center())
], [])
collection_dict = pystencils.AssignmentCollection(
{
y[1, 0]: x.center(),
y[0, 0]: x.center()
}, {})
assert str(collection_normal) == str(collection_dict)
assert collection_dict.main_assignments_dict == {
y[1, 0]: x.center(),
y[0, 0]: x.center()
}
assert collection_dict.subexpressions_dict == {}
def test_replace_and_subs_for_assignment_collection():
x, y = pystencils.fields('x, y: float32[3d]')
a, b, c, d = sp.symbols('a, b, c, d')
assignments = pystencils.AssignmentCollection({
a: sp.floor(1),
b: 2,
c: a + c,
y.center(): sp.ceiling(x.center()) + sp.floor(x.center())
})
expected_assignments = pystencils.AssignmentCollection({
a: sp.floor(3),
b: 2,
c: a + c,
y.center(): sp.ceiling(x.center()) + sp.floor(x.center())
})
assert expected_assignments == assignments.replace(1, 3)
assert expected_assignments == assignments.subs({1: 3})
expected_assignments = pystencils.AssignmentCollection({
d: sp.floor(1),
b: 2,
c: d + c,
y.center(): sp.ceiling(x.center()) + sp.floor(x.center())
})
print(expected_assignments)
print(assignments.subs(a, d))
assert expected_assignments == assignments.subs(a, d)
def add_fixed_constant_boundary_handling(assignments, with_cse=True):
field_accesses = set().union(
itertools.chain.from_iterable(
[a.atoms(Field.Access) for a in assignments]))
if all(all(o == 0 for o in a.offsets) for a in field_accesses):
return assignments
common_shape = next(iter(field_accesses)).field.spatial_shape
ndim = len(common_shape)
def is_out_of_bound(access, shape):
return sp.Or(*[sp.Or(a < 0, a >= s) for a, s in zip(access, shape)])
safe_assignments = [
pystencils.Assignment(
assignment.lhs,
assignment.rhs.subs({
a: ConditionalFieldAccess(
a,
is_out_of_bound(
sp.Matrix(a.offsets) + x_vector(ndim), common_shape))
for a in assignment.rhs.atoms(Field.Access)
if not a.is_absolute_access
})) for assignment in assignments.all_assignments
]
if with_cse:
safe_assignments = sympy_cse(
pystencils.AssignmentCollection(safe_assignments))
return safe_assignments
else:
return pystencils.AssignmentCollection(safe_assignments)
def test_vector_assignment_collection():
"""From #17 (https://i10git.cs.fau.de/pycodegen/pystencils/issues/17)"""
y_m, x_m = sp.Matrix([a, b, c]), sp.Matrix([1, 2, 3])
assignments = ps.AssignmentCollection({y_m: x_m})
print(assignments)
assignments = ps.AssignmentCollection([ps.Assignment(y_m, x_m)])
print(assignments)
def test_native_tensorflow_compilation_cpu():
tf = pytest.importorskip('tensorflow')
module_name = "Ololol"
target = 'cpu'
z, y, x = pystencils.fields("z, y, x: [20,40]")
a = sympy.Symbol('a')
forward_assignments = pystencils.AssignmentCollection({
z[0, 0]: x[0, 0] * sympy.log(a * x[0, 0] * y[0, 0])
})
backward_assignments = create_backward_assignments(forward_assignments)
forward_ast = pystencils.create_kernel(forward_assignments, target)
forward_ast.function_name = 'forward'
backward_ast = pystencils.create_kernel(backward_assignments, target)
backward_ast.function_name = 'backward'
module = TensorflowModule(module_name, [forward_ast, backward_ast])
print(module)
# temp_file = write_cached_content(str(module), '.cpp')
# command = ['c++', '-fPIC', temp_file, '-O2', '-shared',
# '-o', 'foo.so'] + compile_flags + link_flags + extra_flags
# print(command)
# subprocess.check_call(command, env=_compile_env)
lib = module.compile()
assert 'call_forward' in dir(lib)
assert 'call_backward' in dir(lib)
def test_match_for_assignment_collection():
x, y = pystencils.fields('x, y: float32[3d]')
a, b, c, d = sp.symbols('a, b, c, d')
assignments = pystencils.AssignmentCollection({
a: sp.floor(1),
b: 2,
c: a + c,
y.center(): sp.ceiling(x.center()) + sp.floor(x.center())
})
w1 = sp.Wild('w1')
w2 = sp.Wild('w2')
w3 = sp.Wild('w3')
wild_ceiling = sp.ceiling(w1)
wild_addition = w1 + w2
assert assignments.match(pystencils.Assignment(w3, wild_ceiling + w2))[w1] == x.center()
assert assignments.match(pystencils.Assignment(w3, wild_ceiling + w2)) == {
w3: y.center(),
w2: sp.floor(x.center()),
w1: x.center()
}
assert assignments.find(wild_ceiling) == {sp.ceiling(x.center())}
assert len([a for a in assignments.find(wild_addition) if isinstance(a, sp.Add)]) == 2
def test_free_and_bound_symbols():
a1 = ps.Assignment(a, d[0, 0](0))
a2 = ps.Assignment(f[0, 0](1), b * c)
ac = ps.AssignmentCollection([a2], subexpressions=[a1])
assert f[0, 0](1) in ac.bound_symbols
assert d[0, 0](0) in ac.free_symbols
def test_copy():
a1 = ps.Assignment(f[0, 0](0), a * b)
a2 = ps.Assignment(f[0, 0](1), b * c)
ac = ps.AssignmentCollection([a1, a2], subexpressions=[])
ac2 = ac.copy()
assert ac2 == ac
def test_pybind11_compilation_cpu(with_python_bindings):
pytest.importorskip('pybind11')
pytest.importorskip('cppimport')
module_name = "Olololsada"
target = 'cpu'
z, y, x = pystencils.fields("z, y, x: [20,40]")
a = sympy.Symbol('a')
forward_assignments = pystencils.AssignmentCollection(
{z[0, 0]: x[0, 0] * sympy.log(a * x[0, 0] * y[0, 0])})
backward_assignments = create_backward_assignments(forward_assignments)
forward_ast = pystencils.create_kernel(forward_assignments, target)
forward_ast.function_name = 'forward'
backward_ast = pystencils.create_kernel(backward_assignments, target)
backward_ast.function_name = 'backward'
module = PybindModule(module_name, [forward_ast, backward_ast],
with_python_bindings=with_python_bindings)
print(module)
if with_python_bindings:
pybind_extension = module.compile()
assert pybind_extension is not None
assert 'call_forward' in dir(pybind_extension)
assert 'call_backward' in dir(pybind_extension)
def test_torch_native_compilation_cpu():
from torch.utils.cpp_extension import load
module_name = "Ololol"
target = 'cpu'
z, y, x = pystencils.fields("z, y, x: [20,40]")
a = sympy.Symbol('a')
forward_assignments = pystencils.AssignmentCollection(
{z[0, 0]: x[0, 0] * sympy.log(a * x[0, 0] * y[0, 0])})
backward_assignments = create_backward_assignments(forward_assignments)
forward_ast = pystencils.create_kernel(forward_assignments, target)
forward_ast.function_name = 'forward'
backward_ast = pystencils.create_kernel(backward_assignments, target)
backward_ast.function_name = 'backward'
module = TorchModule(module_name, [forward_ast, backward_ast])
print(module)
temp_file = write_cached_content(str(module), '.cpp')
torch_extension = load(module_name, [temp_file])
assert torch_extension is not None
assert 'call_forward' in dir(torch_extension)
assert 'call_backward' in dir(torch_extension)
torch_extension = module.compile()
assert torch_extension is not None
assert 'call_forward' in dir(torch_extension)
assert 'call_backward' in dir(torch_extension)
def test_tfmad_gradient_check_torch():
torch = pytest.importorskip('torch')
a, b, out = ps.fields("a, b, out: float[5,7]")
cont = 2 * ps.fd.Diff(a, 0) - 1.5 * ps.fd.Diff(a, 1) \
- ps.fd.Diff(b, 0) + 3 * ps.fd.Diff(b, 1)
discretize = ps.fd.Discretization2ndOrder(dx=1)
discretization = discretize(cont) + 1.2 * a.center
assignment = ps.Assignment(out.center(), discretization)
assignment_collection = ps.AssignmentCollection([assignment], [])
print('Forward')
print(assignment_collection)
print('Backward')
auto_diff = pystencils_autodiff.AutoDiffOp(assignment_collection,
diff_mode='transposed-forward')
backward = auto_diff.backward_assignments
print(backward)
print('Forward output fields (to check order)')
print(auto_diff.forward_input_fields)
a_tensor = torch.zeros(*a.shape, dtype=torch.float64, requires_grad=True)
b_tensor = torch.zeros(*b.shape, dtype=torch.float64, requires_grad=True)
function = auto_diff.create_tensorflow_op({
a: a_tensor,
b: b_tensor
},
backend='torch')
torch.autograd.gradcheck(function.apply, [a_tensor, b_tensor])
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_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_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_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_tensorflow_jit_cpu():
pytest.importorskip('tensorflow')
module_name = "Ololol"
target = 'cpu'
z, y, x = pystencils.fields("z, y, x: [20,40]")
a = sympy.Symbol('a')
forward_assignments = pystencils.AssignmentCollection(
{z[0, 0]: x[0, 0] * sympy.log(a * x[0, 0] * y[0, 0])})
backward_assignments = create_backward_assignments(forward_assignments)
forward_ast = pystencils.create_kernel(forward_assignments, target)
forward_ast.function_name = 'forward_jit'
backward_ast = pystencils.create_kernel(backward_assignments, target)
backward_ast.function_name = 'backward_jit'
module = TensorflowModule(module_name, [forward_ast, backward_ast])
lib = pystencils_autodiff.tensorflow_jit.compile_sources_and_load(
[str(module)])
assert 'call_forward_jit' in dir(lib)
assert 'call_backward_jit' in dir(lib)
lib = module.compile()
assert 'call_forward_jit' in dir(lib)
assert 'call_backward_jit' in dir(lib)
def test_simple_2d_check_assignment_collection():
# use simply example
z, y, x = ps.fields("z, y, x: [2d]")
forward_assignments = ps.AssignmentCollection([ps.Assignment(
z[0, 0], x[0, 0]*sp.log(x[0, 0]*y[0, 0]))], [])
jac = pystencils_autodiff.get_jacobian_of_assignments(
forward_assignments, [x[0, 0], y[0, 0]])
assert jac.shape == (len(forward_assignments.bound_symbols),
len(forward_assignments.free_symbols))
print(repr(jac))
assert repr(jac) == 'Matrix([[log(x_C*y_C) + 1, x_C/y_C]])'
for diff_mode in DiffModes:
pystencils_autodiff.create_backward_assignments(
forward_assignments, diff_mode=diff_mode)
pystencils_autodiff.create_backward_assignments(
pystencils_autodiff.create_backward_assignments(forward_assignments), diff_mode=diff_mode)
result1 = pystencils_autodiff.create_backward_assignments(
forward_assignments, diff_mode=DiffModes.TRANSPOSED)
result2 = pystencils_autodiff.create_backward_assignments(
forward_assignments, diff_mode=DiffModes.TF_MAD)
assert result1 == result2
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_module_printing_parameter():
module_name = "Ololol"
for target in ('cpu', 'gpu'):
z, y, x = pystencils.fields("z, y, x: [20,40]")
a = sympy.Symbol('a')
forward_assignments = pystencils.AssignmentCollection(
{z[0, 0]: x[0, 0] * sympy.log(a * x[0, 0] * y[0, 0])})
backward_assignments = create_backward_assignments(forward_assignments)
forward_ast = pystencils.create_kernel(forward_assignments, target)
forward_ast.function_name = 'forward'
backward_ast = pystencils.create_kernel(backward_assignments, target)
backward_ast.function_name = 'backward'
module = TorchModule(module_name, [forward_ast, backward_ast])
print(module)
module = TensorflowModule(module_name, {forward_ast: backward_ast})
print(module)
if target == 'cpu':
module = PybindModule(module_name, [forward_ast, backward_ast])
print(module)
module = PybindModule(module_name, forward_ast)
print(module)
def test_staggered(vectorized):
"""Make sure that the RNG counter can be substituted during loop cutting"""
dh = ps.create_data_handling((8, 8), default_ghost_layers=0, default_target=Target.CPU)
j = dh.add_array("j", values_per_cell=dh.dim, field_type=ps.FieldType.STAGGERED_FLUX)
a = ps.AssignmentCollection([ps.Assignment(j.staggered_access(n), 0) for n in j.staggered_stencil])
rng_symbol_gen = random_symbol(a.subexpressions, dim=dh.dim, rng_node=PhiloxTwoDoubles)
a.main_assignments[0] = ps.Assignment(a.main_assignments[0].lhs, next(rng_symbol_gen))
kernel = ps.create_staggered_kernel(a, target=dh.default_target).compile()
if not vectorized:
return
if not instruction_sets:
pytest.skip("cannot detect CPU instruction set")
pytest.importorskip('islpy')
cpu_vectorize_info = {'assume_inner_stride_one': True, 'assume_aligned': False,
'instruction_set': instruction_sets[-1]}
dh.fill(j.name, 867)
dh.run_kernel(kernel, seed=5, time_step=309)
ref_data = dh.gather_array(j.name)
kernel2 = ps.create_staggered_kernel(a, target=dh.default_target, cpu_vectorize_info=cpu_vectorize_info).compile()
dh.fill(j.name, 867)
dh.run_kernel(kernel2, seed=5, time_step=309)
data = dh.gather_array(j.name)
assert np.allclose(ref_data, data)
def test_fixed_constant_bh(num_ghost_layers):
ndim = 2
offsets = list(itertools.product(range(num_ghost_layers + 1), repeat=ndim))
x, y = pystencils.fields(f'x, y: float64[{ndim}d]')
assignments = pystencils.AssignmentCollection({
y.center:
sp.Add(*[x.__getitem__(o) for o in offsets]) / len(offsets)
})
kernel = pystencils.create_kernel(assignments).compile()
print(kernel.code)
bh_assignments = add_fixed_constant_boundary_handling(
assignments, num_ghost_layers)
bh_kernel = pystencils.create_kernel(bh_assignments,
ghost_layers=0).compile()
print(bh_kernel.code)
noise = np.random.rand(*[20, 30, 40][:ndim])
out1 = np.zeros_like(noise)
out2 = np.zeros_like(noise)
kernel(x=noise, y=out1)
bh_kernel(x=noise, y=out2)
def test_reproducability():
from sympy.core.cache import clear_cache
output_0 = None
for i in range(10):
module_name = "Ololol"
target = 'cpu'
z, y, x = pystencils.fields("z, y, x: [20,40]")
a = sympy.Symbol('a')
forward_assignments = pystencils.AssignmentCollection(
{z[0, 0]: x[0, 0] * sympy.log(a * x[0, 0] * y[0, 0])})
backward_assignments = create_backward_assignments(forward_assignments)
forward_ast = pystencils.create_kernel(forward_assignments, target)
forward_ast.function_name = 'forward'
backward_ast = pystencils.create_kernel(backward_assignments, target)
backward_ast.function_name = 'backward'
new_output = str(TorchModule(module_name, [forward_ast, backward_ast]))
TorchModule(module_name, [forward_ast, backward_ast]).compile()
clear_cache()
if not output_0:
output_0 = new_output
assert output_0 == new_output
def __init__(self, assignments, perform_cse=True, *args, **kwargs):
if isinstance(assignments, pystencils.AssignmentCollection):
assignments = assignments.all_assignments
assignments = pystencils.AssignmentCollection(assignments, {})
if perform_cse:
main_assignments = [a for a in assignments if not hasattr(
a, 'lhs') or isinstance(a.lhs, pystencils.Field.Access)]
subexpressions = [a for a in assignments if hasattr(
a, 'lhs') and not isinstance(a.lhs, pystencils.Field.Access)]
assignments = pystencils.AssignmentCollection(main_assignments, subexpressions)
assignments = pystencils.simp.sympy_cse(assignments)
super(AssignmentCollection, self).__init__(assignments.all_assignments, {}, *args, **kwargs)
self.args = []
self.kwargs = {}
self._autodiff = None
self._kernel = None
def test_execute_torch(target):
import pycuda.autoinit
module_name = "Ololol" + target
z, y, x = pystencils.fields("z, y, x: [20,40]")
a = sympy.Symbol('a')
forward_assignments = pystencils.AssignmentCollection(
{z[0, 0]: x[0, 0] * sympy.log(a * x[0, 0] * y[0, 0])})
# backward_assignments = create_backward_assignments(forward_assignments)
if target == 'cpu':
x = np.random.rand(20, 40)
y = np.random.rand(20, 40)
z = np.zeros((20, 40))
else:
gpuarray = pytest.importorskip('pycuda.gpuarray')
x = gpuarray.to_gpu(np.random.rand(20, 40))
y = gpuarray.to_gpu(np.random.rand(20, 40))
z = gpuarray.zeros((20, 40), np.float64)
kernel = pystencils.create_kernel(forward_assignments, target=target)
kernel.function_name = 'forward'
torch_module = TorchModule(module_name, [kernel]).compile()
pystencils_module = kernel.compile()
pystencils_module(x=x, y=y, z=z, a=5.)
if target == 'gpu':
x = x.get()
y = y.get()
z = z.get()
z_pystencils = np.copy(z)
import torch
x = torch.Tensor(x)
y = torch.Tensor(y)
z = torch.Tensor(z)
if target == 'gpu':
x = x.double().cuda()
y = y.double().cuda()
z = z.double().cuda()
else:
x = x.double()
y = y.double()
z = z.double()
torch_module.call_forward(x=x, y=y, z=z, a=5.)
if target == 'gpu':
z = z.cpu()
z_torch = np.copy(z)
assert np.allclose(z_torch[1:-1, 1:-1],
z_pystencils[1:-1, 1:-1],
atol=1e-6)
def test_cuda_unknown():
x, y = pystencils.fields('x,y: float32 [2d]')
assignments = pystencils.AssignmentCollection({
get_dummy_symbol():
sympy.Function('wtf')(address_of(y.center()), 2),
})
ast = pystencils.create_kernel(assignments, target=Target.GPU)
pystencils.show_code(ast)
def test_destructuring_alternative_field_class():
z, x, y = pystencils.fields("z, y, x: [2d]")
normal_assignments = pystencils.AssignmentCollection([
pystencils.Assignment(z[0, 0], x[0, 0] * sympy.log(x[0, 0] * y[0, 0]))
], [])
ast = pystencils.create_kernel(normal_assignments, target='gpu')
ast.body = DestructuringEmojiClass(ast.body)
print(pystencils.show_code(ast))
def test_native_tensorflow_compilation_gpu():
tf = pytest.importorskip('tensorflow')
module_name = "Ololol"
target = 'gpu'
z, y, x = pystencils.fields("z, y, x: [20,40]")
a = sympy.Symbol('a')
forward_assignments = pystencils.AssignmentCollection({
z[0, 0]: x[0, 0] * sympy.log(a * x[0, 0] * y[0, 0])
})
backward_assignments = create_backward_assignments(forward_assignments)
forward_ast = pystencils.create_kernel(forward_assignments, target)
forward_ast.function_name = 'forward2'
backward_ast = pystencils.create_kernel(backward_assignments, target)
backward_ast.function_name = 'backward2'
module = TensorflowModule(module_name, [forward_ast, backward_ast])
print(str(module))
# temp_file = write_cached_content(str(module), '.cu')
# if 'tensorflow_host_compiler' not in get_compiler_config():
# get_compiler_config()['tensorflow_host_compiler'] = get_compiler_config()['command']
# # on my machine g++-6 and clang-7 are working
# # '-ccbin',
# # 'g++-6',
# command = ['nvcc',
# temp_file.name,
# '--expt-relaxed-constexpr',
# '-ccbin',
# get_compiler_config()['tensorflow_host_compiler'],
# '-std=c++14',
# '-x',
# 'cu',
# '-Xcompiler',
# '-fPIC',
# '-c',
# '-o',
# 'foo_gpu.o'] + compile_flags + extra_flags
# subprocess.check_call(command)
# command = ['c++', '-fPIC', 'foo_gpu.o',
# '-shared', '-o', 'foo_gpu.so'] + link_flags
# subprocess.check_call(command)
lib = module.compile()
assert 'call_forward2' in dir(lib)
#
assert 'call_backward2' in dir(lib)
def eigenvalues_3d(eig1, eig2, eig3, xx, xy, xz, yy, yz, zz):
H = sympy.Matrix([[xx.center, xy.center, xz.center],
[xy.center, yy.center, yz.center],
[xz.center, yz.center, zz.center]])
eigenvalues = list(H.eigenvals())
assignments = pystencils.AssignmentCollection(
{[eig1.center, eig2.center, eig3.center][i]: sympy.re(eigenvalues[i])
for i in range(3)})
class complex_symbol_generator():
def __iter__(self):
counter = 0
while True:
yield TypedSymbol('xi_%i' % counter, create_type(np.complex64))
counter += 1
assignments = pystencils.AssignmentCollection(
optimize_assignments(assignments, [evaluate_constant_terms]))
assignments.subexpression_symbol_generator = complex_symbol_generator()
assignments = sympy_cse(assignments)
assignments = optimize_assignments(assignments, [use_complex_sqrt])
# complex_rhs = []
# for a in assignments:
# if isinstance(a.lhs, pystencils.Field.Access):
# complex_rhs.append(a.rhs)
# assignments = pystencils_reco.AssignmentCollection(assignments).subs(
# {c: sympy.Function('real')(c) for c in complex_rhs})
# print(assignments)
# complex_symbols = [(a.lhs, a.rhs) for a in assignments if any(atom.is_real is False for atom in a.rhs.atoms())]
# assignments = assignments.subs({a.lhs: a.rhs for a in assignments if any(
# atom.is_real is False for atom in a.rhs.atoms())})
# print(complex_symbols)
assignments = pystencils_reco.AssignmentCollection(assignments,
perform_cse=False)
return assignments
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)
