Skip to content

Commit

Permalink
[mlir][tensor] Introduce TensorRelayoutOpInterface (#125823)
Browse files Browse the repository at this point in the history
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):
  * #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 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.
  • Loading branch information
banach-space authored Feb 6, 2025
1 parent d2b45ce commit 80fd902
Show file tree
Hide file tree
Showing 6 changed files with 49 additions and 1 deletion.
6 changes: 6 additions & 0 deletions mlir/include/mlir/Dialect/Tensor/IR/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -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)
6 changes: 6 additions & 0 deletions mlir/include/mlir/Dialect/Tensor/IR/Tensor.h
Original file line number Diff line number Diff line change
Expand Up @@ -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
//===----------------------------------------------------------------------===//
Expand Down
33 changes: 33 additions & 0 deletions mlir/include/mlir/Dialect/Tensor/IR/TensorInterfaces.td
Original file line number Diff line number Diff line change
@@ -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
2 changes: 2 additions & 0 deletions mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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">])> {
Expand Down
1 change: 1 addition & 0 deletions mlir/lib/Dialect/Tensor/IR/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ add_mlir_dialect_library(MLIRTensorDialect

DEPENDS
MLIRTensorOpsIncGen
MLIRTensorInterfacesIncGen

LINK_LIBS PUBLIC
MLIRAffineDialect
Expand Down
2 changes: 1 addition & 1 deletion mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down

0 comments on commit 80fd902

Please sign in to comment.