[mlir][tensor] Introduce TensorRelayoutOpInterface

mlir/include/mlir/Dialect/Tensor/IR/CMakeLists.txt

@@ -1,2 +1,8 @@
 add_mlir_dialect(TensorOps tensor)
 add_mlir_doc(TensorOps TensorOps Dialects/ -gen-dialect-doc)
+
+set(LLVM_TARGET_DEFINITIONS TensorInterfaces.td)
+mlir_tablegen(TensorInterfaces.h.inc -gen-op-interface-decls)
+mlir_tablegen(TensorInterfaces.cpp.inc -gen-op-interface-defs)
+add_public_tablegen_target(MLIRTensorInterfacesIncGen)
+add_dependencies(mlir-headers MLIRTensorInterfacesIncGen)

mlir/include/mlir/Dialect/Tensor/IR/Tensor.h

@@ -46,6 +46,12 @@ SmallVector<Range, 8> getOrCreateRanges(OffsetSizeAndStrideOpInterface op,
 
 #include "mlir/Dialect/Tensor/IR/TensorOpsDialect.h.inc"
 
+//===----------------------------------------------------------------------===//
+// Tensor Interfaces
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Tensor/IR/TensorInterfaces.h.inc"
+
 //===----------------------------------------------------------------------===//
 // Tensor Dialect Operations
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/Tensor/IR/TensorInterfaces.td

@@ -0,0 +1,33 @@
+//===- TensorInterfaces.td - Tensor Interfaces Declaration -*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This is the definition file for the structured interface sfor Tensor ops.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef TENSOR_IR_TENSORINTERFACES
+#define TENSOR_IR_TENSORINTERFACES
+
+include "mlir/Interfaces/DestinationStyleOpInterface.td"
+include "mlir/IR/OpBase.td"
+
+// TODO: To be moved to LinalgInterfaces.td, see:
+//  * https://github.com/llvm/llvm-project/pull/123902
+//  * https://discourse.llvm.org/t/rfc-move-tensor-pack-and-tensor-unpack-into-linalg/
+def TensorRelayoutOpInterface : OpInterface<"RelayoutOpInterface"> {
+  let description = [{
+    A Tensor (soon to be Linalg) relayout-op is either tensor.pack or
+    tensor.unpack.
+
+    While we could extend this interface with methods from Tensor_RelayoutOp,
+    this is currently not needed and left as a TODO.
+  }];
+  let cppNamespace = "::mlir::tensor";
+}
+
+#endif // TENSOR_IR_TENSORINTERFACES

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

@@ -10,6 +10,7 @@
 #define TENSOR_OPS
 
 include "mlir/Dialect/Tensor/IR/TensorBase.td"
+include "mlir/Dialect/Tensor/IR/TensorInterfaces.td"
 include "mlir/Interfaces/CastInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/DestinationStyleOpInterface.td"
@@ -1833,6 +1834,7 @@ class Tensor_RelayoutOp<string mnemonic, list<Trait> traits = []> :
         DestinationStyleOpInterface,
         ConditionallySpeculatable, NoMemoryEffect,
         DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>,
+        TensorRelayoutOpInterface,
         TypesMatchWith<"result type matches type of dest",
                    "dest", "result",
                    "$_self">])> {

mlir/lib/Dialect/Tensor/IR/CMakeLists.txt

@@ -16,6 +16,7 @@ add_mlir_dialect_library(MLIRTensorDialect
 
   DEPENDS
   MLIRTensorOpsIncGen
+  MLIRTensorInterfacesIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

@@ -4976,7 +4976,7 @@ struct FoldTensorCastProducerOp
 
     // Reject tensor::PackOp - there's dedicated pattern for that instead.
     if (!foldTensorCastPrecondition(op) ||
-        isa<tensor::PackOp, tensor::UnPackOp>(*op))
+        isa<tensor::RelayoutOpInterface>(*op))
       return failure();
 
     SmallVector<Type> newResultTypes(op->getResultTypes());