export pipeline added

CMakeLists.txt

@@ -0,0 +1,11 @@
+cmake_minimum_required(VERSION 3.18 FATAL_ERROR)
+project(aoti_example)
+
+find_package(CUDA)
+
+find_package(Torch REQUIRED)
+
+add_executable(inference inference.cpp)
+
+target_link_libraries(inference "${TORCH_LIBRARIES}" m)
+set_property(TARGET inference PROPERTY CXX_STANDARD 17)

Makefile

@@ -0,0 +1,18 @@
+train:
+	python train.py
+
+build:
+	cmake -Bbuild -DCMAKE_PREFIX_PATH=`python3 -c 'import torch;print(torch.utils.cmake_prefix_path)'`
+	cmake --build build
+
+clean:
+	rm -rf build
+
+profile:
+	nsys profile --capture-range=cudaProfilerApi --cuda-graph-trace=node --capture-range-end=stop -o profile -f true python train.py
+
+run:
+	./build/inference export/model.so
+
+.PHONY: run nsys build clean
+

inference.cpp

@@ -0,0 +1,25 @@
+#include <iostream>
+#include <vector>
+
+#include <torch/torch.h>
+#include <torch/csrc/inductor/aoti_runner/model_container_runner_cuda.h>
+
+int main(int argc, char *argv[]) {
+
+    char *model_path = NULL;  // e.g. mode.so
+    model_path = argv[1];
+
+    c10::InferenceMode mode;
+    torch::inductor::AOTIModelContainerRunnerCuda *runner;
+    runner = new torch::inductor::AOTIModelContainerRunnerCuda(model_path, 1);
+    std::vector<torch::Tensor> inputs = {
+        torch::randn({32,1}, at::kCUDA),
+        torch::randn({32,3}, at::kCUDA),
+        torch::randn({2,50}, at::kCUDA),
+        torch::randn({32,1}, at::kCUDA)
+    };
+    std::vector<torch::Tensor> outputs = runner->run(inputs);
+    std::cout << "Result from the first inference:"<< std::endl;
+    std::cout << outputs[0] << std::endl;
+    return 0;
+}

model.py

@@ -6,6 +6,7 @@
 
 from e3nn import o3
 
+from torch_geometric.data import Data
 from torch_scatter import scatter_mean, scatter_sum
 
 class Layer(nn.Module):
@@ -102,9 +103,13 @@ def __init__(self):
 
       self.layers = torch.nn.ModuleList(layers)
 
-  def forward(self, graphs):
-
-    node_features, pos, edge_index, batch, num_nodes = graphs.numbers, graphs.pos, graphs.edge_index, graphs.batch, graphs.num_nodes
+  def forward(self,
+              node_features,
+              pos,
+              edge_index,
+              batch):
+
+    # Passing in graphs make dynamo angry
     senders, receivers = edge_index
     relative_positions= pos[receivers] - pos[senders]