Support aten::bmm converter

Signed-off-by: Ruoqian Guo <[email protected]>

Author: ruoqianguo

core/conversion/converters/impl/matrix_multiply.cpp

@@ -8,24 +8,66 @@ namespace converters {
 namespace impl {
 namespace {
 
-auto mm_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns().pattern(
-    {"aten::matmul(Tensor self, Tensor other) -> (Tensor)",
-     [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
-       auto self = args[0].ITensorOrFreeze(ctx);
-       LOG_DEBUG("self tensor shape: " << self->getDimensions());
-
-       auto other = args[1].ITensorOrFreeze(ctx);
-       LOG_DEBUG("other tensor shape: " << other->getDimensions());
-
-       auto mm_layer = ctx->net->addMatrixMultiply(
-           *self, nvinfer1::MatrixOperation::kNONE, *other, nvinfer1::MatrixOperation::kNONE);
-       TRTORCH_CHECK(mm_layer, "Unable to create matrix multiplication node: " << *n);
-       mm_layer->setName(util::node_info(n).c_str());
-       auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], mm_layer->getOutput(0));
-
-       LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
-       return true;
-     }});
+auto mm_registrations TRTORCH_UNUSED =
+    RegisterNodeConversionPatterns()
+        .pattern({"aten::matmul(Tensor self, Tensor other) -> (Tensor)",
+                  [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+                    auto self = args[0].ITensorOrFreeze(ctx);
+                    LOG_DEBUG("self tensor shape: " << self->getDimensions());
+
+                    auto other = args[1].ITensorOrFreeze(ctx);
+                    LOG_DEBUG("other tensor shape: " << other->getDimensions());
+
+                    auto mm_layer = ctx->net->addMatrixMultiply(
+                        *self, nvinfer1::MatrixOperation::kNONE, *other, nvinfer1::MatrixOperation::kNONE);
+                    TRTORCH_CHECK(mm_layer, "Unable to create matrix multiplication node: " << *n);
+                    mm_layer->setName(util::node_info(n).c_str());
+                    auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], mm_layer->getOutput(0));
+
+                    LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
+                    return true;
+                  }})
+        .pattern(
+            {"aten::bmm(Tensor self, Tensor mat2) -> (Tensor)",
+             [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+               auto self = args[0].ITensorOrFreeze(ctx);
+               nvinfer1::Dims selfDims = self->getDimensions();
+               auto mat2 = args[1].ITensorOrFreeze(ctx);
+               nvinfer1::Dims mat2Dims = mat2->getDimensions();
+
+               // check dimensions
+               TRTORCH_CHECK(
+                   selfDims.nbDims == 3,
+                   "Expected 3-dimensional tensor, but got "
+                       << selfDims.nbDims
+                       << "-dimensional tensor for argument #1 'batch1' (while checking arguments for bmm)");
+               TRTORCH_CHECK(
+                   mat2Dims.nbDims == 3,
+                   "Expected 3-dimensional tensor, but got "
+                       << mat2Dims.nbDims
+                       << "-dimensional tensor for argument #2 'batch2' (while checking arguments for bmm)");
+
+               // Self and mat2 should have same size at dimension 0
+               TRTORCH_CHECK(
+                   selfDims.d[0] == mat2Dims.d[0],
+                   "Expected tensor to have size " << selfDims.d[0] << " at dimension 0, but got size " << mat2Dims.d[0]
+                                                   << " for argument #2 'batch2' (while checking arguments for bmm)");
+               // The size of mat2 at dimension 1 should be the same as that of self at dimension 2.
+               TRTORCH_CHECK(
+                   selfDims.d[2] == mat2Dims.d[1],
+                   "Expected tensor to have size " << selfDims.d[2] << " at dimension 1, but got size " << mat2Dims.d[1]
+                                                   << " for argument #2 'batch2' (while checking arguments for bmm)");
+
+               auto mm_layer = ctx->net->addMatrixMultiply(
+                   *self, nvinfer1::MatrixOperation::kNONE, *mat2, nvinfer1::MatrixOperation::kNONE);
+               TRTORCH_CHECK(mm_layer, "Unable to create matrix multiplication node: " << *n);
+
+               mm_layer->setName(util::node_info(n).c_str());
+               auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], mm_layer->getOutput(0));
+
+               LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
+               return true;
+             }});
 } // namespace
 } // namespace impl
 } // namespace converters

tests/core/conversion/converters/test_matrix_multiply.cpp

@@ -25,3 +25,24 @@ TEST(Converters, ATenMMConvertsCorrectly) {
 
   ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt, 2e-6));
 }
+
+TEST(Converters, ATenBMMConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor, %1 : Tensor):
+        %2 : Tensor = aten::bmm(%0, %1)
+        return (%2))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  auto in1 = at::randint(0, 5, {4, 64, 128}, {at::kCUDA});
+  auto in2 = at::randint(0, 5, {4, 128, 64}, {at::kCUDA});
+  auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto jit_results = trtorch::tests::util::RunGraph(g, params, {in1, in2});
+
+  params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in1, in2});
+  auto trt = trt_results[0].reshape_as(jit_results[0]);
+
+  ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt, 2e-6));
+}