[Mlir-commits] [mlir] [mlir][tensor] Introduce `TensorRelayoutOpInterface` (PR #125823)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Feb 5 01:07:19 PST 2025
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-tensor
Author: Andrzej Warzyński (banach-space)
<details>
<summary>Changes</summary>
The newly introduced `TensorRelayoutOpInterface` is created specifically
for `tensor.pack` + `tensor.unpack`. Although the interface is
currently empty, it enables us to refactor the logic in
`FoldTensorCastProducerOp` within the Tensor dialect as follows:
```cpp
// OLD
// Reject tensor::PackOp - there's dedicated pattern for that instead.
if (!foldTensorCastPrecondition(op) ||
isa<tensor::PackOp, tensor::UnPackOp>(*op))
return failure();
```
is replaced with:
```cpp
// NEW
// Reject tensor::PackOp - there's dedicated pattern for that instead.
if (!foldTensorCastPrecondition(op) ||
isa<tensor::RelayoutOpInterface>(*op))
return failure();
```
This will be crucial once `tensor.pack` + `tensor.pack` are replaced
with `linalg.pack` + `linalg.unpack` (i.e. moved to Linalg):
* https://github.com/llvm/llvm-project/pull/123902,
* https://discourse.llvm.org/t/rfc-move-tensor-pack-and-tensor-unpack-into-linalg/.
Note that the interface itself will later be moved to the Linalg
dialect. This decoupling ensures that the Tensor dialect does not
require an understanding of Linalg ops, thus keeping the dependency
lightweight.
This PR is a preparatory step for moving Pack and Unpack ops to Linalg.
This PR is effectively a preparatory step for moving PackOp and UnpackOp
to Linalg. Once that's completed, most CMake changes from this PR will
be effectively reverted.
---
Full diff: https://github.com/llvm/llvm-project/pull/125823.diff
6 Files Affected:
- (modified) mlir/include/mlir/Dialect/Tensor/IR/CMakeLists.txt (+6)
- (modified) mlir/include/mlir/Dialect/Tensor/IR/Tensor.h (+6)
- (added) mlir/include/mlir/Dialect/Tensor/IR/TensorInterfaces.td (+33)
- (modified) mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td (+2)
- (modified) mlir/lib/Dialect/Tensor/IR/CMakeLists.txt (+1)
- (modified) mlir/lib/Dialect/Tensor/IR/TensorOps.cpp (+1-1)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Tensor/IR/CMakeLists.txt b/mlir/include/mlir/Dialect/Tensor/IR/CMakeLists.txt
index cd14fe5c0456118..74a05291376b3cb 100644
--- a/mlir/include/mlir/Dialect/Tensor/IR/CMakeLists.txt
+++ b/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)
diff --git a/mlir/include/mlir/Dialect/Tensor/IR/Tensor.h b/mlir/include/mlir/Dialect/Tensor/IR/Tensor.h
index bd96337a55407a8..1bd0f6553fc8d05 100644
--- a/mlir/include/mlir/Dialect/Tensor/IR/Tensor.h
+++ b/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
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorInterfaces.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorInterfaces.td
new file mode 100644
index 000000000000000..522a9c56f3c92c2
--- /dev/null
+++ b/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
diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
index 69db333f2c6912f..f6927f5ebcfb8e0 100644
--- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
+++ b/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">])> {
diff --git a/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt b/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt
index 5425615dac39324..d9d09d6361a2f0f 100644
--- a/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt
+++ b/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt
@@ -16,6 +16,7 @@ add_mlir_dialect_library(MLIRTensorDialect
DEPENDS
MLIRTensorOpsIncGen
+ MLIRTensorInterfacesIncGen
LINK_LIBS PUBLIC
MLIRAffineDialect
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 2a7d0c1a6ebfad5..fda6246334e15ce 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/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());
``````````
</details>
https://github.com/llvm/llvm-project/pull/125823
More information about the Mlir-commits
mailing list