[Mlir-commits] [mlir] [mlir][Interfaces] Add interface methods to allow reifying single result/single dim of result. (PR #162924)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Oct 10 14:22:09 PDT 2025
https://github.com/MaheshRavishankar created https://github.com/llvm/llvm-project/pull/162924
Current implementation of `reifyResultShapes` forces all implementations to return all dimensions of all results. This can be wasteful when you only require dimensions of one result, or a single dimension of a result. This was initially done with the restriction in mind that it might not always be possible to compute dimension of a single result or one dimension of a single result in all cases. To handle such cases
- `reifyShapeOfResult` (which allows reifying dimensions of just one result), has a default implementation that calls `reifyResultShapes` and returns the dimensions of a single result.
- `reifyDimOfResult` (which allows reifying a single dimension of a single result) has a default implementation that calls `reifyDimOfResult` and returns the value for the dimension of the result (which in turn for the default case would call `reifyDimOfResult`).
Changes required downstream to integrate this change
1. In operation definitions in .td files, for those operations that implement the `ReifyRankedShapedTypeOpInterface`.
```
def <op-name> : Op<..., [...,
DeclareOpInterfaceMethods[ReifyRankedShapedTypeOpInterface]]>
```
should be changed to
```
def <op-name> : Op<..., [...,
DeclareOpInterfaceMethods[ReifyRankedShapedTypeOpInterface, [
"reifyResultShapes"]]]>
```
>From b9b193f062b0472992411578937c8558296e50d3 Mon Sep 17 00:00:00 2001
From: MaheshRavishankar <mahesh.ravishankar at gmail.com>
Date: Fri, 10 Oct 2025 14:12:59 -0700
Subject: [PATCH] [mlir][Interfaces] Add interface methods to allow reifying
single result/single dim of result.
Current implementation of `reifyResultShapes` forces all
implementations to return all dimensions of all results. This can be
wasteful when you only require dimensions of one result, or a single
dimension of a result. This was initially done with the restriction
in mind that it might not always be possible to compute dimension of a
single result or one dimension of a single result in all cases. To
handle such cases
- `reifyShapeOfResult` (which allows reifying dimensions of
just one result), has a default implementation that calls
`reifyResultShapes` and returns the dimensions of a single result.
- `reifyDimOfResult` (which allows reifying a single dimension of a
single result) has a default implementation that calls
`reifyDimOfResult` and returns the value for the dimension of the
result (which in turn for the default case would call
`reifyDimOfResult`).
Changes required downstream to integrate this change
1. In operation definitions in .td files, for those operations that
implement the `ReifyRankedShapedTypeOpInterface`.
```
def <op-name> : Op<..., [...,
DeclareOpInterfaceMethods[ReifyRankedShapedTypeOpInterface]]>
```
should be changed to
```
def <op-name> : Op<..., [...,
DeclareOpInterfaceMethods[ReifyRankedShapedTypeOpInterface, [
"reifyResultShapes"]]]>
```
Signed-off-by: MaheshRavishankar <mahesh.ravishankar at gmail.com>
---
.../Bufferization/IR/BufferizationOps.td | 6 +-
.../mlir/Dialect/Linalg/IR/LinalgOps.td | 3 +-
.../Dialect/Linalg/IR/LinalgRelayoutOps.td | 3 +-
.../mlir/Dialect/MemRef/IR/MemRefOps.td | 3 +-
.../mlir/Interfaces/InferTypeOpInterface.td | 64 +++++++++++++++++--
.../IR/TensorInferTypeOpInterfaceImpl.cpp | 3 +
6 files changed, 73 insertions(+), 9 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
index 6724d4c483101..a9b2b9f39519d 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
@@ -28,7 +28,8 @@ class Bufferization_Op<string mnemonic, list<Trait> traits = []>
def Bufferization_AllocTensorOp : Bufferization_Op<"alloc_tensor",
[AttrSizedOperandSegments, BufferizableOpInterface,
- DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>]> {
+ DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface, [
+ "reifyResultShapes"]>]> {
let summary = "allocate buffer for a tensor";
let description = [{
@@ -219,7 +220,8 @@ def Bufferization_MaterializeInDestinationOp
: Bufferization_Op<"materialize_in_destination",
[AllElementTypesMatch<["source", "dest"]>,
BufferizableOpInterface, DestinationStyleOpInterface,
- DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>,
+ DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface, [
+ "reifyResultShapes"]>,
DeclareOpInterfaceMethods<SubsetOpInterface,
["operatesOnEquivalentSubset", "operatesOnDisjointSubset"]>,
DeclareOpInterfaceMethods<SubsetInsertionOpInterface,
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td
index 7ff44c2e1d2ed..2754ee3b4f586 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td
@@ -94,7 +94,8 @@ def Linalg_IndexOp : Linalg_Op<"index", [Pure]>,
def Linalg_SoftmaxOp : Linalg_Op<"softmax",
[DestinationStyleOpInterface,
PredOpTrait<"input and output have same element type", TCopVTEtIsSameAs<0, 1>>,
- DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>,
+ DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface,
+ ["reifyResultShapes"]>,
DeclareOpInterfaceMethods<AggregatedOpInterface, ["decomposeOperation"]>,
DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
DeclareOpInterfaceMethods<TilingInterface,
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td
index 6504ca8664d49..238fa42cae427 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td
@@ -35,7 +35,8 @@ class Linalg_RelayoutOp<string mnemonic, list<Trait> traits = []> :
DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
DestinationStyleOpInterface, LinalgRelayoutOpInterface,
ConditionallySpeculatable, NoMemoryEffect,
- DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>,
+ DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface, [
+ "reifyResultShapes"]>,
TypesMatchWith<"result type matches type of dest",
"dest", "result",
"$_self">])> {
diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
index 40b7d7e33d5c2..a9911bdb1dcea 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
@@ -1776,7 +1776,8 @@ class MemRef_ReassociativeReshapeOp<string mnemonic, list<Trait> traits = []> :
def MemRef_ExpandShapeOp : MemRef_ReassociativeReshapeOp<"expand_shape", [
DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
DeclareOpInterfaceMethods<MemorySpaceCastConsumerOpInterface>,
- DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>]> {
+ DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface,
+ ["reifyResultShapes"]>]> {
let summary = "operation to produce a memref with a higher rank.";
let description = [{
The `memref.expand_shape` op produces a new view with a higher rank whose
diff --git a/mlir/include/mlir/Interfaces/InferTypeOpInterface.td b/mlir/include/mlir/Interfaces/InferTypeOpInterface.td
index 1a2c05fc16ed5..c5732f134e201 100644
--- a/mlir/include/mlir/Interfaces/InferTypeOpInterface.td
+++ b/mlir/include/mlir/Interfaces/InferTypeOpInterface.td
@@ -361,20 +361,76 @@ def ReifyRankedShapedTypeOpInterface :
let methods = [
InterfaceMethod<
/*desc=*/[{
- Reify the shape of the result of an operation (typically in terms of the
- shape of its operands).
+ Reify the shapes of all the result of an operation (typically in terms
+ of the shape of its operands).
`reifiedReturnShapes` is populated with one vector per op result. Each
of those vectors contains an OpFoldResult for each dimension of the
shaped type. The given builder may be used to insert ops that compute
result shapes.
- If the shape of a particular result cannot be computed it must be empty.
+ If the shape of a particular result cannot be computed it in terms of
+ its operands it must be left empty. If any dimension of the result cannot
+ be computed it must be set to OpFoldResult().
}],
/*retTy=*/"::llvm::LogicalResult",
/*methodName=*/"reifyResultShapes",
/*args=*/(ins "::mlir::OpBuilder &":$builder,
- "::mlir::ReifiedRankedShapedTypeDims &":$reifiedReturnShapes)
+ "::mlir::ReifiedRankedShapedTypeDims &":$reifiedReturnShapes),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{ return ::mlir::failure(); }]
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ Reify the shape of a single result of an operation (typically in terms
+ of the shape of its operands).
+
+ Returns the shape of a single result of the operation as a
+ `SmallVector<OpFoldResult>`, one per dimension of the shaped type. The
+ given builder may be used to insert ops that compute result shapes.
+
+ If any dimension of the result cannot be computed it must be set to
+ OpFoldResult().
+ }],
+ /*retTy=*/"::llvm::FailureOr<SmallVector<::mlir::OpFoldResult>>",
+ /*methodName=*/"reifyShapeOfResult",
+ /*args=*/(ins "::mlir::OpBuilder &":$builder,
+ "int":$resultIndex),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ ReifiedRankedShapedTypeDims reifiedShapes;
+ if (failed($_op.reifyResultShapes(builder, reifiedShapes)))
+ return failure();
+ if (resultIndex < 0 || resultIndex >= (int)(reifiedShapes.size()))
+ return $_op.emitOpError("invalid result index");
+ return reifiedShapes[resultIndex];
+ }]
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ Reify the shape of a dimension of a given result of an operation
+ (typically in terms of the shape of its operands).
+
+ Returns the shape of a specific dimension of a result of the operation as
+ an OpFoldResult. The given builder may be used to insert ops that compute
+ the shapes.
+
+ If the dimension of the result cannot be computed the method must return
+ `failure()`.
+ }],
+ /*retTy=*/"::llvm::FailureOr<::mlir::OpFoldResult>",
+ /*methodName=*/"reifyDimOfResult",
+ /*args=*/(ins "::mlir::OpBuilder &":$builder,
+ "int":$resultIndex, "int":$dim),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ auto shapes = $_op.reifyShapeOfResult(builder, resultIndex);
+ if (failed(shapes))
+ return failure();
+ if (dim < 0 || dim >= (int)((*shapes).size()))
+ return $_op.emitOpError("invalid dimension");
+ return (*shapes)[dim];
+ }]
>
];
}
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.cpp
index 4ec13e189f621..686f6eed1f8c7 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.cpp
@@ -77,6 +77,9 @@ namespace {
struct ReifyExpandShapeOp
: public ReifyRankedShapedTypeOpInterface::ExternalModel<ReifyExpandShapeOp,
ExpandShapeOp> {
+ using Base =
+ ReifyRankedShapedTypeOpInterface::ExternalModel<ReifyExpandShapeOp,
+ ExpandShapeOp>;
LogicalResult
reifyResultShapes(Operation *op, OpBuilder &b,
ReifiedRankedShapedTypeDims &reifyResultShapes) const {
More information about the Mlir-commits
mailing list