From 374d0fd5a54cd84b7f862e79b17a7e4e4b12badb Mon Sep 17 00:00:00 2001
From: Stephan Seitz <stephan.seitz@fau.de>
Date: Wed, 7 Aug 2019 15:20:22 +0200
Subject: [PATCH] Make Torch CUDA compilation pass
---
.../backends/torch_native_cuda.tmpl.cu | 52 ++++++++++++-------
1 file changed, 34 insertions(+), 18 deletions(-)
diff --git a/src/pystencils_autodiff/backends/torch_native_cuda.tmpl.cu b/src/pystencils_autodiff/backends/torch_native_cuda.tmpl.cu
index 5edd28b..bca5d19 100644
--- a/src/pystencils_autodiff/backends/torch_native_cuda.tmpl.cu
+++ b/src/pystencils_autodiff/backends/torch_native_cuda.tmpl.cu
@@ -20,10 +20,12 @@ __global__ void {{ kernel_name }}_cuda_forward_kernel(
{% for i in range(tensor.spatial_dimensions )-%}
int _stride_{{ tensor.name }}_{{ i }} {{- ", " }}
{% endfor -%}
+ {% for i in range(tensor.spatial_dimensions )-%}
+ int _size_{{ tensor.name }}_{{ i }} {{- ", " }}
+ {% endfor -%}
{% endfor -%}
- {% for i in range(forward_output_tensors[0].spatial_dimensions )-%}
- int _size_{{ forward_output_tensors[0] }}_{{ i }} {{- "," if not loop.last }}
- {% endfor %})
+ int _unused
+ )
{
{{forward_kernel}}
}
@@ -36,10 +38,12 @@ __global__ void {{ kernel_name }}_cuda_backward_kernel(
{% for i in range(tensor.spatial_dimensions )-%}
int _stride_{{ tensor.name }}_{{ i }} {{- ", " }}
{% endfor -%}
+ {% for i in range(tensor.spatial_dimensions )-%}
+ int _size_{{ tensor.name }}_{{ i }} {{- ", " }}
+ {% endfor -%}
{% endfor -%}
- {% for i in range(forward_output_tensors[0].spatial_dimensions )-%}
- int _size_{{ forward_output_tensors[0].name }}_{{ i }} {{- "," if not loop.last }}
- {% endfor %})
+ int _unused
+ )
{
{{backward_kernel}}
}
@@ -50,9 +54,14 @@ void {{ kernel_name }}_cuda_forward(
{%- endfor -%})
{
- {% for i in range(forward_output_tensors[0].spatial_dimensions )-%}
- int _size_{{ forward_output_tensors[0].name }}_{{ i }} = {{ forward_output_tensors[0].name }}.size({{ i }});
- {% endfor %}
+ {% for tensor in forward_tensors -%}
+ {% for i in dimensions -%}
+ int _stride_{{tensor}}_{{i}} = {{tensor}}.strides()[{{ i }}];
+ {% endfor -%}
+ {% for i in dimensions -%}
+ int _size_{{tensor}}_{{i}} = {{tensor}}.size({{ i }});
+ {% endfor -%}
+ {% endfor -%}
/*at:: at::device(at::kCUDA).dtype(k{{ dtype }})*/
AT_DISPATCH_FLOATING_TYPES({{ forward_input_tensors[0].name }}.type(), "{{ kernel_name }}_forward_cuda", ([&] {
@@ -63,10 +72,11 @@ void {{ kernel_name }}_cuda_forward(
{% for i in range(tensor.spatial_dimensions) -%}
{{tensor.name}}.strides()[{{ i }}] {{- "," }}
{% endfor -%}
+ {% for i in range(tensor.spatial_dimensions) -%}
+ {{tensor.name}}.size({{ i }}) {{- "," }}
+ {% endfor -%}
{% endfor -%}
- {% for i in range(forward_output_tensors[0].spatial_dimensions) -%}
- {{ forward_output_tensors[0].name }}.size({{ i }}) {{- "," if not loop.last }}
- {% endfor %}
+ 0
);
}));
cudaError_t err = cudaGetLastError();
@@ -82,9 +92,14 @@ void {{ kernel_name }}_cuda_backward(
{%- endfor %})
{
- {% for i in range(backward_output_tensors[0].spatial_dimensions )-%}
- int _size_{{ backward_output_tensors[0].name }}_{{ i }} = {{ backward_output_tensors[0].name }}.size({{ i }});
- {% endfor %}
+ {% for tensor in backward_tensors -%}
+ {% for i in dimensions -%}
+ int _stride_{{tensor}}_{{i}} = {{tensor}}.strides()[{{ i }}];
+ {% endfor -%}
+ {% for i in dimensions -%}
+ int _size_{{tensor}}_{{i}} = {{tensor}}.size({{ i }});
+ {% endfor -%}
+ {% endfor -%}
/*at:: at::device(at::kCUDA).dtype(k{{ dtype }})*/
AT_DISPATCH_FLOATING_TYPES({{ backward_input_tensors[0].name }}.type(), "{{ kernel_name }}_backward_cuda", ([&] {
@@ -95,10 +110,11 @@ void {{ kernel_name }}_cuda_backward(
{% for i in range(tensor.spatial_dimensions )-%}
{{tensor.name}}.strides()[{{ i }}]{{- ", " }}
{% endfor -%}
+ {% for i in range(tensor.spatial_dimensions )-%}
+ {{tensor.name}}.size({{ i }}){{- ", " }}
+ {% endfor -%}
{% endfor -%}
- {% for i in range(backward_output_tensors[0].spatial_dimensions )-%}
- {{ backward_output_tensors[0].name }}.size({{ i }}) {{- "," if not loop.last }}
- {% endfor %}
+ 0
);
}));
--
GitLab