[Mlir-commits] [mlir] [mlir] Add pack/unpack transpose foldings for linalg.generic ops, fix bugs (PR #93055)

[llvmlistbot at llvm.org]

================
@@ -48,6 +48,34 @@ static LogicalResult isPackOn1D(RewriterBase &rewriter, Operation *op,
   return success();
 }
 
+// If the `linalgOp` represents a transpose, return the permutation vector for
+// the transpose. Otherwise, return failure.
+static FailureOr<SmallVector<int64_t>>
+getTransposeOpPermutation(linalg::LinalgOp linalgOp) {
+  if (auto transposeOp = dyn_cast<linalg::TransposeOp>(linalgOp.getOperation()))
+    return SmallVector<int64_t>(transposeOp.getPermutation());
+  if (linalgOp.getNumParallelLoops() != linalgOp.getNumLoops())
+    return failure();
+
+  if (linalgOp.getNumDpsInputs() != 1 || linalgOp.getNumDpsInits() != 1)
+    return failure();
+  auto mapRange = linalgOp.getIndexingMapsArray();
+  if (mapRange.size() != 2 || !mapRange.front().isPermutation() ||
+      !mapRange.back().isPermutation() || mapRange.front() == mapRange.back()) {
+    return failure();
+  }
+  if (!llvm::hasSingleElement(linalgOp.getBlock()->getOperations()))
----------------
Max191 wrote:

I think for now we can keep it restricted to pure transpose ops for these patterns. Maybe we could add a new pattern for something like this in a later PR.

In general, a unary element-wise operation like that will get fused with it's producer/consumer in element-wise fusion. After this you end up with multi-input ops, and it may not be as simple to figure out whether it is beneficial to transpose the generic op. If you do this remapping before element-wise fusion, then I think there could be cases where you add extra transposing that would otherwise fold. Overall, I think it is not a bad idea to try to move the transpose work to the pack op, but I think it takes a bit more careful thought/matching, so I think it is better suited for a later PR.

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