Einsum converter

core/conversion/converters/BUILD

@@ -62,6 +62,7 @@ cc_library(
         "impl/constant_pad.cpp",
         "impl/conv_deconv.cpp",
         "impl/cumsum.cpp",
+        "impl/einsum.cpp",
         "impl/element_wise.cpp",
         "impl/expand.cpp",
         "impl/interpolate.cpp",

core/conversion/converters/impl/einsum.cpp

@@ -0,0 +1,59 @@
+#include "core/conversion/converters/converters.h"
+#include "core/conversion/tensorcontainer/TensorContainer.h"
+#include "core/util/prelude.h"
+
+#include <vector>
+
+namespace torch_tensorrt {
+namespace core {
+namespace conversion {
+namespace converters {
+namespace impl {
+namespace {
+
+auto stack_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pattern(
+    {"aten::einsum(str equation, Tensor[] tensors) -> (Tensor)",
+     [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+       // Extract equation and list of tensors
+       auto equation = args[0].unwrapToString();
+       auto in = args[1].IValue()->toListRef();
+
+       std::vector<nvinfer1::ITensor*> tensors;
+
+       // Populate vector of ITensor pointers
+       for (auto t : in) {
+         nvinfer1::ITensor* itensor;
+
+         // Tensor is either an ITensor (wrapped) or PyTorch Tensor
+         if (t.isTensor()) {
+           auto weight = Weights(ctx, t.toTensor());
+
+           auto const_layer = ctx->net->addConstant(weight.shape, weight.data);
+           TORCHTRT_CHECK(const_layer, "Unable to create constant layer from node: " << *n);
+
+           itensor = const_layer->getOutput(0);
+         } else {
+           auto cont = t.toCustomClass<TensorContainer>();
+           itensor = cont->tensor();
+         }
+
+         tensors.push_back(itensor);
+       }
+
+       // Add Tensor-RT Einsum layer
+       auto einsum_layer = ctx->net->addEinsum(tensors.data(), tensors.size(), equation.c_str());
+       TORCHTRT_CHECK(einsum_layer, "Unable to create einsum layer from node: " << *n);
+
+       einsum_layer->setName(util::node_info(n).c_str());
+       auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], einsum_layer->getOutput(0));
+
+       LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
+       return true;
+     }});
+
+} // namespace
+} // namespace impl
+} // namespace converters
+} // namespace conversion
+} // namespace core
+} // namespace torch_tensorrt

tests/core/conversion/converters/BUILD

@@ -51,6 +51,10 @@ converter_test(
     name = "test_cumsum",
 )
 
+converter_test(
+    name = "test_einsum",
+)
+
 converter_test(
     name = "test_element_wise",
 )
@@ -152,6 +156,7 @@ test_suite(
         ":test_conv_deconv",
         ":test_copy",
         ":test_cumsum",
+        ":test_einsum",
         ":test_element_wise",
         ":test_expand",
         ":test_instance_norm",

tests/core/conversion/converters/test_einsum.cpp

@@ -0,0 +1,104 @@
+#include <string>
+#include "core/compiler.h"
+#include "gtest/gtest.h"
+#include "tests/util/util.h"
+#include "torch/csrc/jit/ir/irparser.h"
+
+TEST(Converters, ATenEinsumConvertsMatMulCorrectly) {
+  const auto graph = R"IR(
+      graph(%x.1 : Tensor, %x.2 : Tensor):
+        %0 : str = prim::Constant[value="ij,jk->ik"]()
+        %3 : Tensor[] = prim::ListConstruct(%x.1, %x.2)
+        %4 : Tensor = aten::einsum(%0, %3)
+        return (%4))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  // Test matrix multiplication via einsum
+  auto in_0 = at::rand({12, 17}, {at::kCUDA});
+  auto in_1 = at::rand({17, 35}, {at::kCUDA});
+
+  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in_0, in_1});
+
+  params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in_0, in_1});
+
+  ASSERT_TRUE(
+      torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
+}
+
+TEST(Converters, ATenEinsumConvertsElementwiseProdCorrectly) {
+  const auto graph = R"IR(
+      graph(%x.1 : Tensor, %x.2 : Tensor):
+        %0 : str = prim::Constant[value="abcd,abcd->abcd"]()
+        %3 : Tensor[] = prim::ListConstruct(%x.1, %x.2)
+        %4 : Tensor = aten::einsum(%0, %3)
+        return (%4))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  // Test elementwise tensor product via einsum
+  auto in_0 = at::rand({7, 5, 2, 8}, {at::kCUDA});
+  auto in_1 = at::rand({7, 5, 2, 8}, {at::kCUDA});
+
+  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in_0, in_1});
+
+  params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in_0, in_1});
+
+  ASSERT_TRUE(
+      torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
+}
+
+TEST(Converters, ATenEinsumConvertsTransposeCorrectly) {
+  const auto graph = R"IR(
+      graph(%x.1 : Tensor):
+        %0 : str = prim::Constant[value="jk->kj"]()
+        %3 : Tensor[] = prim::ListConstruct(%x.1)
+        %4 : Tensor = aten::einsum(%0, %3)
+        return (%4))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  // Test single-matrix transpose via einsum
+  auto in_0 = at::rand({25, 28}, {at::kCUDA});
+
+  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in_0});
+
+  params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in_0});
+
+  ASSERT_TRUE(
+      torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
+}
+
+TEST(Converters, ATenEinsumConvertsVectorsCorrectly) {
+  const auto graph = R"IR(
+      graph(%x.1 : Tensor, %x.2 : Tensor):
+        %0 : str = prim::Constant[value="a,b->ab"]()
+        %3 : Tensor[] = prim::ListConstruct(%x.1, %x.2)
+        %4 : Tensor = aten::einsum(%0, %3)
+        return (%4))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  // Test vector outer product via einsum
+  auto in_0 = at::rand({25}, {at::kCUDA});
+  auto in_1 = at::rand({4}, {at::kCUDA});
+
+  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in_0, in_1});
+
+  params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in_0, in_1});
+
+  ASSERT_TRUE(
+      torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
+}