[ONNX] Add support for Onnx.FusedMatMul op (#4147)

This commit adds the Onnx->Torch lowering for
[Onnx.FusedMatMul](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.FusedMatMul) op.

Signed-off-by: Vivek Khandelwal <vivekkhandelwal1424@gmail.com>

Author: vivekkhandelwal1

lib/Conversion/TorchOnnxToTorch/ComMicrosoftDomain.cpp

@@ -1009,4 +1009,60 @@ void mlir::torch::onnx_c::populateComMicrosoftDomain(
                                                           averagePool);
         return success();
       });
+  patterns.onOp(
+      "FusedMatMul", 1,
+      [](OpBinder binder, ConversionPatternRewriter &rewriter) {
+        Torch::ValueTensorType resultType;
+        Value lhs, rhs;
+        int64_t transA, transB, transBatchA, transBatchB;
+        if (binder.tensorOperands(lhs, rhs) ||
+            binder.s64IntegerAttr(transA, "transA", 0) ||
+            binder.s64IntegerAttr(transB, "transB", 0) ||
+            binder.s64IntegerAttr(transBatchA, "transBatchA", 0) ||
+            binder.s64IntegerAttr(transBatchB, "transBatchB", 0) ||
+            binder.tensorResultType(resultType))
+          return failure();
+
+        // Transposing the LHS argument.
+        Value transposedLhs = lhs;
+        if (transA) {
+          // Determine the rank of lhs tensor.
+          std::optional<unsigned> maybeRank = Torch::getTensorRank(lhs);
+          if (!maybeRank)
+            return rewriter.notifyMatchFailure(
+                binder.op, "Unimplemented: unranked lhs tensor");
+          unsigned lhsRank = *maybeRank;
+          if (failed(createTorchTransposeOp(
+                  rewriter, binder.getLoc(), lhs,
+                  /*dimA=*/lhsRank - 2, /*dimB=*/lhsRank - 1, transposedLhs)))
+            return rewriter.notifyMatchFailure(
+                binder.op, "Failed to create TorchTranspose op for lhs");
+        }
+
+        // Transposing the RHS argument.
+        Value transposedRhs = rhs;
+        if (transB) {
+          std::optional<unsigned> maybeRank = Torch::getTensorRank(rhs);
+          if (!maybeRank)
+            return rewriter.notifyMatchFailure(
+                binder.op, "Unimplemented: unranked rhs tensor");
+          unsigned rhsRank = *maybeRank;
+          if (failed(createTorchTransposeOp(
+                  rewriter, binder.getLoc(), rhs,
+                  /*dimA=*/rhsRank - 2, /*dimB=*/rhsRank - 1, transposedRhs)))
+            return rewriter.notifyMatchFailure(
+                binder.op, "Failed to create TorchTranspose op for rhs");
+        }
+
+        // TODO: Add support for `transBatchA` and `transBatchB`
+        // attribute.
+        if (transBatchA || transBatchB)
+          return rewriter.notifyMatchFailure(
+              binder.op, "Unimplemented: support not present for "
+                         "transBatchA and transBatchB attribute");
+
+        rewriter.replaceOpWithNewOp<Torch::AtenMatmulOp>(
+            binder.op, resultType, transposedLhs, transposedRhs);
+        return success();
+      });
 }

test/Conversion/TorchOnnxToTorch/simple_ops_a_to_f.mlir

@@ -2921,3 +2921,17 @@ func.func @test_dft_inverse_real(%arg0: !torch.vtensor<[10,10,1],f32>, %arg1: !t
   %0 = torch.operator "onnx.DFT"(%arg0, %none, %arg1) {torch.onnx.inverse = 1 : si64} : (!torch.vtensor<[10,10,1],f32>, !torch.none, !torch.vtensor<[],si64>) -> !torch.vtensor<[10,10,2],f32>
   return %0 : !torch.vtensor<[10,10,2],f32>
 }
+
+// -----
+
+// CHECK-LABEL: @test_fusedMatmul(
+// CHECK-SAME:                   %[[LHS:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !torch.vtensor<[?,12,256,64],f32>,
+// CHECK-SAME:                   %[[RHS:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !torch.vtensor<[?,12,256,64],f32>) -> !torch.vtensor<[?,12,256,256],f32>
+func.func @test_fusedMatmul(%arg0: !torch.vtensor<[?,12,256,64],f32>, %arg1: !torch.vtensor<[?,12,256,64],f32>) -> !torch.vtensor<[?,12,256,256],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 21 : si64, torch.onnx_meta.opset_versions = {com.microsoft = 1 : si64}} {
+    %0 = torch.operator "onnx.FusedMatMul"(%arg0, %arg1) {torch.onnx.alpha = 1.250000e-01 : f32, torch.onnx.transA = 0 : si64, torch.onnx.transB = 1 : si64} : (!torch.vtensor<[?,12,256,64],f32>, !torch.vtensor<[?,12,256,64],f32>) -> !torch.vtensor<[?,12,256,256],f32>
+    // CHECK: %[[DIMA:.*]] = torch.constant.int 2
+    // CHECK: %[[DIMB:.*]] = torch.constant.int 3
+    // CHECK: %[[TRANSPOSED_RHS:.*]] = torch.aten.transpose.int %[[RHS]], %[[DIMA]], %[[DIMB]] : !torch.vtensor<[?,12,256,64],f32>, !torch.int, !torch.int -> !torch.vtensor<[?,12,64,256],f32>
+    // CHECK: torch.aten.matmul %[[LHS]], %[[TRANSPOSED_RHS]] : !torch.vtensor<[?,12,256,64],f32>, !torch.vtensor<[?,12,64,256],f32> -> !torch.vtensor<[?,12,256,256],f32>
+    return %0 : !torch.vtensor<[?,12,256,256],f32>
+}