[Mlir-commits] [mlir] 9bf7023 - [mlir][linalg] fix indexOp folder to work in genericOp build with createOrFold (#137427)

[llvmlistbot at llvm.org]

Author: Stanley Winata
Date: 2025-04-25T18:52:02-07:00
New Revision: 9bf702337d86727b48adcc00a60159b1650591dc

URL: https://github.com/llvm/llvm-project/commit/9bf702337d86727b48adcc00a60159b1650591dc
DIFF: https://github.com/llvm/llvm-project/commit/9bf702337d86727b48adcc00a60159b1650591dc.diff

LOG: [mlir][linalg] fix indexOp folder to work in genericOp build with createOrFold (#137427)

Currently in torch-mlir, indexOp folder is segfaulting when we call
createOrFold in a genericOp builder. (*this)->getParentOp ended up with
null which is causing the issue.

This is seen in poolSizeCalculator.getPoolSize
being called from createAvgPoolValueCountIncludePadFalseCase. link:

https://github.com/llvm/torch-mlir/blob/80a3dfddd341c72ab9bd6c6688b872bf3a5e4ddb/lib/Conversion/TorchToLinalg/Pooling.cpp#L918-L921

---------

Signed-off-by: Stanley Winata <stanley.winata at amd.com>
Co-authored-by: Jakub Kuderski <kubakuderski at gmail.com>

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 72fb3308a2549..089ccc6680e48 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -2284,7 +2284,12 @@ LogicalResult IndexOp::verify() {
 }
 
 OpFoldResult IndexOp::fold(FoldAdaptor adaptor) {
-  auto linalgOp = cast<LinalgOp>((*this)->getParentOp());
+  auto linalgOp = dyn_cast_or_null<LinalgOp>((*this)->getParentOp());
+  // Bail out if `linalg.index` does not have a proper parent yet at this
+  // point, e.g., when calling `createOrFold` during IR construction in
+  // `genericOp::build`.
+  if (!linalgOp)
+    return OpFoldResult{};
 
   // Index of unit dims is always 0.
   SmallVector<int64_t, 4> loopBounds = linalgOp.getStaticLoopRanges();