Hi, I’m trying to define a new sampler based on torch.poisson which is implemented in cpp. The idea is compile it using cpp_extension from torch. I got compilation errors when trying to include Aten module in cpp. The example code and the returned message can be found below. How can I call that cpp function from pytorch?
cpp_source="""
#include <torch/extension.h>
#include <ATen/ATen.h>
#include <ATen/CPUApplyUtils.h>
#include <ATen/Config.h>
#include <ATen/Dispatch.h>
#include <ATen/ExpandUtils.h>
#include <ATen/NativeFunctions.h>
#include <c10/util/Exception.h>
#include <c10/util/math_compat.h>
#include <ATen/Utils.h>
#include <ATen/CPUGenerator.h>
#include <ATen/core/DistributionsHelper.h>
#include <ATen/native/Distributions.h>
#include <ATen/native/DispatchStub.h>
#include <ATen/native/UnaryOps.h>
#include <ATen/NamedTensorUtils.h>
#include <type_traits>
#include <functional>
#include <assert.h>
#include <cpuinfo.h>
#include <float.h>
int64_t sample_poisson(double lambda, at::CPUGenerator* generator) {
TORCH_CHECK(lambda >= 0, "invalid Poisson rate, expected rate to be non-negative");
at::uniform_real_distribution<double> standard_uniform(0.0, 1.0);
int64_t X;
double prod, U, enlam;
enlam = std::exp(-lambda);
X = 0;
prod = 1.0;
while (1) {
U = standard_uniform(generator);
prod *= U;
if (prod > enlam) {
X += 1;
} else {
return X;
}
}
}
"""
import torch
from torch.utils.cpp_extension import load_inline
my_module = load_inline(
name='my_module',
cpp_sources=[cpp_source],
functions=['sample_poisson'],
build_directory='./build',
verbose=True
)
I got the following error:
Emitting ninja build file ./build.ninja...
Building extension module my_module...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
[1/2] c++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=my_module -DTORCH_API_INCLUDE_EXTENSION_H -isystem /home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/include -isystem /home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -isystem /home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/include/TH -isystem /home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/include/THC -isystem /home/eric/anaconda3/envs/zi/include/python3.8 -D_GLIBCXX_USE_CXX11_ABI=0 -fPIC -std=c++14 -c /home/eric/Desktop/test_1/main.cpp -o main.o
FAILED: main.o
c++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=my_module -DTORCH_API_INCLUDE_EXTENSION_H -isystem /home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/include -isystem /home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -isystem /home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/include/TH -isystem /home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/include/THC -isystem /home/eric/anaconda3/envs/zi/include/python3.8 -D_GLIBCXX_USE_CXX11_ABI=0 -fPIC -std=c++14 -c /home/eric/Desktop/test_1/main.cpp -o main.o
/home/eric/Desktop/test_1/main.cpp:15:10: fatal error: ATen/CPUGenerator.h: No such file or directory
15 | #include <ATen/CPUGenerator.h>
| ^~~~~~~~~~~~~~~~~~~~~
compilation terminated.
ninja: build stopped: subcommand failed.
Traceback (most recent call last):
File "/home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1516, in _run_ninja_build
subprocess.run(
File "/home/eric/anaconda3/envs/zi/lib/python3.8/subprocess.py", line 512, 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 "<stdin>", line 1, in <module>
File "/home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1119, in load_inline
return _jit_compile(
File "/home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1176, in _jit_compile
_write_ninja_file_and_build_library(
File "/home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1280, in _write_ninja_file_and_build_library
_run_ninja_build(
File "/home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1538, in _run_ninja_build
raise RuntimeError(message) from e
RuntimeError: Error building extension 'my_module'