CUDA kernel function that handles Complex Tensors

eduardo4jesus: <p><strong>How to write a CUDA kernel that handles complex arguments?</strong></p>
<hr>
<p>I have been writing a CUDA extension having the input obtained by <code>torch.view_as_real</code> function, which takes a <code>cfloat</code> and represents an extra dim of <code>float</code> of size two (real and imaginary components).</p>
<p>Here is a code snippet:</p>
<pre><code class="lang-auto">import torch
from my_cuda_extension import multiplication_complex 

cuda = torch.device('cuda')
x = torch.view_as_real(torch.rand(size=(1, 1, 4, 4), dtype=torch.cfloat, device=cuda)*10)
h = torch.view_as_real(torch.rand(size=(1, 1, 4, 4), dtype=torch.cfloat, device=cuda)*10)

multiplication_complex(x, h)
</code></pre>
<p>This is the code I currently have inside the kernel.</p>
<pre><code class="lang-auto">template &lt;typename scalar_t&gt;
__device__ __forceinline__ void multiplication_complex(
    scalar_t a, scalar_t b,
    scalar_t c, scalar_t d,
    scalar_t* out_re, scalar_t* out_im) {
    *out_re += a*c - b*d;
    *out_im += a*d + b*c;
}

template &lt;typename scalar_t&gt;
__global__ void multiplication_cuda_kernel(
    const torch::PackedTensorAccessor32&lt;scalar_t, 5, torch::RestrictPtrTraits&gt; x,
    const torch::PackedTensorAccessor32&lt;scalar_t, 5, torch::RestrictPtrTraits&gt; h,
    torch::PackedTensorAccessor32&lt;scalar_t, 5, torch::RestrictPtrTraits&gt; out,
    const int THREADS, const int C, const int W, const int PLANE_SIZE){

    const int b = blockIdx.x;
    const int f = blockIdx.y;

    const int cell_idx = blockIdx.z * THREADS + threadIdx.x;
    if (cell_idx &gt;= PLANE_SIZE) return;

    const int i = cell_idx/W;
    const int j = fmod(cell_idx, W);

    scalar_t out_re = 0.0;
    scalar_t out_im = 0.0;

    for (int c = 0; c &lt; C; ++c) {

        const scalar_t x_re = x[b][c][i][j][0];
        const scalar_t x_im = x[b][c][i][j][1];

        const scalar_t h_re = h[f][c][i][j][0];
        const scalar_t h_im = h[f][c][i][j][1];

        multiplication_complex(x_re, x_im, h_re, h_im, &amp;out_re, &amp;out_im);
    }

    out[b][f][i][j][0] = out_re;
    out[b][f][i][j][1] = out_im;
}
</code></pre>
<p>Which is dispatched by</p>
<pre><code class="lang-auto">printf("About to DISPATCH\n");
AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "multiplication_complex_cuda",
  ([&amp;] {
      multiplication_cuda_kernel&lt;scalar_t&gt;&lt;&lt;&lt;GRID_SIZE, THREADS&gt;&gt;&gt;(
      x.packed_accessor32&lt;scalar_t, 5, torch::RestrictPtrTraits&gt;(),
      h.packed_accessor32&lt;scalar_t, 5, torch::RestrictPtrTraits&gt;(),
      out.packed_accessor32&lt;scalar_t, 5, torch::RestrictPtrTraits&gt;(),
      THREADS, C, W, PLANE_SIZE);
   })
);
</code></pre>
<hr>
<p><strong>I want to</strong> be able to use complex numbers in a more straight-forward manner and use <code>AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES</code> instead. Thus, avoiding <code>torch.view_as_real</code> and calling the function as</p>
<pre><code class="lang-auto">import torch
from my_cuda_extension import multiplication_complex 

cuda = torch.device('cuda')
x = torch.rand(size=(1, 1, 4, 4), dtype=torch.cfloat, device=cuda)*10
h = torch.rand(size=(1, 1, 4, 4), dtype=torch.cfloat, device=cuda)*10

multiplication_complex(x, h)
</code></pre>
<p>The new dispatch should look like:</p>
<pre><code class="lang-auto">printf("About to DISPATCH\n");
 AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(x.scalar_type(), "multiplication_complex_cuda",
    ([&amp;] {
        multiplication_cuda_kernel&lt;scalar_t&gt;&lt;&lt;&lt;GRID_SIZE, THREADS&gt;&gt;&gt;(
        x.packed_accessor32&lt;scalar_t, 4, torch::RestrictPtrTraits&gt;(),
        h.packed_accessor32&lt;scalar_t, 4, torch::RestrictPtrTraits&gt;(),
        out.packed_accessor32&lt;scalar_t, 4, torch::RestrictPtrTraits&gt;(),
        THREADS, C, W, PLANE_SIZE);
    }));
</code></pre>
<p><strong>BUT</strong> I am not able to write a <code>multiplication_cuda_kernel</code> function that gets called.</p>
<pre><code class="lang-auto">template &lt;typename scalar_t&gt;
__global__ void multiplication_cuda_kernel(
    const torch::PackedTensorAccessor32&lt;scalar_t, 4, torch::RestrictPtrTraits&gt; x,
    const torch::PackedTensorAccessor32&lt;scalar_t, 4, torch::RestrictPtrTraits&gt; h,
    torch::PackedTensorAccessor32&lt;scalar_t, 4, torch::RestrictPtrTraits&gt; out,
    const int THREADS, const int C, const int W, const int PLANE_SIZE
){

    const int b = blockIdx.x; // Image position in Batch
    const int f = blockIdx.y; // Filter position

    const int cell_idx = blockIdx.z * THREADS + threadIdx.x; // data point/pixel/cell index in h x w plane
    if (cell_idx &gt;= PLANE_SIZE) return;

    if (threadIdx.x == 0) printf("Hello Block %d\n", blockIdx.x);
}
</code></pre>
<p>However, this <code>Hello Block</code> is not being printed, and so the <code>multiplication_cuda_kernel</code> function is not being called.</p>
<hr>
<p>I have been observing this file as an example (<a href="https://github.com/pytorch/pytorch/blob/213ac4e59cf5ce4c53780cec3f96dbb82aae1f4d/aten/src/ATen/native/cuda/UnaryComplexKernels.cu#L21" rel="noopener nofollow ugc">pytorch/aten/src/ATen/native/cuda/UnaryComplexKernels.cu</a>), in particular the snippet below.</p>
<pre><code class="lang-auto">template&lt;typename T&gt;
__host__ __device__ static inline c10::complex&lt;T&gt; angle_wrapper(c10::complex&lt;T&gt; v) {
  return std::arg(v);
}

void angle_kernel_cuda(TensorIteratorBase&amp; iter) {
  AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(iter.common_dtype(), "angle_cuda", [&amp;]() {
    gpu_kernel(iter, []GPU_LAMBDA(scalar_t a) -&gt; scalar_t {
      return angle_wrapper(a);
    });
  });
}
</code></pre>
<p>However, none of my attempts inspired by the snippet above were able to successfully use a kernel function that was dispatched for the complex dtype. Any help on getting the <code>multiplication_cuda_kernel</code> function called and how to implement it with a complex scalar_t would be very appreciated.</p>
<p>Thank you for your time.</p>

eduardo4jesus: <p>Wondering if <a class="mention" href="/u/tom">@tom</a> would be available to shed some light on this.</p>

tom: <p>If you want people to look at it, I would think that providing a complete cpp extension or ideally a Python script that loads (CUDA) csrc with something like<br>
<code>test_ext = torch.utils.cpp_extension.load_inline("test_ext", "", csrc, is_python_module=True, verbose=True)</code> and runs it would lower the burden of getting your snippets to run. This way, one can dive right into the matter without (re-) writing the glue code.</p>
<p>I’m a bit doubtful about the code being able to compile but not run the function, something seems to be funny about that, but I’m hoping that you’d show a complete C++ snippet to look at. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>
<p>Best regards</p>
<p>Thomas</p>

eduardo4jesus: <p>Hello Thomas, thank you for your feedback.</p>
<p>I just uploaded the complete code <a href="https://github.com/eduardo4jesus/Complex-Multiplication" rel="noopener nofollow ugc">here</a> for people interested into taking a look.</p>
<p>In the link above I had the current working version of this extension. And I want to change it so that I work with the complex numbers in a “cleaner” manner by</p>
<ol>
<li>Being able to call the <code>complex_multiplication(x, h)</code> with <code>cfloat</code> types.</li>
<li>Handle complex types (somehow) with scalar_t template on the cuda kernel. Something similar to the <code>angle_wrapper</code> function from <a href="https://github.com/pytorch/pytorch/blob/213ac4e59cf5ce4c53780cec3f96dbb82aae1f4d/aten/src/ATen/native/cuda/UnaryComplexKernels.cu#L21" rel="noopener nofollow ugc">Pytorch</a> mentioned above.</li>
</ol>
<p>Eduardo.</p>

tom: <p>If you allow me to re-post the code here directly, the following seems to work for me for a PyTorch git checkout from today:</p>
<pre><code class="lang-auto">import torch
import torch.utils.cpp_extension

csrc = """
#include &lt;torch/extension.h&gt;
#include &lt;cuda.h&gt;
#include &lt;cuda_runtime.h&gt;


template &lt;typename scalar_t&gt;
__device__ __forceinline__ void elementwise_operation(
        scalar_t a,
        scalar_t b,
        scalar_t* out
    ) {
    *out += a*b;
}

template &lt;typename scalar_t&gt;
__global__ void complex_multiplication_cuda_kernel_v1(
        const torch::PackedTensorAccessor32&lt;scalar_t, 4, torch::RestrictPtrTraits&gt; x,
        const torch::PackedTensorAccessor32&lt;scalar_t, 4, torch::RestrictPtrTraits&gt; h,
        torch::PackedTensorAccessor32&lt;scalar_t, 4, torch::RestrictPtrTraits&gt; out,
        const int THREADS, const int C, const int W, const int PLANE_SIZE
    ){

    const int b = blockIdx.x; // Image position in Batch
    const int f = blockIdx.y; // Filter position

    const int cell_idx = blockIdx.z * THREADS + threadIdx.x; // data point/pixel/cell index in h x w plane
    if (cell_idx &gt;= PLANE_SIZE) return;

    const int i = cell_idx/W;
    const int j = fmod(cell_idx, W);

    scalar_t out_ = 0.0;

    /****************************************************************
     * Dimensions should be 
     * x   -&gt; (B, C, H, W) ~&gt; Each b is size C*H*W*I
     * h   -&gt; (F, C, H, W) ~&gt;      f is  ""
     * out -&gt; (B, F, H, W) ~&gt;      b is  ""  F*H*W*I
     ****************************************************************/
    for (int c = 0; c &lt; C; ++c) {

        const scalar_t x_ = x[b][c][i][j];

        const scalar_t h_ = h[f][c][i][j];

        elementwise_operation(x_, h_, &amp;out_);
    }

    out[b][f][i][j] = out_;
}

/**
 * Multiplies two tensors of Complex Tensors
 * @param x
 * @param h
 * @param output
 */
at::Tensor complex_multiplication_cuda_v1(at::Tensor x, at::Tensor h) {
    const int THREADS = 1024;

    const int B = x.size(0);
    const int F = h.size(0);
    const int C = x.size(1);
    const int H = x.size(2);
    const int W = x.size(3);
    const int PLANE_SIZE = H*W;

    const auto Z = (H*W + THREADS - 1)/THREADS;
    const dim3 GRID_SIZE(B, F, Z);

    auto output = torch::zeros(
        {B, F, H, W},
        x.options()
        );

    AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(x.scalar_type(), "complex_multiplication_cuda_v1",
        ([&amp;] {
            complex_multiplication_cuda_kernel_v1&lt;scalar_t&gt;&lt;&lt;&lt;GRID_SIZE, THREADS&gt;&gt;&gt;(
                x.packed_accessor32&lt;scalar_t, 4, torch::RestrictPtrTraits&gt;(),
                h.packed_accessor32&lt;scalar_t, 4, torch::RestrictPtrTraits&gt;(),
                output.packed_accessor32&lt;scalar_t, 4, torch::RestrictPtrTraits&gt;(),
                THREADS, C, W, PLANE_SIZE
            );
        })
    );

    return output;
}

/**
 * Multiplies two tensors of Complex Tensors
 * @param x
 * @param h
 * @return ouput
 */
at::Tensor complex_multiplication(at::Tensor x, at::Tensor h) {
    printf("CPP: my_pytorch_extensions_cuda called\\n");
    return complex_multiplication_cuda_v1(x, h);
}

/********************************
 * Binding Functions
 ********************************/

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("complex_multiplication", &amp;complex_multiplication, "CUDA kernel for multiplication of complex tensors");
}
"""

test_ext = torch.utils.cpp_extension.load_inline("test_ext", "", csrc, is_python_module=True, verbose=True)

torch.set_printoptions(precision=2, linewidth=120)
cuda = torch.device('cuda')
x = torch.rand((1, 1, 4, 4), dtype=torch.cfloat, device=cuda)*10
h = torch.rand((1, 1, 4, 4), dtype=torch.cfloat, device=cuda)*10

output = test_ext.complex_multiplication(x, h)

print((x[:, None] * h[None, :]).sum(2))
print(output)
</code></pre>
<p>This gives me the same complex tensor twice, which seemed your aim.<br>
What does (and should) happen behind the scenes is that you’re using PyTorch’s <code>c10::complex&lt;float&gt;</code>/<code>c10::complex&lt;double&gt;</code> and the operator overloads defined for that provide the <code>*</code>, <code>+</code>, and assignment used in the kernel. (I’m using gcc 10 as the compiler, not sure if it matters, PyTorch currently triggers a bad thing in gcc 11.) Of course, in this setup, you should be using c10::complex throughout, and real() and imag() to get and set real and imaginary parts individually. (see the documentation for c10::complex for more details), and they are not a separate dimension.</p>
<p>Note that complex support was incremental, so I would recommend to try with a rather recent PyTorch version (e.g. nightly or testing) if you are not using it already.</p>
<p>Best regards</p>
<p>Thomas</p>

eduardo4jesus: <p>Thank you for your code snippet. I am using <code>torch.__version == 1.10.0</code> and it worked nicely.<br>
I am not sure why the code I had previously didn’t work. It was probably something wrong on how I used the complex types on the kernel functions.</p>
<p>For the sake of example completeness, here are the kernel functions I tested for manipulating the complex numbers. In the example below <code>calculate_multiplication_long</code> is overloaded, which could potentially have a different behaviour from the real types.</p>
<pre><code class="lang-auto">template &lt;typename scalar_t&gt;
__device__ __forceinline__ void calculate_multiplication(
    scalar_t x, scalar_t h, scalar_t* out
    ) {
    *out += x*h;
}

template &lt;typename scalar_t&gt;
__device__ __forceinline__ void calculate_multiplication_long(
    scalar_t x, scalar_t h, scalar_t* out
    ) {
    *out += x*h;
}

template &lt;typename T&gt;
__device__ __forceinline__ void calculate_multiplication_long(
    c10::complex&lt;T&gt; x, c10::complex&lt;T&gt; h, c10::complex&lt;T&gt;* out
    ) {
    const T a=x.real(), b=x.imag(), c=h.real(), d=h.imag();
    out-&gt;real(out-&gt;real() + a*c - b*d);
    out-&gt;imag(out-&gt;imag() + a*d + b*c);
}
</code></pre>

Jakub_Mitura: <p>A bit offtopic - but did you build it with new CUDA 11.6? If yes can you share your build configuration? Becouse I have a CUDA C++ kernel I would like to use in Python Pytorch Code but I have configuration problems.</p>

eduardo4jesus: <p>No, the machine I am using happens to be <code>CUDA Version: 10.2</code>, not <code>11.6</code>.</p>