"megatron/git@developer.sourcefind.cn:OpenDAS/megatron-lm.git" did not exist on "01dffdb557e00e3761111170ee6b11c52165fed6"
Commit 05d058c1 authored by rusty1s's avatar rusty1s
Browse files

added kernel to ffi build

parent 2aa89004
......@@ -8,25 +8,25 @@ abs_path = 'torch_scatter'
headers = ['torch_scatter/src/cpu.h']
sources = ['torch_scatter/src/cpu.c']
includes = ['torch_scatter/src']
defines = []
include_dirs = ['torch_scatter/src', 'torch_scatter/kernel']
define_macros = []
extra_objects = []
with_cuda = False
if torch.cuda.is_available():
headers += ['torch_scatter/src/cuda.h']
sources += ['torch_scatter/src/cuda.c']
defines += [('WITH_CUDA', None)]
define_macros += [('WITH_CUDA', None)]
extra_objects += ['torch_scatter/build/kernel.so']
with_cuda = True
ffi = create_extension(
name='torch_scatter._ext.ffi',
package=True,
verbose=True,
headers=headers,
sources=sources,
include_dirs=includes,
define_macros=defines,
include_dirs=include_dirs,
define_macros=define_macros,
extra_objects=extra_objects,
with_cuda=with_cuda,
relative_to=__file__)
......
#include <THC/THC.h>
#include "kernel.h"
#define scatter_(NAME) TH_CONCAT_4(scatter_, NAME, _cuda_, Real)
#define index_backward TH_CONCAT_2(index_backward_cuda_, Real)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment