How to get the pointer of torch.cuda.HalfTensor?

username: fangwei123456
cooked: <p>Hi everyone, I am trying to write some custom CUDA kernel functions for torch.cuda.HalfTensor. However, I don’t know how to get pointer of torch.cuda.HalfTensor. Here are my codes:</p>

username: fangwei123456
cooked: <p>c.cpp:</p>
<pre><code class="lang-auto">#include &lt;iostream&gt;

#include &lt;torch/extension.h&gt;

torch::Tensor mul(torch::Tensor &amp; x, torch::Tensor &amp; y);

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {

    m.def("mul", &amp;mul);

}
</code></pre>
<p>kernel.cu:</p>
<pre><code class="lang-auto">#include &lt;cuda.h&gt;

#include &lt;cuda_runtime.h&gt;

#include &lt;torch/extension.h&gt;

#include &lt;ATen/ATen.h&gt;

#include &lt;math.h&gt;

#include &lt;stdio.h&gt;

#include &lt;cuda_fp16.h&gt;

__global__ void mul_cuda_kernel(const __half* __restrict__ x, const __half* __restrict__ y, __half* __restrict__ z, const int size)

{

  const int index = blockIdx.x * blockDim.x + threadIdx.x;

  if (index &lt; size) {

      z[index] = __hmul(x[index], y[index]);

  }

}

torch::Tensor mul(torch::Tensor &amp; x, torch::Tensor &amp; y)

{

    TORCH_CHECK(x.device().is_cuda(), "x must be a CUDA tensor");

    TORCH_CHECK(y.device().is_cuda(), "y must be a CUDA tensor");

    if (! x.is_contiguous())

    {

        x = x.contiguous();

    }

    if (! y.is_contiguous())

    {

        y = y.contiguous();

    }

    auto z = torch::zeros_like(x.data());

    const int size = x.numel();

    const int threads = 1024;

    const int blocks = (size + threads - 1) / threads;

    cudaSetDevice(x.get_device());

    mul_cuda_kernel&lt;&lt;&lt;blocks, threads&gt;&gt;&gt;(x.data_ptr&lt;__half&gt;(), y.data_ptr&lt;__half&gt;(), z.data_ptr&lt;__half&gt;(), size);

    return z;   

}
</code></pre>
<p>test.py:</p>
<pre><code class="lang-auto">import torch

import torch.nn as nn

import torch.nn.functional as F

from torch.utils import cpp_extension

cext_temp = cpp_extension.load(name='cext_temp',

    sources=['./c.cpp', './kernel.cu'], verbose=True)

device = 'cuda:0'

x = torch.rand([8], device=device).half()

y = torch.rand([8], device=device).half()

print(x)

print(y)

print(x*y)

print(cext_temp.mul(x, y))
</code></pre>
<p>When I run test.py, I get the error:</p>
<pre><code class="lang-auto">  File "/home/wfang/spikingjellyCppExt/temp/test.py", line 6, in &lt;module&gt;
    cext_temp = cpp_extension.load(name='cext_temp',
  File "/home/wfang/anaconda3/envs/pytorch-env/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 969, in load
    return _jit_compile(
  File "/home/wfang/anaconda3/envs/pytorch-env/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1196, in _jit_compile
    return _import_module_from_library(name, build_directory, is_python_module)
  File "/home/wfang/anaconda3/envs/pytorch-env/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1547, in _import_module_from_library
    return imp.load_module(module_name, file, path, description)
  File "/home/wfang/anaconda3/envs/pytorch-env/lib/python3.8/imp.py", line 242, in load_module
    return load_dynamic(name, filename, file)
  File "/home/wfang/anaconda3/envs/pytorch-env/lib/python3.8/imp.py", line 342, in load_dynamic
    return _load(spec)
  File "&lt;frozen importlib._bootstrap&gt;", line 702, in _load
  File "&lt;frozen importlib._bootstrap&gt;", line 657, in _load_unlocked
  File "&lt;frozen importlib._bootstrap&gt;", line 556, in module_from_spec
  File "&lt;frozen importlib._bootstrap_external&gt;", line 1101, in create_module
  File "&lt;frozen importlib._bootstrap&gt;", line 219, in _call_with_frames_removed
ImportError: /home/wfang/.cache/torch_extensions/cext_temp/cext_temp.so: undefined symbol: _ZNK2at6Tensor8data_ptrI6__halfEEPT_v
</code></pre>

username: fangwei123456
cooked: <p>I find that I must use <code>at::Half</code> rather than <code>half</code>. Another question is how can I use <code>half2</code>? It seems that <code>at::Half2</code> does not exist.</p>

username: hemeng
cooked: <p>Hey, I also encountered this question. May I ask how you convent <code>at::Half</code> to <code>half</code> while computing in cuda kernel? Thanks for your help!</p>

username: fangwei123456
cooked: <p>In c++, you can use <code>x.data_ptr&lt;at::Half&gt;()</code> to get a half pinter <code>half* x</code>, which is used in cuda kernel.</p>

username: hemeng
cooked: <p>Thanks! I’m just not sure whether the <code>at::Half</code> is just the same <code>half</code> or not.<br>
I use the <code>(half*)(x.data_ptr&lt;at::Half&gt;())</code> to get the pointer. It works fine for now <img src="https://discuss.pytorch.org/images/emoji/apple/relaxed.png?v=10" title=":relaxed:" class="emoji" alt=":relaxed:"> .</p>