[Mlir-commits] [mlir] [mlir][tensor] Fold producer linalg transpose with consumer tensor pack (PR #75658)

[lorenzo chelini]

================
@@ -21,6 +21,57 @@ static bool areAllConstantIntValue(ArrayRef<OpFoldResult> ofrs, int64_t value) {
       ofrs, [&](OpFoldResult ofr) { return isConstantIntValue(ofr, value); });
 }
 
+/// Helper function to generate an equivalent permutation map for
+/// `linalg.transpose` and `tensor.pack` which will be used after their folding
+/// into a `tensor.pack`.
+static bool getRemappedPermutationForTransposeAndPack(
+    PackOp packOp, linalg::TransposeOp transposeOp,
+    SmallVector<int64_t> &newOuterDimsPermVec,
+    SmallVector<int64_t> &newInnerDimsPosVec,
+    SmallVector<OpFoldResult> &newMixedInnerTilesVec,
+    bool isTransposeProducer) {
+  bool foldingPossible = true;
+  auto innerDimsPos = packOp.getInnerDimsPos();
+  auto mixedInnerTiles = packOp.getMixedTiles();
+  auto outerDimsPerm = packOp.getOuterDimsPerm();
+  auto transposePerm = transposeOp.getPermutation();
+  int64_t srcRank = packOp.getSourceRank();
+
+  // Note: if isTransposeProducer = true, transposePerm.size() = srcRank, else
+  // transposePerm.size() > srcRank
+
+  // Process transpose operation for non-tiled outer dimensions
+  for (unsigned int i = 0; i < srcRank; ++i) {
+    int64_t remappedPosition =
+        isTransposeProducer ? (!outerDimsPerm.empty() ? outerDimsPerm[i] : i)
+                            : transposePerm[i];
+
+    if (remappedPosition >= srcRank) {
+      foldingPossible = false;
+      return foldingPossible;
+    }
+
+    remappedPosition =
+        isTransposeProducer
+            ? transposePerm[remappedPosition]
+            : (!outerDimsPerm.empty() ? outerDimsPerm[remappedPosition]
+                                      : remappedPosition);
+
+    newOuterDimsPermVec.push_back(remappedPosition);
+  }
+
+  // Process transpose operation for tiled inner dimensions
+  for (unsigned int i = srcRank; i < srcRank + mixedInnerTiles.size(); ++i) {
+    int64_t remappedPosition =
+        isTransposeProducer ? i - srcRank : transposePerm[i] - srcRank;
+
+    newMixedInnerTilesVec.push_back(mixedInnerTiles[remappedPosition]);
+    newInnerDimsPosVec.push_back(innerDimsPos[remappedPosition]);
+  }
+
+  return foldingPossible;
----------------
chelini wrote:

simply return true?

https://github.com/llvm/llvm-project/pull/75658