From 83285487edb33496828af69f76746e7fc1c11fee Mon Sep 17 00:00:00 2001 From: Andrzej Warzynski 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(*op)) return failure(); ``` is replaced with: ```cpp // NEW // Reject tensor::PackOp - there's dedicated pattern for that instead. if (!foldTensorCastPrecondition(op) || isa(*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 cd14fe5c04561..74a05291376b3 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 bd96337a55407..1bd0f6553fc8d 100644 --- a/mlir/include/mlir/Dialect/Tensor/IR/Tensor.h +++ b/mlir/include/mlir/Dialect/Tensor/IR/Tensor.h @@ -46,6 +46,12 @@ SmallVector 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 0000000000000..522a9c56f3c92 --- /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 69db333f2c691..f6927f5ebcfb8 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 traits = []> : DestinationStyleOpInterface, ConditionallySpeculatable, NoMemoryEffect, DeclareOpInterfaceMethods, + 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 5425615dac393..d9d09d6361a2f 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 2a7d0c1a6ebfa..fda6246334e15 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(*op)) + isa(*op)) return failure(); SmallVector newResultTypes(op->getResultTypes());