From f350dd269a195ab791c00d7aa7db147bc7b9ce78 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski@arm.com>
Date: Wed, 22 Jan 2025 13:46:35 +0000
Subject: [PATCH] [mlir][tensor][linalg] Move Pack/Unpack Ops to Linalg (extra)

This is merely moving code around, no new functionality is added.

PATCH 5: Create `LinalgRelayoutOpInterface` to be able to exclude
`linalg::PackOp` + `linalg::UnpackOp` Ops from patterns/folders outside
the Linalg dialect, e.g. `FoldTensorCastProducerOp` from the Tensor
dialect. Note that there's `FoldTensorCastUnPackOp` and
`FoldTensorCastPackOp` in LinalgOps.cpp (i.e. Linalg dialect) that
provides similar folder (but which fold "correctly"). See e.g. #121393
and #114559 for context.

CONTEXT:
This change was discussed in the following RFC:
* https://discourse.llvm.org/t/rfc-move-tensor-pack-and-tensor-unpack-into-linalg
---
 mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td  | 8 ++++++++
 mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td | 3 ++-
 mlir/lib/Dialect/Tensor/IR/TensorOps.cpp                 | 5 +++--
 3 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
index 244db23925ab3..5986626a72729 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
@@ -178,6 +178,14 @@ def LinalgConvolutionOpInterface : OpInterface<"ConvolutionOpInterface"> {
   ];
 }
 
+// TODO:
+def LinalgRelayoutOpInterface : OpInterface<"RelayoutOpInterface"> {
+  let description = [{
+    TODO
+  }];
+  let cppNamespace = "::mlir::linalg";
+}
+
 def LinalgFillOpInterface : OpInterface<"FillOpInterface"> {
   let description = [{
     A fill operation is defined in general terms:
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td
index 845a34e90bc09..fe0e826f6b771 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td
@@ -18,6 +18,7 @@ include "mlir/Dialect/Linalg/IR/LinalgBase.td"
 include "mlir/Interfaces/DestinationStyleOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
+include "mlir/Dialect/Linalg/IR/LinalgInterfaces.td"
 include "mlir/IR/OpAsmInterface.td"
 
 //===----------------------------------------------------------------------===//
@@ -27,7 +28,7 @@ include "mlir/IR/OpAsmInterface.td"
 class Linalg_RelayoutOp<string mnemonic, list<Trait> traits = []> :
       Op<Linalg_Dialect, mnemonic, !listconcat(traits, [
         DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
-        DestinationStyleOpInterface,
+        DestinationStyleOpInterface, LinalgRelayoutOpInterface,
         ConditionallySpeculatable, NoMemoryEffect,
         DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>,
         TypesMatchWith<"result type matches type of dest",
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 011f8601fa95f..34a3ebfe82f22 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -3935,9 +3935,10 @@ struct FoldTensorCastProducerOp
   LogicalResult matchAndRewrite(DestinationStyleOpInterface op,
                                 PatternRewriter &rewriter) const override {
 
-    // Reject PackOp/UnpackOp - there are dedicated patterns for that instead.
+    // Reject PackOp/UnpackOp (i.e. RelayoutOps) - there are dedicated patterns
+    // for that instead.
     if (!foldTensorCastPrecondition(op) ||
-        isa<linalg::PackOp, linalg::UnPackOp>(*op))
+        isa<linalg::RelayoutOpInterface>(*op))
       return failure();
 
     SmallVector<Type> newResultTypes(op->getResultTypes());