handle edge case when there are still axes that need to get processed on the output for the reshape groups

Author: Matthew Francis-Landau

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

@@ -1635,7 +1635,8 @@ class MoveReshapeBeforeTranspose
     SmallVector<int64_t> groupReshapeOut;
     size_t inputNumElems = 1;
     size_t outputNumElems = 1;
-    for (int i = 0, j = 0; i < reshapeInputType.getRank(); i++) {
+    int j = 0;
+    for (int i = 0; i < reshapeInputType.getRank(); i++) {
       inputNumElems *= reshapeInputType.getDimSize(i);
       if (!transposeInAxes.empty() &&
           transposeInAxes.back() + 1 != transposePerm[i]) {
@@ -1740,7 +1741,8 @@ class MoveTransposeBeforeReshape
     SmallVector<int64_t> groupReshapeOut;
     size_t inputNumElems = 1;
     size_t outputNumElems = 1;
-    for (int i = 0, j = 0; i < reshapeInputType.getRank(); i++) {
+    int j = 0;
+    for (int i = 0; i < reshapeInputType.getRank(); i++) {
       inputNumElems *= reshapeInputType.getDimSize(i);
       inputAxes.push_back(i);
       while (j < reshapeOutputType.getRank() &&
@@ -1764,6 +1766,21 @@ class MoveTransposeBeforeReshape
         groupReshapeOut.clear();
       }
     }
+    while (j < reshapeOutputType.getRank()) {
+      outputNumElems *= reshapeOutputType.getDimSize(j);
+      groupReshapeOut.push_back(reshapeOutputType.getDimSize(j));
+      outputAxes.push_back(transposePerm[j++]);
+    }
+
+    assert(inputNumElems == outputNumElems);
+    assert(inputAxes.empty());
+    if (!outputAxes.empty() || !groupReshapeOut.empty()) {
+      reshapeGroups.push_back(ReshapeGroup{
+          .inputAxes = inputAxes,
+          .outputAxes = outputAxes,
+          .reshapeOut = groupReshapeOut,
+      });
+    }
 
     SmallVector<int64_t> newTranspose;
     SmallVector<int64_t> newReshape;
@@ -1776,6 +1793,38 @@ class MoveTransposeBeforeReshape
       return a.outputAxes[0] < b.outputAxes[0];
     });
 
+    // Debug print of reshapeGroups
+    LLVM_DEBUG({
+      std::stringstream out;
+      out << "reshapeGroups:\n";
+      for (size_t idx = 0; idx < reshapeGroups.size(); ++idx) {
+        const auto &group = reshapeGroups[idx];
+        out << "  Group " << idx << ":\n";
+        out << "    inputAxes: [";
+        for (size_t i = 0; i < group.inputAxes.size(); ++i) {
+          out << group.inputAxes[i];
+          if (i + 1 < group.inputAxes.size())
+            out << ", ";
+        }
+        out << "]\n";
+        out << "    outputAxes: [";
+        for (size_t i = 0; i < group.outputAxes.size(); ++i) {
+          out << group.outputAxes[i];
+          if (i + 1 < group.outputAxes.size())
+            DBGS() << ", ";
+        }
+        out << "]\n";
+        out << "    reshapeOut: [";
+        for (size_t i = 0; i < group.reshapeOut.size(); ++i) {
+          out << group.reshapeOut[i];
+          if (i + 1 < group.reshapeOut.size())
+            DBGS() << ", ";
+        }
+        out << "]\n";
+      }
+      DBGS() << out.str();
+    });
+
     for (auto &group : reshapeGroups) {
       for (int64_t i : group.inputAxes)
         newTranspose.push_back(i);
@@ -1789,6 +1838,8 @@ class MoveTransposeBeforeReshape
     Value newReshapeOp = rewriter.createOrFold<tensorrt::ReshapeOp>(
         op.getLoc(), reshapeInputType.clone(newReshape), newTransposeOp);
 
+    assert(op.getType() == newReshapeOp.getType());
+
     rewriter.replaceOp(op, newReshapeOp);
     return success();
   }
@@ -2244,15 +2295,15 @@ class MatrixMultiplyTransposedArguments
           return std::make_tuple(arg, operation);
         bool swapsLastTwo = true;
         for (int64_t i = 0; i < rank - 2; ++i) {
-          auto expr = permVec[i].dyn_cast<AffineDimExpr>();
+          auto expr = dyn_cast<AffineDimExpr>(permVec[i]);
           if (!expr || expr.getPosition() != i) {
             swapsLastTwo = false;
             break;
           }
         }
         if (swapsLastTwo) {
-          auto expr1 = permVec[rank - 2].dyn_cast<AffineDimExpr>();
-          auto expr2 = permVec[rank - 1].dyn_cast<AffineDimExpr>();
+          auto expr1 = dyn_cast<AffineDimExpr>(permVec[rank - 2]);
+          auto expr2 = dyn_cast<AffineDimExpr>(permVec[rank - 1]);
           if (!(expr1 && expr2 && expr1.getPosition() == rank - 1 &&
                 expr2.getPosition() == rank - 2)) {
             swapsLastTwo = false;

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

@@ -187,3 +187,18 @@ func.func @matmul_eliminate_reshape_lhs_2(%arg0: tensor<1x2x3x4x5x6xf16>, %arg1:
     %2 = tensorrt.reshape %1 : tensor<1x2x60x8xf16> to tensor<1x2x3x4x5x8xf16>
     return %2: tensor<1x2x3x4x5x8xf16>
 }
+
+// -----
+
+// CHECK: @elementwise_reshape(%[[arg0:.+]]: tensor<12x3x3xf32>, %[[arg1:.+]]: tensor<12xf32>)
+// CHECK: %[[v0:.+]] = tensorrt.expand_rank %[[arg1]] : tensor<12xf32> to tensor<12x1x1xf32>
+// CHECK: %[[v1:.+]] = tensorrt.element_wise <kDIV>(%[[arg0]], %[[v0]] : tensor<12x3x3xf32>, tensor<12x1x1xf32>) -> tensor<12x3x3xf32>
+// CHECK: %[[v2:.+]] = tensorrt.transpose {permutation = #map} %[[v1]] : tensor<12x3x3xf32> to tensor<12x3x3xf32>
+// CHECK: return %[[v2]]
+#map = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
+func.func @elementwise_reshape(%arg0: tensor<12x3x3xf32>, %arg1: tensor<12xf32>) -> tensor<12x3x3xf32> {
+  %0 = tensorrt.transpose {permutation = #map} %arg0 : tensor<12x3x3xf32> to tensor<12x3x3xf32>
+  %1 = tensorrt.expand_rank %arg1 : tensor<12xf32> to tensor<12x1x1xf32>
+  %2 = tensorrt.element_wise <kDIV>(%0, %1 : tensor<12x3x3xf32>, tensor<12x1x1xf32>) -> tensor<12x3x3xf32>
+  return %2 : tensor<12x3x3xf32>
+}