einsum -> matmul reorder arguments if matches pattern

Author: Matthew Francis-Landau

mlir-tensorrt/tensorrt/lib/TensorRT/Transforms/TransposeReshapeElimination.cpp

@@ -2043,6 +2043,8 @@ class EinsumToMatrixMultiply : public OpRewritePattern<tensorrt::EinsumOp> {
     char multipliedAxis = 0;
     std::string batchAxes = "";
 
+    Value inputs[2] = {op.getInputs()[0], op.getInputs()[1]};
+
     for (size_t i = 0; i < equation.rhs.size(); i++) {
       char c = equation.rhs[i];
       size_t lhsPos = equation.lhsParts[0].find(c);
@@ -2085,6 +2087,15 @@ class EinsumToMatrixMultiply : public OpRewritePattern<tensorrt::EinsumOp> {
       return failure(/* no multiplied axis */);
     }
 
+    if (matrixAxis[0] != 0 && matrixAxis[1] != 0 &&
+        equation.rhs.find(matrixAxis[0]) > equation.rhs.find(matrixAxis[1])) {
+      // the order of the arguments need to get swapped as the order for a
+      // matrix multiply requires the first matrix axis appears first
+      std::swap(equation.lhsParts[0], equation.lhsParts[1]);
+      std::swap(matrixAxis[0], matrixAxis[1]);
+      std::swap(inputs[0], inputs[1]);
+    }
+
     MatrixOperation opType[2];
     for (int i = 0; i < 2; i++) {
       std::string e = batchAxes;
@@ -2121,8 +2132,8 @@ class EinsumToMatrixMultiply : public OpRewritePattern<tensorrt::EinsumOp> {
     }
 
     rewriter.replaceOpWithNewOp<tensorrt::MatrixMultiplyOp>(
-        op, op.getResult().getType(), op.getInputs()[0], op.getInputs()[1],
-        opType[0], opType[1]);
+        op, op.getResult().getType(), inputs[0], inputs[1], opType[0],
+        opType[1]);
 
     return success();
   }

mlir-tensorrt/tensorrt/test/Dialect/TensorRT/transpose-reshape-elimination.mlir

@@ -202,3 +202,22 @@ func.func @elementwise_reshape(%arg0: tensor<12x3x3xf32>, %arg1: tensor<12xf32>)
   %2 = tensorrt.element_wise <kDIV>(%0, %1 : tensor<12x3x3xf32>, tensor<12x1x1xf32>) -> tensor<12x3x3xf32>
   return %2 : tensor<12x3x3xf32>
 }
+
+// -----
+
+// CHECK: @matmul_argument_swap(%[[arg0:.+]]: tensor<1x2x4x3x561xf32>, %[[arg1:.+]]: tensor<1x2x4x3x3xf32>) -> tensor<1x2x4x561x3xf32>
+// CHECK-DAG: %[[v0:.+]]= tensorrt.collapse_rank %[[arg1]] : tensor<1x2x4x3x3xf32> to tensor<2x4x3x3xf32>
+// CHECK-DAG: %[[v1:.+]] = tensorrt.transpose {permutation = #map} %[[arg0]] : tensor<1x2x4x3x561xf32> to tensor<2x4x561x3x1xf32>
+// CHECK-DAG: %[[v2:.+]] = tensorrt.collapse_rank %[[v1]] : tensor<2x4x561x3x1xf32> to tensor<2x4x561x3xf32>
+// CHECK: %[[v3:.+]] = tensorrt.matrix_multiply {op0 = #tensorrt.matrix_operation<kNONE>, op1 = #tensorrt.matrix_operation<kTRANSPOSE>} ins(%2, %0 : tensor<2x4x561x3xf32>, tensor<2x4x3x3xf32>) -> tensor<2x4x561x3xf32>
+// CHECK: %[[v4:.+]] = tensorrt.expand_rank %[[v3]] : tensor<2x4x561x3xf32> to tensor<1x2x4x561x3xf32>
+// CHECK: return %[[v4]]
+#map = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d4, d3)>
+func.func @matmul_argument_swap(%arg0: tensor<1x2x4x3x561xf32>, %arg1: tensor<1x2x4x3x3xf32>) -> tensor<1x2x4x561x3xf32> {
+  %0 = tensorrt.reshape %arg1 : tensor<1x2x4x3x3xf32> to tensor<8x3x3xf32>
+  %1 = tensorrt.reshape %arg0 : tensor<1x2x4x3x561xf32> to tensor<8x3x561xf32>
+  %2 = tensorrt.matrix_multiply {op0 = #tensorrt.matrix_operation<kNONE>, op1 = #tensorrt.matrix_operation<kNONE>} ins(%0, %1 : tensor<8x3x3xf32>, tensor<8x3x561xf32>) -> tensor<8x3x561xf32>
+  %3 = tensorrt.reshape %2 : tensor<8x3x561xf32> to tensor<1x2x4x3x561xf32>
+  %4 = tensorrt.transpose {permutation = #map} %3 : tensor<1x2x4x3x561xf32> to tensor<1x2x4x561x3xf32>
+  return %4 : tensor<1x2x4x561x3xf32>
+}