def test_simple_2d_check_raw_assignments(): # use simply example z, y, x = ps.fields("z, y, x: [2d]") forward_assignments = [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 == (1, 2) 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)
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 __call__(self, pdf_field: pystencils_autodiff.AdjointField, direction_symbol, *args, **kwargs): # apply heuristics if pdf_field.name.startswith('diff'): forward_field = pystencils.Field.new_field_with_different_name( pdf_field, pdf_field.name[len('diff'):]) pdf_field = pystencils_autodiff.AdjointField(forward_field) assert isinstance(pdf_field, pystencils_autodiff.AdjointField), \ f'{pdf_field} should be a pystencils_autodiff.AdjointField to use AdjointBoundaryCondition' forward_field = pdf_field.corresponding_forward_field forward_assignments = self._forward_condition(forward_field, direction_symbol, *args, **kwargs) backward_assignments = pystencils_autodiff.create_backward_assignments( forward_assignments, diff_fields_prefix=pdf_field.name_prefix, time_constant_fields=self._time_constant_fields, constant_fields=self._constant_fields) assert backward_assignments.all_assignments, ( "Must have a at least one read field in forward boundary " "to have an meaningful adjoint boundary condition") return backward_assignments
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_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_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_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 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_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 test_tfmad_stencil(): f, out = ps.fields("f, out: double[2D]") cont = ps.fd.Diff(f, 0) - ps.fd.Diff(f, 1) discretize = ps.fd.Discretization2ndOrder(dx=1) discretization = discretize(cont) assignment = ps.Assignment(out.center(), discretization) assignment_collection = ps.AssignmentCollection([assignment], []) print('Forward') print(assignment_collection) print('Backward') backward = pystencils_autodiff.create_backward_assignments( assignment_collection, diff_mode='transposed-forward') print(backward)
def test_tensorflow_jit_gpu(): pytest.importorskip('tensorflow') module_name = "Ololols" 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 = 'forward_jit_gpu' # must be different from CPU names backward_ast = pystencils.create_kernel(backward_assignments, target) backward_ast.function_name = 'backward_jit_gpu' module = TensorflowModule(module_name, [forward_ast, backward_ast]) lib = pystencils_autodiff.tensorflow_jit.compile_sources_and_load( [], [str(module)]) assert 'call_forward_jit_gpu' in dir(lib) assert 'call_backward_jit_gpu' in dir(lib) module = TensorflowModule(module_name, [forward_ast, backward_ast]) lib = module.compile() assert 'call_forward_jit_gpu' in dir(lib) assert 'call_backward_jit_gpu' in dir(lib) file_name = pystencils_autodiff.tensorflow_jit.compile_sources_and_load( [], [str(module)], compile_only=True) print(file_name) assert exists(file_name)