test_torch_native_compilation.py 7.54 KiB
# -*- coding: utf-8 -*-
#
# Copyright © 2019 seitz_local <seitz_local@lmeXX>
#
import os
import subprocess
from os.path import dirname, isfile, join
# TODO: from pystencils.backends.cudabackend import generate_cuda
import appdirs
import jinja2
import numpy as np
import pytest
import pystencils
import pystencils_autodiff
# from pystencils.cpu.kernelcreation import create_kernel
from pystencils.backends.cbackend import generate_c
from pystencils.gpucuda.kernelcreation import create_cuda_kernel
from pystencils_autodiff.backends._torch_native import create_autograd_function, generate_torch
torch = pytest.importorskip('torch')
pytestmark = pytest.mark.skipif(subprocess.call(['ninja', '--v']) != 0,
reason='torch compilation requires ninja')
PROJECT_ROOT = dirname
def read_file(file):
with open(file, 'r') as f:
return f.read()
def test_jit():
"""
Test JIT compilation from example on git@github.com:pytorch/extension-cpp.git
`ninja-build` is required
"""
# os.environ['CUDA_HOME'] = "/usr/local/cuda-10.0"
cpp_file = join(dirname(__file__), 'lltm_cuda.cpp')
cuda_file = join(dirname(__file__), 'lltm_cuda_kernel.cu')
assert isfile(cpp_file)
assert isfile(cuda_file)
from torch.utils.cpp_extension import CUDAExtension
lltm_cuda = CUDAExtension(join(dirname(__file__), 'lltm_cuda'),
[cpp_file, cuda_file])
assert lltm_cuda is not None
def test_torch_native_compilation():
x, y = pystencils.fields('x, y: float32[2d]')
assignments = pystencils.AssignmentCollection({y.center(): x.center()**2},
{})
autodiff = pystencils_autodiff.AutoDiffOp(assignments)
backward_assignments = autodiff.backward_assignments
print(assignments)
print(backward_assignments)
template_string = read_file(
join(
dirname(__file__),
'../../src/pystencils_autodiff/backends/torch_native_cuda.tmpl.cpp'
))
template = jinja2.Template(template_string)
print(template_string)
forward_kernel = create_cuda_kernel(assignments.all_assignments).body
backward_kernel = create_cuda_kernel(
backward_assignments.all_assignments).body
forward_code = generate_c(forward_kernel)
backward_code = generate_c(backward_kernel)
output = template.render(
forward_tensors=[f.name for f in autodiff.forward_fields],
forward_input_tensors=[f.name for f in autodiff.forward_input_fields],
forward_output_tensors=[
f.name for f in autodiff.forward_output_fields
],
backward_tensors=[
f.name
for f in autodiff.backward_fields + autodiff.forward_input_fields
],
backward_input_tensors=[
f.name for f in autodiff.backward_input_fields
],
backward_output_tensors=[
f.name for f in autodiff.backward_output_fields
],
forward_kernel=forward_code,
backward_kernel=backward_code,
dimensions=range(2),
kernel_name="square",
dtype="float")
print(output)
template_string = read_file(
join(
dirname(__file__),
'../../src/pystencils_autodiff/backends/torch_native_cuda.tmpl.cu')
)
template = jinja2.Template(template_string)
print(template_string)
output = template.render(
forward_tensors=[f for f in autodiff.forward_fields],
forward_input_tensors=[f for f in autodiff.forward_input_fields],
forward_output_tensors=[f for f in autodiff.forward_output_fields],
backward_tensors=[
f for f in autodiff.backward_fields + autodiff.forward_input_fields
],
backward_input_tensors=[f for f in autodiff.backward_input_fields],
backward_output_tensors=[f for f in autodiff.backward_output_fields],
forward_kernel=forward_code,
backward_kernel=backward_code,
backward_blocks=str({1, 1, 1}),
backward_threads=str({1, 1, 1}),
forward_blocks=str({1, 1, 1}),
forward_threads=str({1, 1, 1}),
kernel_name="square",
dimensions=range(2))
print(output)
template_string = read_file(
join(
dirname(__file__),
'../../src/pystencils_autodiff/backends/torch_native_cpu.tmpl.cpp')
)
template = jinja2.Template(template_string)
print(template_string)
output = template.render(
forward_tensors=[f.name for f in autodiff.forward_fields],
forward_input_tensors=[f.name for f in autodiff.forward_input_fields],
forward_output_tensors=[
f.name for f in autodiff.forward_output_fields
],
backward_tensors=[
f.name
for f in autodiff.backward_fields + autodiff.forward_input_fields
],
backward_input_tensors=[
f.name for f in autodiff.backward_input_fields
],
backward_output_tensors=[
f.name for f in autodiff.backward_output_fields
],
forward_kernel=forward_code,
backward_kernel=backward_code,
kernel_name="square",
dtype="float",
dimensions=range(2))
print(output)
@pytest.mark.skipif("TRAVIS" in os.environ, reason="Temporary skip")
def test_generate_torch_gpu():
x, y = pystencils.fields('x, y: float32[2d]')
assignments = pystencils.AssignmentCollection({y.center(): x.center()**2},
{})
autodiff = pystencils_autodiff.AutoDiffOp(assignments)
op_cuda = generate_torch(appdirs.user_cache_dir('pystencils'),
autodiff,
is_cuda=True,
dtype=np.float32)
assert op_cuda is not None
def test_generate_torch_cpu():
x, y = pystencils.fields('x, y: float32[2d]')
assignments = pystencils.AssignmentCollection({y.center(): x.center()**2},
{})
autodiff = pystencils_autodiff.AutoDiffOp(assignments)
op_cpp = generate_torch(appdirs.user_cache_dir('pystencils'),
autodiff,
is_cuda=False,
dtype=np.float32)
assert op_cpp is not None
def test_execute_torch():
x, y = pystencils.fields('x, y: float64[32,32]')
assignments = pystencils.AssignmentCollection({y.center(): 5 + x.center()},
{})
autodiff = pystencils_autodiff.AutoDiffOp(assignments)
x_tensor = pystencils_autodiff.torch_tensor_from_field(x, 1, cuda=False)
y_tensor = pystencils_autodiff.torch_tensor_from_field(y, 1, cuda=False)
op_cpp = create_autograd_function(autodiff, {x: x_tensor, y: y_tensor})
foo = op_cpp.forward()
print(foo)
assert op_cpp is not None
@pytest.mark.skipif('NO_GPU_EXECUTION' in os.environ,
reason='Skip GPU execution tests')
def test_execute_torch_gpu():
x, y = pystencils.fields('x, y: float64[32,32]')
assignments = pystencils.AssignmentCollection({y.center(): 5 + x.center()},
{})
autodiff = pystencils_autodiff.AutoDiffOp(assignments)
x_tensor = pystencils_autodiff.torch_tensor_from_field(x, 3, cuda=True)
y_tensor = pystencils_autodiff.torch_tensor_from_field(y, 4, cuda=True)
assert y_tensor.is_cuda
assert torch.cuda.is_available()
op_cuda = create_autograd_function(autodiff, {x: x_tensor, y: y_tensor})
assert op_cuda is not None
rtn = op_cuda.forward()
print(y_tensor)
print(rtn)