From 6a5b3fef13890acd1a9413a838c371be344c3e4e Mon Sep 17 00:00:00 2001
From: Stephan Seitz <stephan.seitz@fau.de>
Date: Wed, 26 Feb 2020 13:58:23 +0100
Subject: [PATCH] Try to make module independent from pystencils
---
src/pyronn_torch/__init__.py | 2 +-
src/pyronn_torch/codegen.py | 316 ++++++++++++++++++++---------------
tests/test_projection.py | 5 +-
3 files changed, 189 insertions(+), 134 deletions(-)
diff --git a/src/pyronn_torch/__init__.py b/src/pyronn_torch/__init__.py
index 57356e9..9e1ab5a 100644
--- a/src/pyronn_torch/__init__.py
+++ b/src/pyronn_torch/__init__.py
@@ -24,6 +24,6 @@ except Exception as e:
import warnings
warnings.warn(str(e))
import pyronn_torch.codegen
- cpp_extension = pyronn_torch.codegen.generate_shared_object()
+ cpp_extension = pyronn_torch.codegen.compile_shared_object()
__all__ = ['ConeBeamProjector', 'cpp_extension']
diff --git a/src/pyronn_torch/codegen.py b/src/pyronn_torch/codegen.py
index 4783c70..b7f11ef 100644
--- a/src/pyronn_torch/codegen.py
+++ b/src/pyronn_torch/codegen.py
@@ -9,142 +9,149 @@
import argparse
+import tempfile
from glob import glob
from os import makedirs
from os.path import basename, dirname, join
from shutil import copyfile, copytree, rmtree
-import pystencils
-from pystencils.astnodes import Block
-from pystencils.cpu.cpujit import get_cache_config
-from pystencils.data_types import TypedSymbol, create_type
-from pystencils.kernelparameters import FieldPointerSymbol, FieldShapeSymbol
-from pystencils_autodiff.backends.astnodes import TorchModule
-from pystencils_autodiff.framework_integration.astnodes import (
- CustomFunctionCall, WrapperFunction)
-from pystencils_autodiff.framework_integration.printer import \
- FrameworkIntegrationPrinter
-
-volume = pystencils.fields('volume: float32[3d]')
-projection = pystencils.fields('projection: float32[3d]')
-projection_matrices = pystencils.fields('matrices: float32[3d]')
-inv_matrices = pystencils.fields('inv_matrices: float32[3d]')
-source_points = pystencils.fields('source_points: float32[1d]')
-volume_slice = pystencils.fields('volume_slice: float32[2d]')
-projections_1d = pystencils.fields('projections_1d: float32[2d]')
-ray_vectors = pystencils.fields('ray_vectors: float32[2d]')
-
-FUNCTIONS = {
- 'Cone_Backprojection3D_Kernel_Launcher': CustomFunctionCall('Cone_Backprojection3D_Kernel_Launcher',
- FieldPointerSymbol(projection.name, projection.dtype, const=True),
- FieldPointerSymbol(volume.name, volume.dtype, const=False),
- FieldPointerSymbol(projection_matrices.name,
- projection_matrices.dtype, const=True),
- FieldShapeSymbol(['matrices'], 0),
- *[FieldShapeSymbol(['volume'], i) for i in range(2, -1, -1)],
- TypedSymbol('volume_spacing_x', create_type('float32'), const=True),
- TypedSymbol('volume_spacing_y', create_type('float32'), const=True),
- TypedSymbol('volume_spacing_z', create_type('float32'), const=True),
- TypedSymbol('volume_origin_x', create_type('float32'), const=True),
- TypedSymbol('volume_origin_y', create_type('float32'), const=True),
- TypedSymbol('volume_origin_z', create_type('float32'), const=True),
- *[FieldShapeSymbol(['projection'], i) for i in range(2, 0, -1)],
- TypedSymbol('projection_multiplier', create_type('float32'), const=True),
- fields_accessed=[volume, projection, projection_matrices], custom_signature="""
-void Cone_Backprojection3D_Kernel_Launcher(const float *sinogram_ptr, float *out, const float *projection_matrices, const int number_of_projections,
- const int volume_width, const int volume_height, const int volume_depth,
- const float volume_spacing_x, const float volume_spacing_y, const float volume_spacing_z,
- const float volume_origin_x, const float volume_origin_y, const float volume_origin_z,
- const int detector_width, const int detector_height, const float projection_multiplier);
-"""), # noqa
-'Cone_Projection_Kernel_Launcher': CustomFunctionCall('Cone_Projection_Kernel_Launcher',
- FieldPointerSymbol(volume.name, volume.dtype, const=True),
- FieldPointerSymbol(projection.name, projection.dtype, const=False),
- FieldPointerSymbol(inv_matrices.name,
- inv_matrices.dtype, const=True),
- FieldPointerSymbol(source_points.name,
- source_points.dtype, const=True),
- FieldShapeSymbol([source_points.name], 0),
- *[FieldShapeSymbol(['volume'], i) for i in range(2, -1, -1)],
- TypedSymbol('volume_spacing_x', create_type('float32'), const=True),
- TypedSymbol('volume_spacing_y', create_type('float32'), const=True),
- TypedSymbol('volume_spacing_z', create_type('float32'), const=True),
- *[FieldShapeSymbol(['projection'], i) for i in range(2, 0, -1)],
- TypedSymbol('step_size', create_type('float32'), const=True),
- fields_accessed=[volume, projection, inv_matrices, source_points], custom_signature="""
-void Cone_Projection_Kernel_Launcher(const float* volume_ptr, float *out, const float *inv_AR_matrix, const float *src_points,
- const int number_of_projections, const int volume_width, const int volume_height, const int volume_depth,
- const float volume_spacing_x, const float volume_spacing_y, const float volume_spacing_z,
- const int detector_width, const int detector_height, const float step_size);
-"""), # noqa
-'Cone_Projection_Kernel_Tex_Interp_Launcher': CustomFunctionCall('Cone_Projection_Kernel_Tex_Interp_Launcher',
- FieldPointerSymbol(volume.name, volume.dtype, const=True),
- FieldPointerSymbol(projection.name, projection.dtype, const=False),
- FieldPointerSymbol(inv_matrices.name,
- inv_matrices.dtype, const=True),
- FieldPointerSymbol(source_points.name,
- source_points.dtype, const=True),
- FieldShapeSymbol([source_points.name], 0),
- *[FieldShapeSymbol(['volume'], i) for i in range(2, -1, -1)],
- TypedSymbol('volume_spacing_x', create_type('float32'), const=True),
- TypedSymbol('volume_spacing_y', create_type('float32'), const=True),
- TypedSymbol('volume_spacing_z', create_type('float32'), const=True),
- *[FieldShapeSymbol(['projection'], i) for i in range(2, 0, -1)],
- TypedSymbol('step_size', create_type('float32'), const=True),
- fields_accessed=[volume, projection, inv_matrices, source_points], custom_signature="""
-void Cone_Projection_Kernel_Tex_Interp_Launcher(
- const float *__restrict__ volume_ptr, float *out,
- const float *inv_AR_matrix, const float *src_points,
- const int number_of_projections, const int volume_width,
- const int volume_height, const int volume_depth,
- const float volume_spacing_x, const float volume_spacing_y,
- const float volume_spacing_z, const int detector_width,
- const int detector_height, const float step_size);
-"""), # noqa
-'Parallel_Projection2D_Kernel_Launcher': CustomFunctionCall('Parallel_Projection2D_Kernel_Launcher',
- FieldPointerSymbol(volume_slice.name, volume_slice.dtype, const=True),
- FieldPointerSymbol(projections_1d.name, projections_1d.dtype, const=False),
- FieldPointerSymbol(ray_vectors.name,
- ray_vectors.dtype, const=True),
- FieldShapeSymbol([ray_vectors.name], 0),
- *[FieldShapeSymbol(['volume_slice'], i) for i in range(1, -1, -1)],
- TypedSymbol('volume_spacing_x', create_type('float32'), const=True),
- TypedSymbol('volume_spacing_y', create_type('float32'), const=True),
- TypedSymbol('volume_origin_x', create_type('float32'), const=True),
- TypedSymbol('volume_origin_y', create_type('float32'), const=True),
- FieldShapeSymbol([projections_1d.name], 1),
- TypedSymbol('detector_spacing', create_type('float32'), const=True),
- TypedSymbol('detector_origin', create_type('float32'), const=True),
- fields_accessed=[volume_slice, projections_1d, ray_vectors], custom_signature="""
-void Parallel_Projection2D_Kernel_Launcher(
- const float *volume_ptr, float *out, const float *ray_vectors,
- const int number_of_projections, const int volume_width,
- const int volume_height, const float volume_spacing_x,
- const float volume_spacing_y, const float volume_origin_x,
- const float volume_origin_y, const int detector_size,
- const float detector_spacing, const float detector_origin);
-"""), # noqa
-'Parallel_Backprojection2D_Kernel_Launcher': CustomFunctionCall('Parallel_Backprojection2D_Kernel_Launcher',
- FieldPointerSymbol(projections_1d.name, projections_1d.dtype, const=True),
- FieldPointerSymbol(volume_slice.name, volume_slice.dtype, const=False),
- FieldPointerSymbol(ray_vectors.name,
- ray_vectors.dtype, const=True),
- FieldShapeSymbol([ray_vectors.name], 0),
- *[FieldShapeSymbol(['volume_slice'], i) for i in range(1, -1, -1)],
- TypedSymbol('volume_spacing_x', create_type('float32'), const=True),
- TypedSymbol('volume_spacing_y', create_type('float32'), const=True),
- TypedSymbol('volume_origin_x', create_type('float32'), const=True),
- TypedSymbol('volume_origin_y', create_type('float32'), const=True),
- FieldShapeSymbol([projections_1d.name], 1),
- TypedSymbol('detector_spacing', create_type('float32'), const=True),
- TypedSymbol('detector_origin', create_type('float32'), const=True),
- fields_accessed=[volume_slice, projections_1d, ray_vectors], custom_signature="""
-void Parallel_Backprojection2D_Kernel_Launcher(const float *sinogram_ptr, float *out, const float *ray_vectors, const int number_of_projections,
- const int volume_width, const int volume_height, const float volume_spacing_x, const float volume_spacing_y,
- const float volume_origin_x, const float volume_origin_y,
- const int detector_size, const float detector_spacing, const float detector_origin);
-"""), # noqa
- }
+try:
+ import pystencils
+ from pystencils.cpu.cpujit import get_cache_config
+ from pystencils.data_types import TypedSymbol, create_type
+ from pystencils.kernelparameters import FieldPointerSymbol, FieldShapeSymbol
+ from pystencils_autodiff.backends.astnodes import TorchModule
+ from pystencils_autodiff.framework_integration.astnodes import \
+ CustomFunctionCall
+ from pystencils_autodiff.framework_integration.printer import \
+ FrameworkIntegrationPrinter
+
+ volume = pystencils.fields('volume: float32[3d]')
+ projection = pystencils.fields('projection: float32[3d]')
+ projection_matrices = pystencils.fields('matrices: float32[3d]')
+ inv_matrices = pystencils.fields('inv_matrices: float32[3d]')
+ source_points = pystencils.fields('source_points: float32[1d]')
+ volume_slice = pystencils.fields('volume_slice: float32[2d]')
+ projections_1d = pystencils.fields('projections_1d: float32[2d]')
+ ray_vectors = pystencils.fields('ray_vectors: float32[2d]')
+
+ FUNCTIONS = {
+ 'Cone_Backprojection3D_Kernel_Launcher': CustomFunctionCall('Cone_Backprojection3D_Kernel_Launcher',
+ FieldPointerSymbol(projection.name, projection.dtype, const=True),
+ FieldPointerSymbol(volume.name, volume.dtype, const=False),
+ FieldPointerSymbol(projection_matrices.name,
+ projection_matrices.dtype, const=True),
+ FieldShapeSymbol(['matrices'], 0),
+ *[FieldShapeSymbol(['volume'], i) for i in range(2, -1, -1)],
+ TypedSymbol('volume_spacing_x', create_type('float32'), const=True),
+ TypedSymbol('volume_spacing_y', create_type('float32'), const=True),
+ TypedSymbol('volume_spacing_z', create_type('float32'), const=True),
+ TypedSymbol('volume_origin_x', create_type('float32'), const=True),
+ TypedSymbol('volume_origin_y', create_type('float32'), const=True),
+ TypedSymbol('volume_origin_z', create_type('float32'), const=True),
+ *[FieldShapeSymbol(['projection'], i) for i in range(2, 0, -1)],
+ TypedSymbol('projection_multiplier',
+ create_type('float32'), const=True),
+ fields_accessed=[volume, projection, projection_matrices], custom_signature="""
+ void Cone_Backprojection3D_Kernel_Launcher(const float *sinogram_ptr, float *out, const float *projection_matrices, const int number_of_projections,
+ const int volume_width, const int volume_height, const int volume_depth,
+ const float volume_spacing_x, const float volume_spacing_y, const float volume_spacing_z,
+ const float volume_origin_x, const float volume_origin_y, const float volume_origin_z,
+ const int detector_width, const int detector_height, const float projection_multiplier);
+ """), # noqa
+ 'Cone_Projection_Kernel_Launcher': CustomFunctionCall('Cone_Projection_Kernel_Launcher',
+ FieldPointerSymbol(volume.name, volume.dtype, const=True),
+ FieldPointerSymbol(projection.name, projection.dtype, const=False),
+ FieldPointerSymbol(inv_matrices.name,
+ inv_matrices.dtype, const=True),
+ FieldPointerSymbol(source_points.name,
+ source_points.dtype, const=True),
+ FieldShapeSymbol([source_points.name], 0),
+ *[FieldShapeSymbol(['volume'], i) for i in range(2, -1, -1)],
+ TypedSymbol('volume_spacing_x', create_type('float32'), const=True),
+ TypedSymbol('volume_spacing_y', create_type('float32'), const=True),
+ TypedSymbol('volume_spacing_z', create_type('float32'), const=True),
+ *[FieldShapeSymbol(['projection'], i) for i in range(2, 0, -1)],
+ TypedSymbol('step_size', create_type('float32'), const=True),
+ fields_accessed=[volume, projection, inv_matrices, source_points], custom_signature="""
+ void Cone_Projection_Kernel_Launcher(const float* volume_ptr, float *out, const float *inv_AR_matrix, const float *src_points,
+ const int number_of_projections, const int volume_width, const int volume_height, const int volume_depth,
+ const float volume_spacing_x, const float volume_spacing_y, const float volume_spacing_z,
+ const int detector_width, const int detector_height, const float step_size);
+ """), # noqa
+ 'Cone_Projection_Kernel_Tex_Interp_Launcher': CustomFunctionCall('Cone_Projection_Kernel_Tex_Interp_Launcher',
+ FieldPointerSymbol(volume.name, volume.dtype, const=True),
+ FieldPointerSymbol(projection.name, projection.dtype, const=False),
+ FieldPointerSymbol(inv_matrices.name,
+ inv_matrices.dtype, const=True),
+ FieldPointerSymbol(source_points.name,
+ source_points.dtype, const=True),
+ FieldShapeSymbol([source_points.name], 0),
+ *[FieldShapeSymbol(['volume'], i) for i in range(2, -1, -1)],
+ TypedSymbol('volume_spacing_x', create_type('float32'), const=True),
+ TypedSymbol('volume_spacing_y', create_type('float32'), const=True),
+ TypedSymbol('volume_spacing_z', create_type('float32'), const=True),
+ *[FieldShapeSymbol(['projection'], i) for i in range(2, 0, -1)],
+ TypedSymbol('step_size', create_type('float32'), const=True),
+ fields_accessed=[volume, projection, inv_matrices, source_points], custom_signature="""
+ void Cone_Projection_Kernel_Tex_Interp_Launcher(
+ const float *__restrict__ volume_ptr, float *out,
+ const float *inv_AR_matrix, const float *src_points,
+ const int number_of_projections, const int volume_width,
+ const int volume_height, const int volume_depth,
+ const float volume_spacing_x, const float volume_spacing_y,
+ const float volume_spacing_z, const int detector_width,
+ const int detector_height, const float step_size);
+ """), # noqa
+ 'Parallel_Projection2D_Kernel_Launcher': CustomFunctionCall('Parallel_Projection2D_Kernel_Launcher',
+ FieldPointerSymbol(volume_slice.name, volume_slice.dtype, const=True),
+ FieldPointerSymbol(projections_1d.name,
+ projections_1d.dtype, const=False),
+ FieldPointerSymbol(ray_vectors.name,
+ ray_vectors.dtype, const=True),
+ FieldShapeSymbol([ray_vectors.name], 0),
+ *[FieldShapeSymbol(['volume_slice'], i) for i in range(1, -1, -1)],
+ TypedSymbol('volume_spacing_x', create_type('float32'), const=True),
+ TypedSymbol('volume_spacing_y', create_type('float32'), const=True),
+ TypedSymbol('volume_origin_x', create_type('float32'), const=True),
+ TypedSymbol('volume_origin_y', create_type('float32'), const=True),
+ FieldShapeSymbol([projections_1d.name], 1),
+ TypedSymbol('detector_spacing', create_type('float32'), const=True),
+ TypedSymbol('detector_origin', create_type('float32'), const=True),
+ fields_accessed=[volume_slice, projections_1d, ray_vectors], custom_signature="""
+ void Parallel_Projection2D_Kernel_Launcher(
+ const float *volume_ptr, float *out, const float *ray_vectors,
+ const int number_of_projections, const int volume_width,
+ const int volume_height, const float volume_spacing_x,
+ const float volume_spacing_y, const float volume_origin_x,
+ const float volume_origin_y, const int detector_size,
+ const float detector_spacing, const float detector_origin);
+ """), # noqa
+ 'Parallel_Backprojection2D_Kernel_Launcher': CustomFunctionCall('Parallel_Backprojection2D_Kernel_Launcher',
+ FieldPointerSymbol(projections_1d.name,
+ projections_1d.dtype, const=True),
+ FieldPointerSymbol(volume_slice.name, volume_slice.dtype, const=False),
+ FieldPointerSymbol(ray_vectors.name,
+ ray_vectors.dtype, const=True),
+ FieldShapeSymbol([ray_vectors.name], 0),
+ *[FieldShapeSymbol(['volume_slice'], i) for i in range(1, -1, -1)],
+ TypedSymbol('volume_spacing_x', create_type('float32'), const=True),
+ TypedSymbol('volume_spacing_y', create_type('float32'), const=True),
+ TypedSymbol('volume_origin_x', create_type('float32'), const=True),
+ TypedSymbol('volume_origin_y', create_type('float32'), const=True),
+ FieldShapeSymbol([projections_1d.name], 1),
+ TypedSymbol('detector_spacing', create_type('float32'), const=True),
+ TypedSymbol('detector_origin', create_type('float32'), const=True),
+ fields_accessed=[volume_slice, projections_1d, ray_vectors], custom_signature="""
+ void Parallel_Backprojection2D_Kernel_Launcher(const float *sinogram_ptr, float *out, const float *ray_vectors, const int number_of_projections,
+ const int volume_width, const int volume_height, const float volume_spacing_x, const float volume_spacing_y,
+ const float volume_origin_x, const float volume_origin_y,
+ const int detector_size, const float detector_spacing, const float detector_origin);
+ """), # noqa
+ }
+except Exception as e:
+ import warnings
+ warnings.warn(e)
def generate_shared_object(output_folder=None, source_files=None, show_code=False):
@@ -189,6 +196,50 @@ def generate_shared_object(output_folder=None, source_files=None, show_code=Fals
return extension
+def compile_shared_object(output_folder=None, source_files=None):
+ """
+ Same as generate_shared_object without pystencils dependency
+ """
+ object_cache = tempfile.gettempdir()
+ module_name = 'pyronn_torch_cpp'
+
+ if not output_folder:
+ output_folder = dirname(__file__)
+
+ if not source_files:
+ source_files = glob(join(dirname(__file__), 'PYRO-NN-Layers', '*.cu.cc'))
+
+ generated_file = join(dirname(__file__), 'pyronn_torch.cpp')
+
+ cuda_sources = []
+ makedirs(join(object_cache, module_name), exist_ok=True)
+ rmtree(join(object_cache, module_name, 'helper_headers'), ignore_errors=True)
+ copytree(join(dirname(__file__), 'PYRO-NN-Layers', 'helper_headers'),
+ join(object_cache, module_name, 'helper_headers'))
+
+ for s in source_files:
+ dst = join(object_cache, module_name, basename(s).replace('.cu.cc', '.cu'))
+ copyfile(s, dst) # Torch only accepts *.cu as CUDA
+ cuda_sources.append(dst)
+
+ # extension = module.compile(extra_source_files=cuda_sources,
+ # extra_cuda_flags=['-arch=sm_35'],
+ # with_cuda=True,
+ # compile_module_name=module_name)
+
+ from torch.utils.cpp_extension import load
+ torch_extension = load(module_name,
+ [generated_file] + cuda_sources,
+ with_cuda=True,
+ extra_cflags=['--std=c++14'],
+ extra_cuda_cflags=['-std=c++14', '-arch=sm_35'],
+ build_directory=join(object_cache, module_name),
+ extra_include_paths=[])
+ shared_object_file = join(object_cache, module_name, module_name + '.so')
+ copyfile(shared_object_file, join(output_folder, module_name + '.so'))
+ return torch_extension
+
+
def main():
parser = argparse.ArgumentParser()
@@ -196,7 +247,8 @@ def main():
parser.add_argument('--source-files', default=None)
args = parser.parse_args()
- generate_shared_object(args.output_folder, args.source_files, show_code=True)
+ # generate_shared_object(args.output_folder, args.source_files, show_code=True)
+ compile_shared_object(args.output_folder, args.source_files)
if __name__ == '__main__':
diff --git a/tests/test_projection.py b/tests/test_projection.py
index 6fac28f..fad1a3a 100644
--- a/tests/test_projection.py
+++ b/tests/test_projection.py
@@ -6,6 +6,8 @@
"""
"""
+import pytest
+
import pyronn_torch
@@ -13,7 +15,8 @@ def test_init():
assert pyronn_torch.cpp_extension
-def test_projection():
+@pytest.mark.parametrize('with_texture', ('with_texture', False))
+def test_projection(with_texture):
projector = pyronn_torch.ConeBeamProjector.from_conrad_config()
volume = projector.new_volume_tensor()
--
GitLab