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_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_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 native_tensorflowop_from_autodiffop(autodiff_obj: pystencils_autodiff.AutoDiffOp,
use_cuda):
import tensorflow as tf
if use_cuda:
forward_ast = autodiff_obj.forward_ast_gpu
backward_ast = autodiff_obj.backward_ast_gpu
else:
forward_ast = autodiff_obj.forward_ast_cpu
backward_ast = autodiff_obj.backward_ast_cpu
op_name = f'{autodiff_obj.op_name.lower()}_hash{_hash(str((autodiff_obj, autodiff_obj.forward_input_fields, autodiff_obj.backward_input_fields)).encode()).hexdigest()}' # noqa
if use_cuda:
op_name += '_cuda'
forward_ast.function_name = op_name + "_forward"
backward_ast.function_name = op_name + "_backward"
module = TensorflowModule(op_name, [forward_ast, backward_ast])
compiled_op = module.compile()
backward_func = getattr(compiled_op, stringcase.snakecase(
stringcase.pascalcase("call_" + backward_ast.function_name)))
grad_fields = [f for f in autodiff_obj.backward_input_fields if f not in autodiff_obj.forward_input_fields]
def gradient_calculation(op, *grad):
if not isinstance(grad, Iterable):
grad = [grad]
return backward_func(**{grad_fields[i].name: g for i, g in enumerate(grad)},
**{autodiff_obj.forward_input_fields[i].name: inp for i, inp in enumerate(op.inputs)
if autodiff_obj.forward_input_fields[i] in backward_ast.fields_accessed})
try:
tf.RegisterGradient(stringcase.pascalcase("call_" + forward_ast.function_name))(
gradient_calculation
)
except Exception:
pass
return getattr(compiled_op, stringcase.snakecase(stringcase.pascalcase("call_" + forward_ast.function_name)))
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)
def test_detect_cpu_vs_cpu():
for target in ('cpu', 'gpu'):
z, y, x = pystencils.fields("z, y, x: [20,40]")
a = sympy.Symbol('a')
assignments = pystencils.AssignmentCollection({
z[0, 0]: x[0, 0] * sympy.log(a * x[0, 0] * y[0, 0])
})
kernel_ast = pystencils.create_kernel(assignments, target=target)
module = TensorflowModule('my_module', [kernel_ast])
print(module)
assert 'DEVICE_' + target.upper() in str(module)