Cpp_extension and Aten modules: compiling errors

ekoplin: <p>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?</p>
<pre><code class="lang-auto">cpp_source="""

#include &lt;torch/extension.h&gt;
#include &lt;ATen/ATen.h&gt;
#include &lt;ATen/CPUApplyUtils.h&gt;
#include &lt;ATen/Config.h&gt;
#include &lt;ATen/Dispatch.h&gt;
#include &lt;ATen/ExpandUtils.h&gt;
#include &lt;ATen/NativeFunctions.h&gt;
#include &lt;c10/util/Exception.h&gt;
#include &lt;c10/util/math_compat.h&gt;

#include &lt;ATen/Utils.h&gt;
#include &lt;ATen/CPUGenerator.h&gt;
#include &lt;ATen/core/DistributionsHelper.h&gt;
#include &lt;ATen/native/Distributions.h&gt;
#include &lt;ATen/native/DispatchStub.h&gt;
#include &lt;ATen/native/UnaryOps.h&gt;
#include &lt;ATen/NamedTensorUtils.h&gt;

#include &lt;type_traits&gt;
#include &lt;functional&gt;
#include &lt;assert.h&gt;
#include &lt;cpuinfo.h&gt;
#include &lt;float.h&gt;


int64_t sample_poisson(double lambda, at::CPUGenerator* generator) {
  TORCH_CHECK(lambda &gt;= 0, "invalid Poisson rate, expected rate to be non-negative");
  at::uniform_real_distribution&lt;double&gt; 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 &gt; enlam) {
		X += 1;
	  } else {
		return X;
	  }
	}
}

"""
</code></pre>
<pre><code class="lang-auto">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
)
</code></pre>
<p>I got the following error:</p>
<pre><code class="lang-auto">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 &lt;ATen/CPUGenerator.h&gt;
      |          ^~~~~~~~~~~~~~~~~~~~~
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.
</code></pre>
<p>The above exception was the direct cause of the following exception:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "&lt;stdin&gt;", line 1, in &lt;module&gt;
  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'
</code></pre>

ptrblck: <p>The error is raised in:</p>
<pre><code class="lang-python">/home/eric/Desktop/test_1/main.cpp:15:10: fatal error: ATen/CPUGenerator.h: No such file or directory
</code></pre>
<p>since <code>ATen/CPUGenerator.h</code> was moved in <a href="https://github.com/pytorch/pytorch/pull/36026">this PR</a>, so you would need to use the new file name.</p>

ekoplin: <p>Thank you very much. Now is working.</p>