Fix memref.expand_shape verifier

[bjacob]

Torch-mlir integration is currently blocked on memref.expand_shape verifier errors of the form

'memref.expand_shape' op invalid output shape provided at pos 1

The verifier code generating these errors was introduced in #91245. I have commented there why I believe it's incorrect. This PR has my suggested fix.

Unfortunately, this does not seem to be directly testable on memref IR, because static_output_shape is not directly exposed in the custom assembly format.

[llvmbot]

@llvm/pr-subscribers-mlir-memref

@llvm/pr-subscribers-mlir

Author: Benoit Jacob (bjacob)

Changes

Torch-mlir integration is currently blocked on memref.expand_shape verifier errors of the form

'memref.expand_shape' op invalid output shape provided at pos 1

The verifier code generating these errors was introduced in #91245. I have commented there why I believe it's incorrect. This PR has my suggested fix.

Unfortunately, this does not seem to be directly testable on memref IR, because static_output_shape is not directly exposed in the custom assembly format.

---

Full diff: https://github.com/llvm/llvm-project/pull/91501.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp (+4-5)

diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index 78201ae29cd9..dc6fd770d9bd 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -2356,12 +2356,11 @@ 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++])
+  for (auto [pos, shape] : llvm::enumerate(resShape)) {
+    if (!ShapedType::isDynamic(shape) && shape != staticOutputShapes[pos]) {
       emitOpError("invalid output shape provided at pos ") << pos;
+    }
+  }
 
   return success();
 }