Error when building a Custom C++ and CUDA Extension

Bradipo_Eremita: <p>I’m trying to implement a 3D NMS function in CUDA following the example outlined in <a href="https://pytorch.org/tutorials/advanced/cpp_extension.html" class="inline-onebox" rel="noopener nofollow ugc">Custom C++ and CUDA Extensions — PyTorch Tutorials 2.1.1+cu121 documentation</a>.</p>
<p>My CUDA and C++ experience is limited and thus it may be that there is, as I hope, a pretty stupid error somewhere that one of you guys could spot. Anyway my code is made by four scripts:</p>
<ol>
<li><strong>nms_kernel.cu</strong> which contains the cuda functions;</li>
<li><strong>nms_cuda.cpp</strong> which contains the c++ function;</li>
<li><strong>setup.py</strong></li>
</ol>
<p>Hereafter I’m posting the four scripts:</p>
<p><strong>nms_kernel.cu</strong></p>
<pre><code class="lang-auto">#include &lt;torch/extension.h&gt;
#include &lt;cuda.h&gt;
#include &lt;cuda_runtime.h&gt;

__device__ inline float devIoU(float const * const a, float const * const b) {
  float left = fmaxf(a[0], b[0]), right = fminf(a[2], b[2]);
  float top = fmaxf(a[1], b[1]), bottom = fminf(a[3], b[3]);
  float front = fmaxf(a[4], b[4]), back = fminf(a[5], b[5]);

  float width = fmaxf(right - left + 1, 0.f), height = fmaxf(bottom - top + 1, 0.f), depth = fmaxf(back - front + 1, 0.f);
  float interS = width * height * depth;
  float Sa = (a[2] - a[0] + 1) * (a[3] - a[1] + 1) * (a[5] - a[4] + 1);
  float Sb = (b[2] - b[0] + 1) * (b[3] - b[1] + 1) * (b[5] - b[4] + 1);
  //printf("IoU 3D %f \n", interS / (Sa + Sb - interS));

  return interS / (Sa + Sb - interS);
}

__global__ void nms_kernel(const int n_boxes, const float nms_overlap_thresh,
                           const float *dev_boxes, unsigned long long *dev_mask) {
  const int row_start = blockIdx.y;
  const int col_start = blockIdx.x;

  // if (row_start &gt; col_start) return;

  const int row_size =
        fminf(n_boxes - row_start * threadsPerBlock, threadsPerBlock);
  const int col_size =
        fminf(n_boxes - col_start * threadsPerBlock, threadsPerBlock);

  __shared__ float block_boxes[threadsPerBlock * 7];
  if (threadIdx.x &lt; col_size) {
    block_boxes[threadIdx.x * 7 + 0] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 7 + 0];
    block_boxes[threadIdx.x * 7 + 1] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 7 + 1];
    block_boxes[threadIdx.x * 7 + 2] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 7 + 2];
    block_boxes[threadIdx.x * 7 + 3] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 7 + 3];
    block_boxes[threadIdx.x * 7 + 4] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 7 + 4];
    block_boxes[threadIdx.x * 7 + 5] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 7 + 5];
    block_boxes[threadIdx.x * 7 + 6] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 7 + 6];
  }
  __syncthreads();

  if (threadIdx.x &lt; row_size) {
    const int cur_box_idx = threadsPerBlock * row_start + threadIdx.x;
    const float *cur_box = dev_boxes + cur_box_idx * 7;
    int i = 0;
    unsigned long long t = 0;
    int start = 0;
    if (row_start == col_start) {
      start = threadIdx.x + 1;
    }
    for (i = start; i &lt; col_size; i++) {
      if (devIoU(cur_box, block_boxes + i * 7) &gt; nms_overlap_thresh) {
        t |= 1ULL &lt;&lt; i;
      }
    }
    const int col_blocks = DIVUP(n_boxes, threadsPerBlock);
    dev_mask[cur_box_idx * col_blocks + col_start] = t;
  }
}


void _nms(int boxes_num, float * boxes_dev,
          unsigned long long * mask_dev, float nms_overlap_thresh) {


  dim3 blocks(DIVUP(boxes_num, threadsPerBlock),
              DIVUP(boxes_num, threadsPerBlock));
  dim3 threads(threadsPerBlock);
  nms_kernel&lt;&lt;&lt;blocks, threads&gt;&gt;&gt;(boxes_num,
                                  nms_overlap_thresh,
                                  boxes_dev,
                                  mask_dev);
}


void gpu_nms(at::Tensor keep, at::Tensor boxes, at::Tensor num_out,  float nms_overlap_thresh){
    // params boxes: (N, 7) [x, y, z, dx, dy, dz, heading]
    // params keep: (N)

    int boxes_num = boxes.size(0);
    const float * boxes_data = boxes.data&lt;float&gt;();
    long * keep_data = keep.data&lt;long&gt;();

    const int col_blocks = DIVUP(boxes_num, threadsPerBlock);

    unsigned long long mask_data[boxes_num * col_blocks]
    cudaMalloc((void**)&amp;mask_data, boxes_num * col_blocks * sizeof(unsigned long long));
    _nms(boxes_num, boxes_data, mask_data, nms_overlap_thresh);

    unsigned long long mask_cpu[boxes_num * col_blocks];
    //unsigned long long * mask_cpu = new unsigned long long [boxes_num * col_blocks];
    std::vector&lt;unsigned long long&gt; mask_cpu(boxes_num * col_blocks);

    //printf("boxes_num=%d, col_blocks=%d\n", boxes_num, col_blocks);
    cudaMemcpy(&amp;mask_cpu[0], mask_data, boxes_num * col_blocks * sizeof(unsigned long long),
                           cudaMemcpyDeviceToHost);

    cudaFree(mask_data);

    unsigned long long remv_cpu[col_blocks];
    cudaMemset(remv_cpu, 0, col_blocks * sizeof(unsigned long long));

    int num_to_keep = 0;
    int i, j;
    for (int i = 0; i &lt; boxes_num; i++){
        int nblock = i / threadsPerBlock;
        int inblock = i % threadsPerBlock;

        if (!(remv_cpu[nblock] &amp; (1ULL &lt;&lt; inblock))){
            keep_data[num_to_keep++] = i;
            unsigned long long *p = &amp;mask_cpu[0] + i * col_blocks;
            for (int j = nblock; j &lt; col_blocks; j++){
                remv_cpu[j] |= p[j];
            }
        }
    }
    if ( cudaSuccess != cudaGetLastError() ) printf( "Error!\n" );

    return num_to_keep;
}

</code></pre>
<p><strong>nms_cuda.cpp</strong></p>
<pre><code class="lang-auto">#include &lt;torch/extension.h&gt;

#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)


int gpu_nms(at::Tensor keep, at::Tensor boxes, at::Tensor num_out,  float nms_overlap_thresh)

inst gpu_nms_forward(at::Tensor keep, at::Tensor boxes, at::Tensor num_out,  float nms_overlap_thresh){
  CHECK_INPUT(keep)
  CHECK_INPUT(boxes)
  return gpu_nms(keep, boxes, num_out, nms_overlap_threshold)
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){
  m.def("forward", &amp;gpu_nms_forward, "gpu_nms ('CUDA')");
}

</code></pre>
<p><strong>setup.py</strong></p>
<pre><code class="lang-auto">from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension


setup(
    name='nms_3d',
    ext_modules=[
        CUDAExtension('nms_3d', [
            'nms_cuda.cpp',
            'nms_kernel.cu'

        ])],
    cmdclass={'build_ext': BuildExtension})

</code></pre>
<p><strong>System information</strong>: I’m running this is a conda environment with the following relevant packages and libraries:</p>
<ol>
<li>pytorch=1.9.1 (stable);</li>
<li>CUDA=10.2;</li>
<li>cudatoolkit=10.2.89;</li>
<li>cudatoolkit-dev=10.1.243;</li>
<li>python=3.8.8;</li>
<li>cudnn=7.6.5;</li>
<li>Pytorch C++ API: libtorch-shared-with-deps-1.9.1+cu102</li>
<li>Ubuntu 20.04.3 LTS</li>
</ol>
<p>When I try to do</p>
<blockquote>
<p>python setup.py install</p>
</blockquote>
<p>I get the following error:</p>
<pre><code class="lang-auto">running install
running bdist_egg
running egg_info
writing nms_3d.egg-info/PKG-INFO
writing dependency_links to nms_3d.egg-info/dependency_links.txt
writing top-level names to nms_3d.egg-info/top_level.txt
reading manifest file 'nms_3d.egg-info/SOURCES.txt'
writing manifest file 'nms_3d.egg-info/SOURCES.txt'
installing library code to build/bdist.linux-x86_64/egg
running install_lib
running build_ext
building 'nms_3d' extension
Emitting ninja build file /home/michele/Documents/GitHub/MicheleDelliVeneri/Thesis/Code/nms_3D/build/temp.linux-x86_64-3.8/build.ninja...
Compiling objects...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
ninja: error: '/home/michele/Documents/GitHub/MicheleDelliVeneri/Thesis/Code/nms_3D/nms_kernel.cu', needed by '/home/michele/Documents/GitHub/MicheleDelliVeneri/Thesis/Code/nms_3D/build/temp.linux-x86_64-3.8/nms_kernel.o', missing and no known rule to make it
Traceback (most recent call last):
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1666, in _run_ninja_build
    subprocess.run(
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/subprocess.py", line 516, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['ninja', '-v']' returned non-zero exit status 1.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "setup.py", line 5, in &lt;module&gt;
    setup(
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/site-packages/setuptools/__init__.py", line 153, in setup
    return distutils.core.setup(**attrs)
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/distutils/core.py", line 148, in setup
    dist.run_commands()
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/distutils/dist.py", line 966, in run_commands
    self.run_command(cmd)
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/distutils/dist.py", line 985, in run_command
    cmd_obj.run()
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/site-packages/setuptools/command/install.py", line 67, in run
    self.do_egg_install()
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/site-packages/setuptools/command/install.py", line 109, in do_egg_install
    self.run_command('bdist_egg')
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/distutils/cmd.py", line 313, in run_command
    self.distribution.run_command(command)
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/distutils/dist.py", line 985, in run_command
    cmd_obj.run()
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/site-packages/setuptools/command/bdist_egg.py", line 164, in run
    cmd = self.call_command('install_lib', warn_dir=0)
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/site-packages/setuptools/command/bdist_egg.py", line 150, in call_command
    self.run_command(cmdname)
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/distutils/cmd.py", line 313, in run_command
    self.distribution.run_command(command)
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/distutils/dist.py", line 985, in run_command
    cmd_obj.run()
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/site-packages/setuptools/command/install_lib.py", line 11, in run
    self.build()
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/distutils/command/install_lib.py", line 107, in build
    self.run_command('build_ext')
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/distutils/cmd.py", line 313, in run_command
    self.distribution.run_command(command)
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/distutils/dist.py", line 985, in run_command
    cmd_obj.run()
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/site-packages/setuptools/command/build_ext.py", line 79, in run
    _build_ext.run(self)
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/site-packages/Cython/Distutils/old_build_ext.py", line 186, in run
    _build_ext.build_ext.run(self)
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/distutils/command/build_ext.py", line 340, in run
    self.build_extensions()
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 709, in build_extensions
    build_ext.build_extensions(self)
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/site-packages/Cython/Distutils/old_build_ext.py", line 195, in build_extensions
    _build_ext.build_ext.build_extensions(self)
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/distutils/command/build_ext.py", line 449, in build_extensions
    self._build_extensions_serial()
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/distutils/command/build_ext.py", line 474, in _build_extensions_serial
    self.build_extension(ext)
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/site-packages/setuptools/command/build_ext.py", line 196, in build_extension
    _build_ext.build_extension(self, ext)
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/distutils/command/build_ext.py", line 528, in build_extension
    objects = self.compiler.compile(sources,
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 530, in unix_wrap_ninja_compile
    _write_ninja_file_and_compile_objects(
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1355, in _write_ninja_file_and_compile_objects
    _run_ninja_build(
  File "/home/michele/anaconda3/envs/with_cudatoolkit/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1682, in _run_ninja_build
    raise RuntimeError(message) from e
RuntimeError: Error compiling objects for extension
</code></pre>
<p>From the error message, I cannot seem to understand which may be the problem, any help would be greatly appreciated.</p>
<p>Thank you,<br>
Michele</p>

googlebot: <p>manually run ninja from console on this:</p>
<blockquote>
<p>Emitting ninja build file /home/michele/Documents/GitHub/MicheleDelliVeneri/Thesis/Code/nms_3D/build/temp.linux-x86_64-3.8/build.ninja</p>
</blockquote>
<p>and you should see c++ error(s)</p>

Bradipo_Eremita: <p>First of all thank you for the reply, I tried what you suggested and by going in the mentioned folder and running the command:</p>
<blockquote>
<p>ninja</p>
</blockquote>
<p>I get the same output message unfortunately:</p>
<pre><code class="lang-auto">ninja: error: '/home/michele/Documents/GitHub/MicheleDelliVeneri/Thesis/Code/nms_3D/nms_kernel.cu', needed by '/home/michele/Documents/GitHub/MicheleDelliVeneri/Thesis/Code/nms_3D/build/temp.linux-x86_64-3.8/nms_kernel.o', missing and no known rule to make it

</code></pre>
<p>am I misunderstanding what you meant by manually run ninja?</p>

googlebot: <p>If you view ninja.build, it has relevant compiler commands (in your case nms_kernel.cu build is failing). I think your build file doesn’t contain absolute filenames, then try either running from directory with sources “ninja -f &lt;full_path_to_ninja.build&gt;” or copy compiler command from build file.</p>
<p>Actually, I think there is some setup.py flag for verbose output, but I can’t tell this specifically right now.</p>

Bradipo_Eremita: <p>Thank you Alex for your responses, it gave me an idea on how to debug it and at the end I managed to do it. Given that maybe it could be instructional for someone else,  here it is how it went down.</p>
<p>Instead of going in the setup.py direction, I made a CMakeLists.txt file which looks like this:</p>
<pre><code class="lang-auto">cmake_minimum_required(VERSION 3.10 FATAL_ERROR)
# here I decide the name of the library
project(nms_cmake LANGUAGES CXX CUDA)
set(CMAKE_CXX_STANDARD 14)

# I make sure to point to the correct libraries
set(CMAKE_PREFIX_PATH /home/michele/pytorch_cpp/libtorch-shared-with-deps-1.9.1+cu102/libtorch)

# I make sure that it has all the python and torch directories available
set(TORCH_INCLUDE_DIRS
   /home/michele/anaconda3/envs/with_cudatoolkit/include/
   /home/michele/anaconda3/envs/with_cudatoolkit/include/torch/csrc/api/include)

# I make sure that both Python and Torch are found, also all the py-dev stuff
find_package(Torch REQUIRED)
find_package(Python REQUIRED COMPONENTS Development)

# add all the sources to the library
# List all your code files here
add_library(nms_cmake SHARED
  main.cu
)

target_link_libraries(nms_cmake "${TORCH_LIBRARIES}")

# To really make sure that it is using the correct language and gcc to make stuff.
target_compile_options(nms_cmake PRIVATE $&lt;$&lt;COMPILE_LANGUAGE:CUDA&gt;:-ccbin g++-7&gt;)

</code></pre>
<p><strong>main.cu</strong>  at this stage was just a combination of the previous nms_kernel.cu and nms_cuda.cpp scripts and full of unknown syntactical errors.<br>
After creating the build fonder and launching:</p>
<pre><code class="lang-auto"> cmake -DCMAKE_PREFIX_PATH=$PWD/home/michele/pytorch_cpp/libtorch-shared-with-deps-1.9.1+cu102/libtorch ..
</code></pre>
<p>and</p>
<blockquote>
<p>make</p>
</blockquote>
<p>I could see may errors in the file and after some debugging it ended up like this:</p>
<pre><code class="lang-auto">#include &lt;c10/cuda/CUDAException.h&gt;
#include &lt;torch/torch.h&gt;
#include &lt;torch/all.h&gt;
#include &lt;torch/python.h&gt;
#include &lt;torch/extension.h&gt;

using namespace at;

#define DIVUP(m,n) ((m) / (n) + ((m) % (n) &gt; 0))
int const threadsPerBlock = sizeof(unsigned long long) * 8;


__device__ inline float devIoU(float const * const a, float const * const b) {
  float left = fmaxf(a[0], b[0]), right = fminf(a[2], b[2]);
  float top = fmaxf(a[1], b[1]), bottom = fminf(a[3], b[3]);
  float front = fmaxf(a[4], b[4]), back = fminf(a[5], b[5]);

  float width = fmaxf(right - left + 1, 0.f), height = fmaxf(bottom - top + 1, 0.f), depth = fmaxf(back - front + 1, 0.f);
  float interS = width * height * depth;
  float Sa = (a[2] - a[0] + 1) * (a[3] - a[1] + 1) * (a[5] - a[4] + 1);
  float Sb = (b[2] - b[0] + 1) * (b[3] - b[1] + 1) * (b[5] - b[4] + 1);
  //printf("IoU 3D %f \n", interS / (Sa + Sb - interS));

  return interS / (Sa + Sb - interS);
}

__global__ void nms_kernel(const int n_boxes, const float nms_overlap_thresh,
                           const float *dev_boxes, unsigned long long *dev_mask) {
  const int row_start = blockIdx.y;
  const int col_start = blockIdx.x;

  // if (row_start &gt; col_start) return;

  const int row_size =
        fminf(n_boxes - row_start * threadsPerBlock, threadsPerBlock);
  const int col_size =
        fminf(n_boxes - col_start * threadsPerBlock, threadsPerBlock);

  __shared__ float block_boxes[threadsPerBlock * 7];
  if (threadIdx.x &lt; col_size) {
    block_boxes[threadIdx.x * 7 + 0] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 7 + 0];
    block_boxes[threadIdx.x * 7 + 1] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 7 + 1];
    block_boxes[threadIdx.x * 7 + 2] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 7 + 2];
    block_boxes[threadIdx.x * 7 + 3] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 7 + 3];
    block_boxes[threadIdx.x * 7 + 4] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 7 + 4];
    block_boxes[threadIdx.x * 7 + 5] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 7 + 5];
    block_boxes[threadIdx.x * 7 + 6] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 7 + 6];
  }
  __syncthreads();

  if (threadIdx.x &lt; row_size) {
    const int cur_box_idx = threadsPerBlock * row_start + threadIdx.x;
    const float *cur_box = dev_boxes + cur_box_idx * 7;
    int i = 0;
    unsigned long long t = 0;
    int start = 0;
    if (row_start == col_start) {
      start = threadIdx.x + 1;
    }
    for (i = start; i &lt; col_size; i++) {
      if (devIoU(cur_box, block_boxes + i * 7) &gt; nms_overlap_thresh) {
        t |= 1ULL &lt;&lt; i;
      }
    }
    const int col_blocks = DIVUP(n_boxes, threadsPerBlock);
    dev_mask[cur_box_idx * col_blocks + col_start] = t;
  }
}


void _nms(int boxes_num, float * boxes_dev,
          unsigned long long * mask_dev, float nms_overlap_thresh) {


  dim3 blocks(DIVUP(boxes_num, threadsPerBlock),
              DIVUP(boxes_num, threadsPerBlock));
  dim3 threads(threadsPerBlock);
  nms_kernel&lt;&lt;&lt;blocks, threads&gt;&gt;&gt;(boxes_num,
                                  nms_overlap_thresh,
                                  boxes_dev,
                                  mask_dev);
}


int64_t gpu_nms(at::Tensor keep, at::Tensor boxes, at::Tensor num_out,  double nms_overlap_thresh){
    // params boxes: (N, 7) [x, y, z, dx, dy, dz, heading]
    // params keep: (N)

    int boxes_num = boxes.size(0);
    float * boxes_data = boxes.data_ptr&lt;float&gt;();
    long * keep_data = keep.data_ptr&lt;long&gt;();

    const int col_blocks = DIVUP(boxes_num, threadsPerBlock);

    //unsigned long long mask_data[boxes_num * col_blocks];
    unsigned long long *mask_data = NULL;
    cudaMalloc((void**)&amp;mask_data, boxes_num * col_blocks * sizeof(unsigned long long));
    _nms(boxes_num, boxes_data, mask_data, nms_overlap_thresh);

    //unsigned long long mask_cpu[boxes_num * col_blocks];
    //unsigned long long * mask_cpu = new unsigned long long [boxes_num * col_blocks];
    std::vector&lt;unsigned long long&gt; mask_cpu(boxes_num * col_blocks);

    //printf("boxes_num=%d, col_blocks=%d\n", boxes_num, col_blocks);
    cudaMemcpy(&amp;mask_cpu[0], mask_data, boxes_num * col_blocks * sizeof(unsigned long long),
                           cudaMemcpyDeviceToHost);

    cudaFree(mask_data);

    unsigned long long remv_cpu[col_blocks];
    cudaMemset(remv_cpu, 0, col_blocks * sizeof(unsigned long long));

    int num_to_keep = 0;
    for (int i = 0; i &lt; boxes_num; i++){
        int nblock = i / threadsPerBlock;
        int inblock = i % threadsPerBlock;

        if (!(remv_cpu[nblock] &amp; (1ULL &lt;&lt; inblock))){
            keep_data[num_to_keep++] = i;
            unsigned long long *p = &amp;mask_cpu[0] + i * col_blocks;
            for (int j = nblock; j &lt; col_blocks; j++){
                remv_cpu[j] |= p[j];
            }
        }
    }
    if ( cudaSuccess != cudaGetLastError() ) printf( "Error!\n" );

    return num_to_keep;
}

TORCH_LIBRARY(pytorch_cmake_example, m) {
  m.def("gpu_nms(Tensor keep, Tensor boxes, Tensor num_out, float nms_overlap_thresh) -&gt; int num_to_keep");
  m.impl("gpu_nms", c10::DispatchKey::CUDA, TORCH_FN(gpu_nms));
  //c10::DispatchKey::CPU is also an option
}
</code></pre>
<p>To load the library you can then import it in a python file as have a possibly working 3D NMS as follows:</p>
<pre><code class="lang-auto">import torch
torch.ops.load_library("build/libnms_cmake.so")

def nms_gpu(dets, thresh):
  """
  dets has to be a tensor
  """

  scores = dets[:, -1]
  order = scores.sort(0, descending=True)[1]
  dets = dets[order].contiguous()

  keep = torch.LongTensor(dets.size(0))
  num_out = torch.LongTensor(1)
  torch.ops.libnms_cmake.gpu_nms(keep, dets, num_out, thresh)
  return order[keep[:num_out[0]].cuda()].contiguous()
</code></pre>
<p>Thanks for the help.</p>