Skip to content

[mlir][memref.expand_shape] Add verifier check to ensure correct output_shape is provided by user #91245

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
May 7, 2024
+21 −0
Merged

[mlir][memref.expand_shape] Add verifier check to ensure correct output_shape is provided by user #91245

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2353,6 +2353,16 @@ LogicalResult ExpandShapeOp::verify() {
<< " dynamic dims while output_shape has " << getOutputShape().size()
<< " values";

// Verify if provided output shapes are in agreement with output type.
DenseI64ArrayAttr staticOutputShapes = getStaticOutputShapeAttr();
ArrayRef<int64_t> resShape = getResult().getType().getShape();
unsigned staticShapeNum = 0;

for (auto [pos, shape] : llvm::enumerate(resShape))
if (!ShapedType::isDynamic(shape) &&
shape != staticOutputShapes[staticShapeNum++])
emitOpError("invalid output shape provided at pos ") << pos;

Comment on lines +2359 to +2365
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this correct? Concretely, this verifier is producing an error when resShape and staticOutputShape are equal, and contain a mix of static and kDynamic values. Indeed, since in C++ the && operator is a sequence point, the staticShapeNum++ increment is only happening if the left hand side, !ShapedType::isDynamic(shape), evaluated to true. In other words, this verifier code requires staticOutputShape to only list static dimensions and omit kDynamic values.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If my understanding is correct, I would suggest this modification:

diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index 78201ae29cd9..4f33137dcff4 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -2356,11 +2356,8 @@ LogicalResult ExpandShapeOp::verify() {
   // Verify if provided output shapes are in agreement with output type.
   DenseI64ArrayAttr staticOutputShapes = getStaticOutputShapeAttr();
   ArrayRef<int64_t> resShape = getResult().getType().getShape();
-  unsigned staticShapeNum = 0;
-
   for (auto [pos, shape] : llvm::enumerate(resShape))
-    if (!ShapedType::isDynamic(shape) &&
-        shape != staticOutputShapes[staticShapeNum++])
+    if (!ShapedType::isDynamic(shape) && shape != staticOutputShapes[pos])
       emitOpError("invalid output shape provided at pos ") << pos;
 
   return success();

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That makes sense... does this fix the issue you were hitting?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes it does.

return success();
}

Expand Down
11 changes: 11 additions & 0 deletions mlir/test/Dialect/MemRef/invalid.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1103,3 +1103,14 @@ func.func @subview_invalid_strides_rank_reduction(%m: memref<7x22x333x4444xi32>)
: memref<7x22x333x4444xi32> to memref<7x11x4444xi32>
return
}

// -----

func.func @expand_shape_invalid_output_shape(
%arg0: memref<30x20xf32, strided<[4000, 2], offset: 100>>) {
// expected-error @+1 {{invalid output shape provided at pos 2}}
%0 = memref.expand_shape %arg0 [[0, 1], [2]] output_shape [2, 15, 21] :
memref<30x20xf32, strided<[4000, 2], offset: 100>>
into memref<2x15x20xf32, strided<[60000, 4000, 2], offset: 100>>
return
}