How to implement a GPU neighbor sampling kernel properly in pyg-lib

[Gibyeng]

I am writing a GPU-based neighbor sampling kernel, but I got a problem.
This is what I have done.

• Create file "neighbor_kenel.cu" in "csrc/sampler/cuda"
• In the "neighbor_kernel.cu", I defined namespace "pyg" and "sampler" and "neighbor_sample_kernel" function, just like the one in "random_walk_kernel.cu".
• Register the kernel as follows

TORCH_LIBRARY_IMPL(pyg, CUDA, m) {
  m.impl(TORCH_SELECTIVE_NAME("pyg::neighbor_sample"),
         TORCH_FN(neighbor_sample_kernel));
}

• In the python side, I am trying to call the "neighbor_sample_kernel" in this way.

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
d_rowptr = rowptr.to(device)
d_col = col.to(device)
d_seed = seed.to(device)
for seed in tqdm(node_perm.split(batch_size)):
    pyg_lib.sampler.neighbor_sample(
        d_rowptr,
        d_col,
        d_seed,
        num_neighbors,
        node_time=node_time,
        edge_time=None,
        seed_time=None,
        edge_weight=edge_weight,
        replace=args.replace,
        directed=args.directed,
        disjoint=args.disjoint,
        temporal_strategy=args.temporal_strategy,
        return_edge_id=True,
    )

Then I got this error

NotImplementedError: Could not run 'pyg::neighbor_sample' with arguments from the 'CUDA' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build).

Thank you in advance.

[rusty1s]

This looks correct. Can you confirm that your kernel gets actually compiled during building?

[rusty1s]

Can you show me the installation log of it? I also think you need to run via python setup.py install.

[Gibyeng]

this is the log after I run python ./setup.py build

running build
running build_py
running build_ext
-- Caffe2: CUDA detected: 11.8
-- Caffe2: CUDA nvcc is: /usr/local/cuda/bin/nvcc
-- Caffe2: CUDA toolkit directory: /usr/local/cuda
-- Caffe2: Header version is: 11.8
-- /usr/local/cuda/lib64/libnvrtc.so shorthash is 672ee683
-- USE_CUDNN is set to 0. Compiling without cuDNN support
-- USE_CUSPARSELT is set to 0. Compiling without cuSPARSELt support
-- Autodetected CUDA architecture(s):  7.5 7.5
-- Added CUDA NVCC flags for: -gencode;arch=compute_75,code=sm_75
CMake Warning at /home/xxx/anaconda3/envs/pyG-env/lib/python3.11/site-packages/torch/share/cmake/Torch/TorchConfig.cmake:22 (message):
  static library kineto_LIBRARY-NOTFOUND not found.
Call Stack (most recent call first):
  /home/xxx/anaconda3/envs/pyG-env/lib/python3.11/site-packages/torch/share/cmake/Torch/TorchConfig.cmake:127 (append_torchlib_if_found)
  CMakeLists.txt:79 (find_package)


-- CMAKE_BUILD_TYPE: RELEASE
-- CMAKE_CXX_FLAGS_DEBUG: -g -O0
-- CMAKE_CXX_FLAGS_RELEASE: -O3 -DNDEBUG
-- Configuring done (5.5s)
CMake Warning (dev) in CMakeLists.txt:
  Policy CMP0104 is not set: CMAKE_CUDA_ARCHITECTURES now detected for NVCC,
  empty CUDA_ARCHITECTURES not allowed.  Run "cmake --help-policy CMP0104"
  for policy details.  Use the cmake_policy command to set the policy and
  suppress this warning.

  CUDA_ARCHITECTURES is empty for target "pyg".
This warning is for project developers.  Use -Wno-dev to suppress it.

CMake Warning at CMakeLists.txt:55 (add_library):
  Cannot generate a safe runtime search path for target pyg because files in
  some directories may conflict with libraries in implicit directories:

    runtime library [libgomp.so.1] in /usr/lib/gcc/x86_64-linux-gnu/7 may be hidden by files in:
      /home/xxx/anaconda3/envs/pyG-env/lib

  Some of these libraries may not be found correctly.


-- Generating done (0.2s)
-- Build files have been written to: /home/xxx/Documents/pyg-lib/build/temp.linux-x86_64-cpython-311
[1/2] Building CUDA object CMakeFiles/...csrc/sampler/cuda/neighbor_kernel.cu.o
/home/xxx/Documents/pyg-lib/pyg_lib/csrc/sampler/cuda/neighbor_kernel.cu(16): warning #177-D: function "pyg::sampler::<unnamed>::blocks" was declared but never referenced

[2/2] Linking CXX shared library /home...lib.linux-x86_64-cpython-311/libpyg.so

[Gibyeng]

And this is the log, after I run python setup.py install.

running install
/home/xxx/anaconda3/envs/pyG-env/lib/python3.11/site-packages/setuptools/_distutils/cmd.py:66: SetuptoolsDeprecationWarning: setup.py install is deprecated.
!!

        ********************************************************************************
        Please avoid running ``setup.py`` directly.
        Instead, use pypa/build, pypa/installer or other
        standards-based tools.

        See https://blog.ganssle.io/articles/2021/10/setup-py-deprecated.html for details.
        ********************************************************************************

!!
  self.initialize_options()
/home/xxx/anaconda3/envs/pyG-env/lib/python3.11/site-packages/setuptools/_distutils/cmd.py:66: EasyInstallDeprecationWarning: easy_install command is deprecated.
!!

        ********************************************************************************
        Please avoid running ``setup.py`` and ``easy_install``.
        Instead, use pypa/build, pypa/installer or other
        standards-based tools.

        See https://github.com/pypa/setuptools/issues/917 for details.
        ********************************************************************************

!!
  self.initialize_options()
running bdist_egg
running egg_info
writing pyg_lib.egg-info/PKG-INFO
writing dependency_links to pyg_lib.egg-info/dependency_links.txt
writing requirements to pyg_lib.egg-info/requires.txt
writing top-level names to pyg_lib.egg-info/top_level.txt
reading manifest file 'pyg_lib.egg-info/SOURCES.txt'
reading manifest template 'MANIFEST.in'
warning: no previously-included files matching '*' found under directory 'docs'
warning: no previously-included files matching '*' found under directory 'benchmark'
adding license file 'LICENSE'
adding license file 'LICENSE_radix_sort'
writing manifest file 'pyg_lib.egg-info/SOURCES.txt'
installing library code to build/bdist.linux-x86_64/egg
running install_lib
running build_py
running build_ext
-- Caffe2: CUDA detected: 11.8
-- Caffe2: CUDA nvcc is: /usr/local/cuda/bin/nvcc
-- Caffe2: CUDA toolkit directory: /usr/local/cuda
-- Caffe2: Header version is: 11.8
-- /usr/local/cuda/lib64/libnvrtc.so shorthash is 672ee683
-- USE_CUDNN is set to 0. Compiling without cuDNN support
-- USE_CUSPARSELT is set to 0. Compiling without cuSPARSELt support
-- Autodetected CUDA architecture(s):  7.5 7.5
-- Added CUDA NVCC flags for: -gencode;arch=compute_75,code=sm_75
CMake Warning at /home/xxx/anaconda3/envs/pyG-env/lib/python3.11/site-packages/torch/share/cmake/Torch/TorchConfig.cmake:22 (message):
  static library kineto_LIBRARY-NOTFOUND not found.
Call Stack (most recent call first):
  /home/xxx/anaconda3/envs/pyG-env/lib/python3.11/site-packages/torch/share/cmake/Torch/TorchConfig.cmake:127 (append_torchlib_if_found)
  CMakeLists.txt:79 (find_package)


-- CMAKE_BUILD_TYPE: RELEASE
-- CMAKE_CXX_FLAGS_DEBUG: -g -O0
-- CMAKE_CXX_FLAGS_RELEASE: -O3 -DNDEBUG
-- Configuring done (1.6s)
CMake Warning (dev) in CMakeLists.txt:
  Policy CMP0104 is not set: CMAKE_CUDA_ARCHITECTURES now detected for NVCC,
  empty CUDA_ARCHITECTURES not allowed.  Run "cmake --help-policy CMP0104"
  for policy details.  Use the cmake_policy command to set the policy and
  suppress this warning.

  CUDA_ARCHITECTURES is empty for target "pyg".
This warning is for project developers.  Use -Wno-dev to suppress it.

CMake Warning at CMakeLists.txt:55 (add_library):
  Cannot generate a safe runtime search path for target pyg because files in
  some directories may conflict with libraries in implicit directories:

    runtime library [libgomp.so.1] in /usr/lib/gcc/x86_64-linux-gnu/7 may be hidden by files in:
      /home/xxx/anaconda3/envs/pyG-env/lib

  Some of these libraries may not be found correctly.


-- Generating done (0.0s)
-- Build files have been written to: /home/xxx/Documents/pyg-lib/build/temp.linux-x86_64-cpython-311
ninja: no work to do.
creating build/bdist.linux-x86_64/egg
copying build/lib.linux-x86_64-cpython-311/libpyg.so -> build/bdist.linux-x86_64/egg
creating build/bdist.linux-x86_64/egg/pyg_lib
copying build/lib.linux-x86_64-cpython-311/pyg_lib/_triton.py -> build/bdist.linux-x86_64/egg/pyg_lib
creating build/bdist.linux-x86_64/egg/pyg_lib/partition
copying build/lib.linux-x86_64-cpython-311/pyg_lib/partition/__init__.py -> build/bdist.linux-x86_64/egg/pyg_lib/partition
creating build/bdist.linux-x86_64/egg/pyg_lib/ops
copying build/lib.linux-x86_64-cpython-311/pyg_lib/ops/__init__.py -> build/bdist.linux-x86_64/egg/pyg_lib/ops
copying build/lib.linux-x86_64-cpython-311/pyg_lib/ops/scatter_reduce.py -> build/bdist.linux-x86_64/egg/pyg_lib/ops
copying build/lib.linux-x86_64-cpython-311/pyg_lib/__init__.py -> build/bdist.linux-x86_64/egg/pyg_lib
copying build/lib.linux-x86_64-cpython-311/pyg_lib/home.py -> build/bdist.linux-x86_64/egg/pyg_lib
creating build/bdist.linux-x86_64/egg/pyg_lib/sampler
copying build/lib.linux-x86_64-cpython-311/pyg_lib/sampler/__init__.py -> build/bdist.linux-x86_64/egg/pyg_lib/sampler
copying build/lib.linux-x86_64-cpython-311/pyg_lib/testing.py -> build/bdist.linux-x86_64/egg/pyg_lib
byte-compiling build/bdist.linux-x86_64/egg/pyg_lib/_triton.py to _triton.cpython-311.pyc
byte-compiling build/bdist.linux-x86_64/egg/pyg_lib/partition/__init__.py to __init__.cpython-311.pyc
byte-compiling build/bdist.linux-x86_64/egg/pyg_lib/ops/__init__.py to __init__.cpython-311.pyc
byte-compiling build/bdist.linux-x86_64/egg/pyg_lib/ops/scatter_reduce.py to scatter_reduce.cpython-311.pyc
byte-compiling build/bdist.linux-x86_64/egg/pyg_lib/__init__.py to __init__.cpython-311.pyc
byte-compiling build/bdist.linux-x86_64/egg/pyg_lib/home.py to home.cpython-311.pyc
byte-compiling build/bdist.linux-x86_64/egg/pyg_lib/sampler/__init__.py to __init__.cpython-311.pyc
byte-compiling build/bdist.linux-x86_64/egg/pyg_lib/testing.py to testing.cpython-311.pyc
creating stub loader for libpyg.so
byte-compiling build/bdist.linux-x86_64/egg/libpyg.py to libpyg.cpython-311.pyc
creating build/bdist.linux-x86_64/egg/EGG-INFO
copying pyg_lib.egg-info/PKG-INFO -> build/bdist.linux-x86_64/egg/EGG-INFO
copying pyg_lib.egg-info/SOURCES.txt -> build/bdist.linux-x86_64/egg/EGG-INFO
copying pyg_lib.egg-info/dependency_links.txt -> build/bdist.linux-x86_64/egg/EGG-INFO
copying pyg_lib.egg-info/requires.txt -> build/bdist.linux-x86_64/egg/EGG-INFO
copying pyg_lib.egg-info/top_level.txt -> build/bdist.linux-x86_64/egg/EGG-INFO
writing build/bdist.linux-x86_64/egg/EGG-INFO/native_libs.txt
zip_safe flag not set; analyzing archive contents...
__pycache__.libpyg.cpython-311: module references __file__
pyg_lib.__pycache__.__init__.cpython-311: module references __file__
creating 'dist/pyg_lib-0.3.1-py3.11-linux-x86_64.egg' and adding 'build/bdist.linux-x86_64/egg' to it
removing 'build/bdist.linux-x86_64/egg' (and everything under it)
Processing pyg_lib-0.3.1-py3.11-linux-x86_64.egg
removing '/home/xxx/anaconda3/envs/pyG-env/lib/python3.11/site-packages/pyg_lib-0.3.1-py3.11-linux-x86_64.egg' (and everything under it)
creating /home/xxx/anaconda3/envs/pyG-env/lib/python3.11/site-packages/pyg_lib-0.3.1-py3.11-linux-x86_64.egg
Extracting pyg_lib-0.3.1-py3.11-linux-x86_64.egg to /home/xxx/anaconda3/envs/pyG-env/lib/python3.11/site-packages
Adding pyg-lib 0.3.1 to easy-install.pth file

Installed /home/xxx/anaconda3/envs/pyG-env/lib/python3.11/site-packages/pyg_lib-0.3.1-py3.11-linux-x86_64.egg
Processing dependencies for pyg-lib==0.3.1
Finished processing dependencies for pyg-lib==0.3.1

[Gibyeng]

It seems the kernel is called as expected after I run python setup.py install. I think the problem has been solved. thank you for your help.

[rusty1s]

Super :)