[Mlir-commits] [mlir] [mlir][tensor] Introduce `TensorRelayoutOpInterface` (PR #125823)
Andrzej Warzyński
llvmlistbot at llvm.org
Wed Feb 5 01:05:29 PST 2025
https://github.com/banach-space created https://github.com/llvm/llvm-project/pull/125823
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.
>From 83285487edb33496828af69f76746e7fc1c11fee Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Mon, 3 Feb 2025 08:55:40 +0000
Subject: [PATCH] [mlir][tensor] Introduce `TensorRelayoutOpInterface`
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.
---
.../mlir/Dialect/Tensor/IR/CMakeLists.txt | 6 ++++
mlir/include/mlir/Dialect/Tensor/IR/Tensor.h | 6 ++++
.../Dialect/Tensor/IR/TensorInterfaces.td | 33 +++++++++++++++++++
.../mlir/Dialect/Tensor/IR/TensorOps.td | 2 ++
mlir/lib/Dialect/Tensor/IR/CMakeLists.txt | 1 +
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 2 +-
6 files changed, 49 insertions(+), 1 deletion(-)
create mode 100644 mlir/include/mlir/Dialect/Tensor/IR/TensorInterfaces.td
diff --git a/mlir/include/mlir/Dialect/Tensor/IR/CMakeLists.txt b/mlir/include/mlir/Dialect/Tensor/IR/CMakeLists.txt
index cd14fe5c045611..74a05291376b3c 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 bd96337a55407a..1bd0f6553fc8d0 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 00000000000000..522a9c56f3c92c
--- /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 69db333f2c6912..f6927f5ebcfb8e 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 5425615dac3932..d9d09d6361a2f0 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 2a7d0c1a6ebfad..fda6246334e15c 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());
More information about the Mlir-commits
mailing list