C10::Half float type support for atomicAdd?

YutongZheng: <p>I’m writing a cuda extension for a customized function. It requires an atomicAdd() function during the calculation. Meanwhile, I want to make it possible for half-precision support. for example, in my .cuda file I have something like:</p>
<pre><code class="lang-auto">#include &lt;c10/util/Half.h&gt;
#include "sample.h"

//------------------------------------------------------------------------
// Helpers.

template &lt;class T&gt; struct InternalType;
template &lt;&gt; struct InternalType&lt;double&gt;     { typedef double scalar_t; };
template &lt;&gt; struct InternalType&lt;float&gt;      { typedef float  scalar_t; };
template &lt;&gt; struct InternalType&lt;c10::Half&gt;  { typedef float  scalar_t; };

//------------------------------------------------------------------------
// Generic CUDA implementation.

template &lt;class T&gt; static __global__ void my_kernel(const void* x, void* y)
{
    typedef typename InternalType&lt;T&gt;::scalar_t scalar_t;
    // some operation includes the atomicAdd operation:
    for (...)
    {
        ...
        scalar_t final_value = (scalar_t)(*x);
        atomicAdd(y, (T)final_value);
        ...
    }
}

//------------------------------------------------------------------------
// Kernel select function that's callable from .cpp file. 
template &lt;class T&gt; kernel_spec choose_kernel()
{
    spec = {(void*)my_kernel&lt;T&gt;};
    return spec;
}

//------------------------------------------------------------------------
// Template specializations.

template kernel_spec choose_kernel&lt;double&gt;   ();
template kernel_spec choose_kernel&lt;float&gt;    ();
template kernel_spec choose_kernel&lt;c10::Half&gt;();
</code></pre>
<p>And in my .cpp file, I have something like:</p>
<pre><code class="lang-auto">#include &lt;torch/extension.h&gt;
#include &lt;ATen/cuda/CUDAContext.h&gt;
#include &lt;c10/cuda/CUDAGuard.h&gt;
#include "sample.h"

...

    // Choose CUDA kernel.
    kernel_spec spec;
    AT_DISPATCH_FLOATING_TYPES_AND_HALF(y.scalar_type(), "my_kernel_cuda", [&amp;]
    {
        spec = choose_kernel&lt;scalar_t&gt;();
    });

</code></pre>
<p>But I got error during the compiling:</p>
<pre><code class="lang-auto">RuntimeError: Error building extension 'sample_plugin': [1/3] /usr/local/cuda/bin/nvcc -DTORCH_EXTENSION_NAME=sample_plugin -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1013\" -isystem /opt/conda/lib/python3.8/site-packages/torch/include -isystem /opt/conda/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -isystem /opt/conda/lib/python3.8/site-packages/torch/include/TH -isystem /opt/conda/lib/python3.8/site-packages/torch/include/THC -isystem /usr/local/cuda/include -isystem /opt/conda/include/python3.8 -D_GLIBCXX_USE_CXX11_ABI=1 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_75,code=sm_75 --compiler-options '-fPIC' --use_fast_math -std=c++14 -c /workspace/torch_utils/ops/sample.cu -o sample.cuda.o 
FAILED: sample.cuda.o 
/usr/local/cuda/bin/nvcc -DTORCH_EXTENSION_NAME=sample_plugin -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1013\" -isystem /opt/conda/lib/python3.8/site-packages/torch/include -isystem /opt/conda/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -isystem /opt/conda/lib/python3.8/site-packages/torch/include/TH -isystem /opt/conda/lib/python3.8/site-packages/torch/include/THC -isystem /usr/local/cuda/include -isystem /opt/conda/include/python3.8 -D_GLIBCXX_USE_CXX11_ABI=1 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_75,code=sm_75 --compiler-options '-fPIC' --use_fast_math -std=c++14 -c /workspace/torch_utils/ops/sample.cu -o sample.cuda.o 
/workspace/torch_utils/ops/sample.cu(72): error: no instance of overloaded function "atomicAdd" matches the argument list
            argument types are: (c10::Half *, c10::Half)
          detected during:
            instantiation of "void kernel_backward&lt;T&gt;(kernel_params) [with T=c10::Half]" 
(84): here
            instantiation of "kernel_spec choose_kernel&lt;T&gt;(int) [with T=c10::Half]" 
(93): here
</code></pre>
<p>atomicAdd should support half-precision since compute capability 7.* and I have 7.5 (GeForce RTX 2080) in my case.<br>
Is there anything I did wrong? or any suggestions I can fix it? Thanks!</p>

ptrblck: <p>Yes, <code>atomicAdd</code> is supported and used as seen <a href="https://github.com/pytorch/pytorch/blob/085e2f7bddc45f859fcdb786926d60d709b2daa0/aten/src/ATen/cuda/Atomic.cuh#L181-L190">here</a>.<br>
I don’t know how you are calling it, but based on the error message it seems the passed types are wrong.</p>

YutongZheng: <p>Turns out I should use <code>gpuAtomicAdd</code> rather than <code>atomicAdd</code>. Replacing solved the problem! Thanks!</p>