Torchscript C++ Interface (#110)

* Began writing C++ inerface test.

* Switched to a CMake workflow.

* Making progress getting the test to link.

* We are having trouble linking the loaded library.

* Renamed files, linked in Python 3 appropriately.

* Successfully linked program and loaded in scripted module.

* Working on a nontrivial JITScript test.

* Call needs to be updated to work on a CUDA backend.

* Forward call works!

* Ready for linting, testing, and merge.

* Linted.

* Deleted README tutorial file, already covered by import_test.py.

* Added rpath linking, now need to fix the Python linkage issue.

* Removed print statements.

* Linted.

Author: vbharadwaj-bk

examples/readme_tutorial.py

@@ -10,14 +10,6 @@
 )
 from openequivariance.implementations.utils import torch_to_oeq_dtype
 
-__all__ = [
-    "TPProblem",
-    "Irreps",
-    "TensorProduct",
-    "TensorProductConv",
-    "torch_to_oeq_dtype",
-]
-
 __version__ = version("openequivariance")
 
 
@@ -30,3 +22,18 @@ def _check_package_editable():
 
 
 _editable_install_output_path = Path(__file__).parent.parent / "outputs"
+
+
+def torch_ext_so_path():
+    return openequivariance.extlib.torch_module.__file__
+
+
+__all__ = [
+    "TPProblem",
+    "Irreps",
+    "TensorProduct",
+    "TensorProductConv",
+    "torch_to_oeq_dtype",
+    "_check_package_editable",
+    "torch_ext_so_path",
+]

openequivariance/__init__.py

@@ -416,7 +416,7 @@ tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> jit_conv_doubl
 
 // =========================================================== 
 
-TORCH_LIBRARY_FRAGMENT(torch_tp_jit, m) { 
+TORCH_LIBRARY_FRAGMENT(libtorch_tp_jit, m) { 
     m.class_<TorchJITProduct>("TorchJITProduct")
         .def(torch::init<string, Map_t, Map_t, Map_t, Map_t>())
         .def("__obj_flatten__", &TorchJITProduct::__obj_flatten__)
@@ -437,9 +437,9 @@ TORCH_LIBRARY_FRAGMENT(torch_tp_jit, m) {
                 return c10::make_intrusive<TorchJITProduct>(get<0>(state), get<1>(state), get<2>(state), get<3>(state), get<4>(state));
             });
 
-    m.def("jit_tp_forward(__torch__.torch.classes.torch_tp_jit.TorchJITProduct jit, Tensor L1_in, Tensor L2_in, Tensor W) -> Tensor");
-    m.def("jit_tp_backward(__torch__.torch.classes.torch_tp_jit.TorchJITProduct jit, Tensor L1_in, Tensor L2_in, Tensor W, Tensor L3_grad) -> (Tensor, Tensor, Tensor)");
-    m.def("jit_tp_double_backward(__torch__.torch.classes.torch_tp_jit.TorchJITProduct jit, Tensor L1_in, Tensor L2_in, Tensor W, Tensor L3_grad, Tensor L1_dgrad, Tensor L2_dgrad, Tensor W_dgrad) -> (Tensor, Tensor, Tensor, Tensor)");
+    m.def("jit_tp_forward(__torch__.torch.classes.libtorch_tp_jit.TorchJITProduct jit, Tensor L1_in, Tensor L2_in, Tensor W) -> Tensor");
+    m.def("jit_tp_backward(__torch__.torch.classes.libtorch_tp_jit.TorchJITProduct jit, Tensor L1_in, Tensor L2_in, Tensor W, Tensor L3_grad) -> (Tensor, Tensor, Tensor)");
+    m.def("jit_tp_double_backward(__torch__.torch.classes.libtorch_tp_jit.TorchJITProduct jit, Tensor L1_in, Tensor L2_in, Tensor W, Tensor L3_grad, Tensor L1_dgrad, Tensor L2_dgrad, Tensor W_dgrad) -> (Tensor, Tensor, Tensor, Tensor)");
 
 
     m.class_<TorchJITConv>("TorchJITConv")
@@ -462,12 +462,12 @@ TORCH_LIBRARY_FRAGMENT(torch_tp_jit, m) {
                 return c10::make_intrusive<TorchJITConv>(get<0>(state), get<1>(state), get<2>(state), get<3>(state), get<4>(state));
             });
 
-    m.def("jit_conv_forward(__torch__.torch.classes.torch_tp_jit.TorchJITConv jit, Tensor L1_in, Tensor L2_in, Tensor W, Tensor rows, Tensor cols, Tensor workspace, Tensor transpose_perm) -> Tensor");
-    m.def("jit_conv_backward(__torch__.torch.classes.torch_tp_jit.TorchJITConv jit, Tensor L1_in, Tensor L2_in, Tensor W, Tensor L3_grad, Tensor rows, Tensor cols, Tensor workspace, Tensor transpose_perm) -> (Tensor, Tensor, Tensor)");
-    m.def("jit_conv_double_backward(__torch__.torch.classes.torch_tp_jit.TorchJITConv jit, Tensor L1_in, Tensor L2_in, Tensor W, Tensor L3_grad, Tensor L1_dgrad, Tensor L2_dgrad, Tensor W_dgrad, Tensor rows, Tensor cols, Tensor workspace, Tensor transpose_perm) -> (Tensor, Tensor, Tensor, Tensor)");
+    m.def("jit_conv_forward(__torch__.torch.classes.libtorch_tp_jit.TorchJITConv jit, Tensor L1_in, Tensor L2_in, Tensor W, Tensor rows, Tensor cols, Tensor workspace, Tensor transpose_perm) -> Tensor");
+    m.def("jit_conv_backward(__torch__.torch.classes.libtorch_tp_jit.TorchJITConv jit, Tensor L1_in, Tensor L2_in, Tensor W, Tensor L3_grad, Tensor rows, Tensor cols, Tensor workspace, Tensor transpose_perm) -> (Tensor, Tensor, Tensor)");
+    m.def("jit_conv_double_backward(__torch__.torch.classes.libtorch_tp_jit.TorchJITConv jit, Tensor L1_in, Tensor L2_in, Tensor W, Tensor L3_grad, Tensor L1_dgrad, Tensor L2_dgrad, Tensor W_dgrad, Tensor rows, Tensor cols, Tensor workspace, Tensor transpose_perm) -> (Tensor, Tensor, Tensor, Tensor)");
 };
 
-TORCH_LIBRARY_IMPL(torch_tp_jit, CUDA, m) { 
+TORCH_LIBRARY_IMPL(libtorch_tp_jit, CUDA, m) { 
     m.impl("jit_tp_forward", &jit_tp_forward);
     m.impl("jit_tp_backward", &jit_tp_backward);
     m.impl("jit_tp_double_backward", &jit_tp_double_backward);
@@ -477,4 +477,4 @@ TORCH_LIBRARY_IMPL(torch_tp_jit, CUDA, m) {
     m.impl("jit_conv_double_backward", &jit_conv_double_backward);
 };
 
-PYBIND11_MODULE(torch_tp_jit, m) {}
+PYBIND11_MODULE(libtorch_tp_jit, m) {}

…nequivariance/extension/torch_tp_jit.cpp …uivariance/extension/libtorch_tp_jit.cppopenequivariance/extension/torch_tp_jit.cpp renamed to openequivariance/extension/libtorch_tp_jit.cpp openequivariance/extension/torch_tp_jit.cpp renamed to openequivariance/extension/libtorch_tp_jit.cpp

@@ -0,0 +1,9 @@
+cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
+project(test_oeq_jitscript_load)
+
+find_package(Torch REQUIRED)
+
+add_executable(load_jitscript load_jitscript.cpp)
+target_link_libraries(load_jitscript "${TORCH_LIBRARIES}")
+target_link_libraries(load_jitscript -Wl,--no-as-needed "${OEQ_EXTLIB}") 
+set_property(TARGET load_jitscript PROPERTY CXX_STANDARD 17)

openequivariance/extension/test/CMakeLists.txt

@@ -0,0 +1,62 @@
+#include <torch/script.h>
+
+#include <iostream>
+#include <memory>
+
+/* 
+* This program takes in two JITScript modules that execute 
+* a tensor product in FP32 precision. 
+* The first module is compiled from e3nn, the second is
+* OEQ's compiled module. The program checks that the
+* two outputs are comparable. 
+*/
+
+int main(int argc, const char* argv[]) {
+    if (argc != 7) {
+        std::cerr << "usage: load_jitscript "
+                    << "<path-to-e3nn-module> "
+                    << "<path-to-oeq-module> "
+                    << "<L1_dim> "
+                    << "<L2_dim> "
+                    << "<weight_numel> "
+                    << "<batch_size> "
+                    << std::endl;
+
+        return 1;
+    }
+
+    int64_t L1_dim = std::stoi(argv[3]);
+    int64_t L2_dim = std::stoi(argv[4]);
+    int64_t weight_numel = std::stoi(argv[5]);
+    int64_t batch_size = std::stoi(argv[6]); 
+
+    torch::Device device(torch::kCUDA);
+    std::vector<torch::jit::IValue> inputs;
+    inputs.push_back(torch::randn({batch_size, L1_dim}, device));
+    inputs.push_back(torch::randn({batch_size, L2_dim}, device));
+    inputs.push_back(torch::randn({batch_size, weight_numel}, device));
+
+    torch::jit::script::Module module_e3nn, module_oeq;
+    try {
+        module_e3nn = torch::jit::load(argv[1]);
+        module_oeq = torch::jit::load(argv[2]);
+    }
+    catch (const c10::Error& e) {
+        std::cerr << "error loading script module" << std::endl;
+        return 1;
+    }
+
+    module_e3nn.to(device);
+    module_oeq.to(device);
+
+    at::Tensor output_e3nn = module_e3nn.forward(inputs).toTensor();
+    at::Tensor output_oeq = module_oeq.forward(inputs).toTensor();
+
+    if(at::allclose(output_e3nn, output_oeq, 1e-5, 1e-5)) {
+        return 0;
+    } 
+    else {
+        std::cerr << "torch.allclose returned FALSE comparing model outputs." << std::endl;
+        return 1;
+    }
+}

openequivariance/extension/test/load_jitscript.cpp

@@ -1,9 +1,11 @@
 # ruff: noqa : F401, E402
+import sys
 import os
 import warnings
 from pathlib import Path
 
 from openequivariance.benchmark.logging_utils import getLogger
+from distutils import sysconfig
 
 oeq_root = str(Path(__file__).parent.parent)
 
@@ -12,6 +14,18 @@
 torch_module, generic_module = None, None
 postprocess_kernel = lambda kernel: kernel  # noqa : E731
 
+try:
+    python_lib_dir = sysconfig.get_config_var("LIBDIR")
+    major, minor = sys.version_info.major, sys.version_info.minor
+    python_lib_name = f"python{major}.{minor}"
+
+except Exception as e:
+    print("Error while retrieving Python library information:", file=sys.stderr)
+    print(e, file=sys.stderr)
+    print("Syconfig variable list:", file=sys.stderr)
+    print(sysconfig.get_config_vars(), file=sys.stderr)
+    exit(1)
+
 if not build_ext:
     from openequivariance.extlib.generic_module import (
         GenericTensorProductImpl,
@@ -32,14 +46,23 @@
 
     extra_cflags = ["-O3"]
     generic_sources = ["generic_module.cpp"]
-    torch_sources = ["torch_tp_jit.cpp"]
+    torch_sources = ["libtorch_tp_jit.cpp"]
+
+    include_dirs, extra_link_args = (
+        ["util"],
+        [
+            f"-Wl,--no-as-needed,-rpath,{python_lib_dir}",
+            f"-L{python_lib_dir}",
+            f"-l{python_lib_name}",
+        ],
+    )
 
-    include_dirs, extra_link_args = ["util"], None
     if torch.version.cuda:
-        extra_link_args = ["-Wl,--no-as-needed", "-lcuda", "-lcudart", "-lnvrtc"]
+        extra_link_args.extend(["-lcuda", "-lcudart", "-lnvrtc"])
 
         try:
-            cuda_libs = library_paths("cuda")[1]
+            torch_libs, cuda_libs = library_paths("cuda")
+            extra_link_args.append("-Wl,-rpath," + torch_libs)
             extra_link_args.append("-L" + cuda_libs)
             if os.path.exists(cuda_libs + "/stubs"):
                 extra_link_args.append("-L" + cuda_libs + "/stubs")
@@ -48,7 +71,9 @@
 
         extra_cflags.append("-DCUDA_BACKEND")
     elif torch.version.hip:
-        extra_link_args = ["-Wl,--no-as-needed", "-lhiprtc"]
+        extra_link_args.extend(["-lhiprtc"])
+        torch_libs = library_paths("cuda")[0]
+        extra_link_args.append("-Wl,-rpath," + torch_libs)
 
         def postprocess(kernel):
             kernel = kernel.replace("__syncwarp();", "__threadfence_block();")
@@ -72,7 +97,7 @@ def postprocess(kernel):
 
         try:
             torch_module = torch.utils.cpp_extension.load(
-                "torch_tp_jit",
+                "libtorch_tp_jit",
                 torch_sources,
                 extra_cflags=extra_cflags,
                 extra_include_paths=include_dirs,

openequivariance/extlib/__init__.py

@@ -103,7 +103,7 @@ def generate_double_backward_schedule(warps_per_block):
             global torch
             import torch
 
-            internal_cls = torch.classes.torch_tp_jit.TorchJITProduct
+            internal_cls = torch.classes.libtorch_tp_jit.TorchJITProduct
         else:
             internal_cls = extlib.JITTPImpl
 
@@ -142,7 +142,7 @@ def register_torch_fakes(cls):
         global torch
         import torch
 
-        @torch._library.register_fake_class("torch_tp_jit::TorchJITProduct")
+        @torch._library.register_fake_class("libtorch_tp_jit::TorchJITProduct")
         class TorchJITProduct:
             def __init__(
                 self,
@@ -198,19 +198,19 @@ def backward_rawptr(
             ):
                 pass
 
-        @torch.library.register_fake("torch_tp_jit::jit_tp_forward")
+        @torch.library.register_fake("libtorch_tp_jit::jit_tp_forward")
         def fake_forward(jit, L1_in, L2_in, W):
             return L1_in.new_empty(
                 L1_in.shape[0], jit.wrapped_obj.kernel_dims["L3_dim"]
             )
 
-        @torch.library.register_fake("torch_tp_jit::jit_tp_backward")
+        @torch.library.register_fake("libtorch_tp_jit::jit_tp_backward")
         def fake_backward(jit, L1_in, L2_in, W, L3_grad):
             return torch.empty_like(L1_in), torch.empty_like(L2_in), torch.empty_like(W)
 
     @classmethod
     def register_autograd(cls):
-        backward_op = torch.ops.torch_tp_jit.jit_tp_backward
+        backward_op = torch.ops.libtorch_tp_jit.jit_tp_backward
 
         def setup_context(ctx, inputs, output):
             ctx.jit, ctx.L1_in, ctx.L2_in, ctx.weights = inputs
@@ -222,20 +222,20 @@ def backward(ctx, grad_output):
             return None, L1_grad, L2_grad, W_grad
 
         torch.library.register_autograd(
-            "torch_tp_jit::jit_tp_forward", backward, setup_context=setup_context
+            "libtorch_tp_jit::jit_tp_forward", backward, setup_context=setup_context
         )
 
         def setup_context_double_backward(ctx, inputs, output):
             ctx.jit, ctx.L1_in, ctx.L2_in, ctx.weights, ctx.L3_grad = inputs
 
         def double_backward(ctx, E, F, G):
-            result = torch.ops.torch_tp_jit.jit_tp_double_backward(
+            result = torch.ops.libtorch_tp_jit.jit_tp_double_backward(
                 ctx.jit, ctx.L1_in, ctx.L2_in, ctx.weights, ctx.L3_grad, E, F, G
             )
             return None, result[0], result[1], result[2], result[3]
 
         torch.library.register_autograd(
-            "torch_tp_jit::jit_tp_backward",
+            "libtorch_tp_jit::jit_tp_backward",
             double_backward,
             setup_context=setup_context_double_backward,
         )

openequivariance/implementations/LoopUnrollTP.py

@@ -20,4 +20,4 @@ def name():
     def forward(
         self, L1: torch.Tensor, L2: torch.Tensor, W: torch.Tensor
     ) -> torch.Tensor:
-        return torch.ops.torch_tp_jit.jit_tp_forward(self.internal, L1, L2, W)
+        return torch.ops.libtorch_tp_jit.jit_tp_forward(self.internal, L1, L2, W)