Pytorch c++ extension

(Birendra Kathariya) #1

I have a c++ code that computes morton-code of n 3d points.

std::vector<uint64_t> computeMortonCode(std::vector<std::vector> points) {
std::vector<uint64_t> mortonCode(points.size());
for (int i = 0; i < points.size(); i++) {
auto point = points[i];
mortonCode[i] = mortonAddr(point[0], point[1], point[2]);
return mortonCode;

here function “mortonAddr” returns morton code given x,y,z coordinate of a point.

What changes should I make in this code so that I can call this function from python such that input points and output morton-code are in python torch variable? Many thanks in advance!!
ps: what other changes would this require if I want to leverage power of cuda.


You want to create an extension. This tutorial will get you through all the details.

Basically, there are 3 steps:

  1. Write your extension code
  2. Bind it to python using pybind11
  3. Build it using setuptools and torch.utils.cpp_extension

You can see another example of pytorch extension in the first commit of this repo. There is only one cpp extension there, it’s maybe easier to follow.

Tell me if you need any further help.