fix issues with push down transpose from einsum

Author: Matthew Francis-Landau

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

@@ -164,7 +164,8 @@ struct PushDownTransposeActivationRewriter
     auto activationOp = rewriter.create<ActivationOp>(
         op.getLoc(), producer.getInput(), op.getActivationType(),
         op.getAlphaAttr(), op.getBetaAttr());
-    auto newTranspose = rewriter.create<TransposeOp>(producer.getLoc(), activationOp.getResult(), permutation);
+    auto newTranspose = rewriter.create<TransposeOp>(
+        producer.getLoc(), activationOp.getResult(), permutation);
     rewriter.replaceOp(op, newTranspose.getResult());
     return success();
   }
@@ -181,7 +182,8 @@ struct PushDownTransposeUnary : OpRewritePattern<UnaryOp> {
     AffineMap permutation = producer.getPermutation();
     auto unary = rewriter.create<UnaryOp>(op.getLoc(), producer.getInput(),
                                           op.getUnaryOperationAttr());
-    auto newTranspose = rewriter.create<TransposeOp>(producer.getLoc(), unary.getResult(), permutation);
+    auto newTranspose = rewriter.create<TransposeOp>(
+        producer.getLoc(), unary.getResult(), permutation);
     rewriter.replaceOp(op, newTranspose.getResult());
     return success();
   }
@@ -755,13 +757,23 @@ class EinsumPushDownTranspose : public OpRewritePattern<tensorrt::EinsumOp> {
                 return a.first < b.first;
               });
 
+    LLVM_DEBUG({
+      std::stringstream out;
+      out << "outputAxes: [";
+      for (auto x : outputAxes) {
+        out << x.first << "(" << x.second << ") ";
+      }
+      out << "]\n";
+      DBGS() << out.str();
+    });
+
     SmallVector<int64_t> newEinsumShape;
-    SmallVector<int64_t> outputPerm;
+    SmallVector<int64_t> forwardPerm;
     std::string newEinsumRhs = "";
     for (auto &[c, i] : outputAxes) {
       newEinsumRhs += c;
       newEinsumShape.push_back(op.getType().getDimSize(i));
-      outputPerm.push_back(i);
+      forwardPerm.push_back(i);
     }
     if (newEinsumRhs == equation.rhs)
       return failure(); // no change
@@ -773,10 +785,13 @@ class EinsumPushDownTranspose : public OpRewritePattern<tensorrt::EinsumOp> {
         op.getLoc(), op.getType().clone(newEinsumShape), op.getInputs(),
         newEinsumEquation);
 
+    auto forwardMap =
+        AffineMap::getPermutationMap(forwardPerm, op.getLoc().getContext());
+
     auto newTranspose = rewriter.create<tensorrt::TransposeOp>(
-        op.getLoc(), newEinsum.getResult(),
-        AffineMap::getPermutationMap(outputPerm, op.getLoc().getContext()));
+        op.getLoc(), newEinsum.getResult(), inversePermutation(forwardMap));
 
+    assert(op.getType() == newTranspose.getType());
     rewriter.replaceOp(op, newTranspose.getResult());
 
     return success();
@@ -1662,29 +1677,29 @@ class MoveReshapeBeforeTranspose
       }
     }
     assert(inputNumElems == outputNumElems);
-    while(j < reshapeOutputType.getRank()) {
+    while (j < reshapeOutputType.getRank()) {
       outputNumElems *= reshapeOutputType.getDimSize(j);
       groupReshapeOut.push_back(reshapeOutputType.getDimSize(j));
       transposeOutAxes.push_back(j++);
     }
     assert(inputNumElems == outputNumElems);
     assert(transposeInAxes.empty());
-    if(!transposeOutAxes.empty() || !groupReshapeOut.empty()) {
+    if (!transposeOutAxes.empty() || !groupReshapeOut.empty()) {
       reshapeGroups.push_back(ReshapeGroup{
           .transposeInAxes = transposeInAxes,
           .transposeOutAxes = transposeOutAxes,
           .reshapeOut = groupReshapeOut,
           .startOutputIdx = -1, // set later
-        });
+      });
     }
 
     SmallVector<int64_t> newTranspose;
     SmallVector<int64_t> newReshape;
 
     std::sort(reshapeGroups.begin(), reshapeGroups.end(), [](auto &a, auto &b) {
-      if(a.transposeInAxes.empty())
+      if (a.transposeInAxes.empty())
         return false;
-      if(b.transposeInAxes.empty())
+      if (b.transposeInAxes.empty())
         return true;
       return a.transposeInAxes[0] < b.transposeInAxes[0];
     });
@@ -1713,28 +1728,29 @@ class MoveReshapeBeforeTranspose
         out << "    transposeInAxes: [";
         for (size_t i = 0; i < group.transposeInAxes.size(); ++i) {
           out << group.transposeInAxes[i];
-          if (i + 1 < group.transposeInAxes.size()) out << ", ";
+          if (i + 1 < group.transposeInAxes.size())
+            out << ", ";
         }
         out << "]\n";
         out << "    transposeOutAxes: [";
         for (size_t i = 0; i < group.transposeOutAxes.size(); ++i) {
           out << group.transposeOutAxes[i];
-          if (i + 1 < group.transposeOutAxes.size()) out << ", ";
+          if (i + 1 < group.transposeOutAxes.size())
+            out << ", ";
         }
         out << "]\n";
         out << "    reshapeOut: [";
         for (size_t i = 0; i < group.reshapeOut.size(); ++i) {
           out << group.reshapeOut[i];
-          if (i + 1 < group.reshapeOut.size()) out << ", ";
+          if (i + 1 < group.reshapeOut.size())
+            out << ", ";
         }
         out << "]\n";
         out << "    startOutputIdx: " << group.startOutputIdx << "\n";
       }
       DBGS() << out.str();
     });
 
-
-
     for (auto &group : reshapeGroups) {
       for (size_t i = 0; i < group.reshapeOut.size(); i++)
         newTranspose.push_back(group.startOutputIdx + i);

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

@@ -287,4 +287,28 @@ func.func @reshape_transpose_reorder_ones_dim(%arg0: tensor<2x1x1x1x1xf32>, %arg
     %3 = tensorrt.reshape %2 : tensor<2x1x1x1x1xf32> to tensor<2x1x1x1xf32>
     %4 = tensorrt.deconvolution {dilation = array<i64: 1, 1>, num_groups = 2 : ui32, post_padding = array<i64: 0, 0>, pre_padding = array<i64: 0, 0>, stride = array<i64: 1, 2>} in(%arg1 : tensor<1x2x3x3xf32>) kernelWeights(%3 : tensor<2x1x1x1xf32>) -> tensor<1x2x3x5xf32>
     return %4 : tensor<1x2x3x5xf32>
+}
+
+// -----
+
+
+// CHECK: @push_down_transpose_einsum(%[[arg0:.+]]: tensor<1x6x1500x64xf32>, %[[arg1:.+]]: tensor<1x6x1500x1500xf32>) -> tensor<1x1500x384xf32>
+// CHECK-DAG: %[[const0:.+]] = tensorrt.constant dense<1.000000e+00> : tensor<384x6x64xf32>
+// CHECK-DAG: %[[v0:.+]] = tensorrt.collapse_rank %[[arg0]] : tensor<1x6x1500x64xf32> to tensor<6x1500x64xf32>
+// CHECK-DAG: %[[v1:.+]] = tensorrt.collapse_rank %[[arg1]] : tensor<1x6x1500x1500xf32> to tensor<6x1500x1500xf32>
+// CHECK: %[[v2:.+]] = tensorrt.matrix_multiply [[params:.+]] ins(%[[v0]], %[[v1]] : tensor<6x1500x64xf32>, tensor<6x1500x1500xf32>) -> tensor<6x64x1500xf32>
+// CHECK: %[[v3:.+]] = tensorrt.einsum [[params2:.+]] ins(%[[v2]], %[[const0]] : tensor<6x64x1500xf32>, tensor<384x6x64xf32>) -> tensor<1500x384xf32>
+// CHECK: %[[v4:.+]] = tensorrt.expand_rank %[[v3:.+]] : tensor<1500x384xf32> to tensor<1x1500x384xf32>
+// CHECK: return %[[v4]]
+func.func @push_down_transpose_einsum(%arg0: tensor<1x6x1500x64xf32>, %arg1: tensor<1x6x1500x1500xf32>) -> tensor<1x1500x384xf32> {
+  %cst_f32 = tensorrt.constant dense<1.000000e+00> : tensor<384x384xf32>
+  %0 = tensorrt.reshape %arg0 : tensor<1x6x1500x64xf32> to tensor<6x1500x64xf32>
+  %1 = tensorrt.reshape %arg1 : tensor<1x6x1500x1500xf32> to tensor<6x1500x1500xf32>
+  %2 = tensorrt.einsum {equation = "bcd,bec->ebd"} ins(%0, %1 : tensor<6x1500x64xf32>, tensor<6x1500x1500xf32>) -> tensor<1500x6x64xf32>
+  %3 = tensorrt.reshape %2 : tensor<1500x6x64xf32> to tensor<1x1500x6x64xf32>
+  %4 = tensorrt.reshape %2 : tensor<1500x6x64xf32> to tensor<1500x384xf32>
+  %cst_f32_0 = tensorrt.constant dense<1.000000e+00> : tensor<384x6x64xf32>
+  %5 = tensorrt.einsum {equation = "bde,cde->bc"} ins(%2, %cst_f32_0 : tensor<1500x6x64xf32>, tensor<384x6x64xf32>) -> tensor<1500x384xf32>
+  %6 = tensorrt.reshape %5 : tensor<1500x384xf32> to tensor<1x1500x384xf32>
+  return %6 : tensor<1x1500x384xf32>
 }