[Mlir-commits] [mlir] [MLIR][Tensor, MemRef] Fold expand_shape and collapse_shape if identity (PR #80658)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Feb 5 02:51:31 PST 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-memref
@llvm/pr-subscribers-mlir
Author: James Newling (newling)
<details>
<summary>Changes</summary>
Before: op verifiers failed if the input and output ranks were the same (i.e. no expansion or collapse). This behaviour requires users of these shape ops to verify manually that they are not creating identity versions of these ops every time they build them -- problematic. This PR removes this strict verification, and introduces folders for the the identity cases.
---
Full diff: https://github.com/llvm/llvm-project/pull/80658.diff
7 Files Affected:
- (modified) mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h (-2)
- (modified) mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp (+4-8)
- (modified) mlir/lib/Dialect/Tensor/IR/TensorOps.cpp (+4-12)
- (modified) mlir/test/Dialect/MemRef/canonicalize.mlir (+19)
- (modified) mlir/test/Dialect/MemRef/invalid.mlir (-14)
- (modified) mlir/test/Dialect/Tensor/canonicalize.mlir (+20-1)
- (modified) mlir/test/Dialect/Tensor/invalid.mlir (-14)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h b/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
index 61c929dee0f27..3a672a1cc6060 100644
--- a/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
+++ b/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
@@ -111,8 +111,6 @@ static LogicalResult verifyReshapeLikeTypes(Op op, T expandedType,
<< " to have higher rank than the type = " << collapsedType;
if (expandedRank == 0)
return op.emitOpError("expected non-zero memref ranks");
- if (expandedRank == collapsedRank)
- return op.emitOpError("expected to collapse or expand dims");
if (collapsedRank == 0) {
// If collapsed rank is 0, then expanded type must be static shaped and of
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index b79ab8f3d671e..b79022d675bd6 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -2223,10 +2223,6 @@ LogicalResult ExpandShapeOp::verify() {
MemRefType srcType = getSrcType();
MemRefType resultType = getResultType();
- if (srcType.getRank() >= resultType.getRank())
- return emitOpError("expected rank expansion, but found source rank ")
- << srcType.getRank() << " >= result rank " << resultType.getRank();
-
// Verify result shape.
if (failed(verifyCollapsedShape(getOperation(), srcType.getShape(),
resultType.getShape(),
@@ -2377,10 +2373,6 @@ LogicalResult CollapseShapeOp::verify() {
MemRefType srcType = getSrcType();
MemRefType resultType = getResultType();
- if (srcType.getRank() <= resultType.getRank())
- return emitOpError("expected rank reduction, but found source rank ")
- << srcType.getRank() << " <= result rank " << resultType.getRank();
-
// Verify result shape.
if (failed(verifyCollapsedShape(getOperation(), resultType.getShape(),
srcType.getShape(), getReassociationIndices(),
@@ -2455,11 +2447,15 @@ void CollapseShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
}
OpFoldResult ExpandShapeOp::fold(FoldAdaptor adaptor) {
+ if (getSrcType() == getType())
+ return getSrc();
return foldReshapeOp<ExpandShapeOp, CollapseShapeOp>(*this,
adaptor.getOperands());
}
OpFoldResult CollapseShapeOp::fold(FoldAdaptor adaptor) {
+ if (getSrcType() == getType())
+ return getSrc();
return foldReshapeOp<CollapseShapeOp, ExpandShapeOp>(*this,
adaptor.getOperands());
}
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index b2fe58099b2fb..591cbcde80084 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -1656,22 +1656,10 @@ static LogicalResult verifyTensorReshapeOp(TensorReshapeOp op,
}
LogicalResult ExpandShapeOp::verify() {
- auto srcType = getSrcType();
- auto resultType = getResultType();
- if (srcType.getRank() >= resultType.getRank())
- return emitOpError("expected rank expansion, but found source rank ")
- << srcType.getRank() << " >= result rank " << resultType.getRank();
-
return verifyTensorReshapeOp(*this, getResultType(), getSrcType());
}
LogicalResult CollapseShapeOp::verify() {
- auto srcType = getSrcType();
- auto resultType = getResultType();
- if (srcType.getRank() <= resultType.getRank())
- return emitOpError("expected rank reduction, but found source rank ")
- << srcType.getRank() << " <= result rank " << resultType.getRank();
-
return verifyTensorReshapeOp(*this, getSrcType(), getResultType());
}
@@ -1872,11 +1860,15 @@ void CollapseShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
}
OpFoldResult ExpandShapeOp::fold(FoldAdaptor adaptor) {
+ if (getSrcType() == getType())
+ return getSrc();
return foldReshapeOp<ExpandShapeOp, CollapseShapeOp>(*this,
adaptor.getOperands());
}
OpFoldResult CollapseShapeOp::fold(FoldAdaptor adaptor) {
+ if (getSrcType() == getType())
+ return getSrc();
return foldReshapeOp<CollapseShapeOp, ExpandShapeOp>(*this,
adaptor.getOperands());
}
diff --git a/mlir/test/Dialect/MemRef/canonicalize.mlir b/mlir/test/Dialect/MemRef/canonicalize.mlir
index eccfc485b2034..8c0b765c1b123 100644
--- a/mlir/test/Dialect/MemRef/canonicalize.mlir
+++ b/mlir/test/Dialect/MemRef/canonicalize.mlir
@@ -1,5 +1,24 @@
// RUN: mlir-opt %s -canonicalize="test-convergence" --split-input-file -allow-unregistered-dialect | FileCheck %s
+
+// CHECK-LABEL: collapse_shape_identity_fold
+// CHECK-NEXT: return
+func.func @collapse_shape_identity_fold(%arg0 : memref<5xi8>) -> memref<5xi8> {
+ %0 = memref.collapse_shape %arg0 [[0]] : memref<5xi8> into memref<5xi8>
+ return %0 : memref<5xi8>
+}
+
+// -----
+
+// CHECK-LABEL: expand_shape_identity_fold
+// CHECK-NEXT: return
+func.func @expand_shape_identity_fold(%arg0 : memref<5x4xi8>) -> memref<5x4xi8> {
+ %0 = memref.expand_shape %arg0 [[0], [1]] : memref<5x4xi8> into memref<5x4xi8>
+ return %0 : memref<5x4xi8>
+}
+
+// -----
+
// CHECK-LABEL: func @subview_of_size_memcast
// CHECK-SAME: %[[ARG0:.[a-z0-9A-Z_]+]]: memref<4x6x16x32xi8>
// CHECK: %[[S:.+]] = memref.subview %[[ARG0]][0, 1, 0, 0] [1, 1, 16, 32] [1, 1, 1, 1] : memref<4x6x16x32xi8> to memref<16x32xi8, strided{{.*}}>
diff --git a/mlir/test/Dialect/MemRef/invalid.mlir b/mlir/test/Dialect/MemRef/invalid.mlir
index 7bb7a2affcbd1..53841fa836da7 100644
--- a/mlir/test/Dialect/MemRef/invalid.mlir
+++ b/mlir/test/Dialect/MemRef/invalid.mlir
@@ -415,20 +415,6 @@ func.func @collapse_shape_out_of_bounds(%arg0: memref<?x?xf32>) {
// -----
-func.func @expand_shape_invalid_ranks(%arg0: memref<?x?xf32>) {
- // expected-error @+1 {{op expected rank expansion, but found source rank 2 >= result rank 2}}
- %0 = memref.expand_shape %arg0 [[0], [1]] : memref<?x?xf32> into memref<?x?xf32>
-}
-
-// -----
-
-func.func @collapse_shape_invalid_ranks(%arg0: memref<?x?xf32>) {
- // expected-error @+1 {{op expected rank reduction, but found source rank 2 <= result rank 2}}
- %0 = memref.collapse_shape %arg0 [[0], [1]] : memref<?x?xf32> into memref<?x?xf32>
-}
-
-// -----
-
func.func @expand_shape_out_of_bounds(%arg0: memref<?xf32>) {
// expected-error @+1 {{op reassociation index 2 is out of bounds}}
%0 = memref.expand_shape %arg0 [[0, 1, 2]] : memref<?xf32> into memref<4x?xf32>
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index ed964071358ac..9b31d858d63c5 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -1,5 +1,24 @@
// RUN: mlir-opt %s -split-input-file -canonicalize="test-convergence" | FileCheck %s
+
+// CHECK-LABEL: expand_shape_identity_fold
+// CHECK-NEXT: return
+func.func @expand_shape_identity_fold(%arg0 : tensor<5xf32>) -> tensor<5xf32> {
+ %0 = tensor.expand_shape %arg0 [[0]] : tensor<5xf32> into tensor<5xf32>
+ return %0 : tensor<5xf32>
+}
+
+// -----
+
+// CHECK-LABEL: collapse_shape_identity_fold
+// CHECK-NEXT: return
+func.func @collapse_shape_identity_fold(%arg0 : tensor<5x4xf32>) -> tensor<5x4xf32> {
+ %0 = tensor.collapse_shape %arg0 [[0], [1]] : tensor<5x4xf32> into tensor<5x4xf32>
+ return %0 : tensor<5x4xf32>
+}
+
+// -----
+
// CHECK-LABEL: @tensor_bitcast_chain_ok
// CHECK-SAME: %[[IN:.*]]: tensor<2xi32>
func.func @tensor_bitcast_chain_ok(%input: tensor<2xi32>) -> tensor<2xf32> {
@@ -1977,7 +1996,7 @@ func.func @unpack_pack(%t: tensor<128x128xf32>) -> tensor<128x128xf32> {
// Chain: NC -> NCnc -> NCnc -> NC
// CHECK: func.func @unpack_pack(
-// CHECK-SAME: %[[T:.+]]: tensor<128x128xf32>,
+// CHECK-SAME: %[[T:.+]]: tensor<128x128xf32>,
// CHECK: return %[[T]] : tensor<128x128xf32>
func.func @unpack_pack(%t: tensor<128x128xf32>, %tile1: index, %tile2: index) -> tensor<128x128xf32> {
%tensor_empty = tensor.empty(%tile1, %tile2) : tensor<16x16x?x?xf32>
diff --git a/mlir/test/Dialect/Tensor/invalid.mlir b/mlir/test/Dialect/Tensor/invalid.mlir
index 735e5146e9dbc..39cd3788cb081 100644
--- a/mlir/test/Dialect/Tensor/invalid.mlir
+++ b/mlir/test/Dialect/Tensor/invalid.mlir
@@ -343,20 +343,6 @@ func.func @illegal_collapsing_reshape_mixed_tensor_2(%arg0 : tensor<?x4x5xf32>)
// -----
-func.func @expand_shape_invalid_ranks(%arg0: tensor<?x?xf32>) {
- // expected-error @+1 {{op expected rank expansion, but found source rank 2 >= result rank 2}}
- %0 = tensor.expand_shape %arg0 [[0], [1]] : tensor<?x?xf32> into tensor<?x?xf32>
-}
-
-// -----
-
-func.func @collapse_shape_invalid_ranks(%arg0: tensor<?x?xf32>) {
- // expected-error @+1 {{op expected rank reduction, but found source rank 2 <= result rank 2}}
- %0 = tensor.collapse_shape %arg0 [[0], [1]] : tensor<?x?xf32> into tensor<?x?xf32>
-}
-
-// -----
-
func.func @rank(%0: f32) {
// expected-error at +1 {{'tensor.rank' op operand #0 must be tensor of any type values}}
"tensor.rank"(%0): (f32)->index
``````````
</details>
https://github.com/llvm/llvm-project/pull/80658
More information about the Mlir-commits
mailing list