[mlir][linalg] fix indexOp folder to work in genericOp build with createOrFold

[raikonenfnu]

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

[llvmbot]

@llvm/pr-subscribers-mlir-linalg

Author: Stanley Winata (raikonenfnu)

Changes

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

---

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

1 Files Affected:

• (modified) mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp (+4-1)

diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 72fb3308a2549..bd2430927fca2 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -2284,7 +2284,10 @@ LogicalResult IndexOp::verify() {
 }
 
 OpFoldResult IndexOp::fold(FoldAdaptor adaptor) {
-  auto linalgOp = cast<LinalgOp>((*this)->getParentOp());
+  auto linalgOp = dyn_cast_or_null<LinalgOp>((*this)->getParentOp());
+  // Early exit if parent op is not linalgOp.
+  if (!linalgOp)
+    return OpFoldResult{};
 
   // Index of unit dims is always 0.
   SmallVector<int64_t, 4> loopBounds = linalgOp.getStaticLoopRanges();

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Stanley Winata (raikonenfnu)

Changes

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

---

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

1 Files Affected:

• (modified) mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp (+4-1)

diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 72fb3308a2549..bd2430927fca2 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -2284,7 +2284,10 @@ LogicalResult IndexOp::verify() {
 }
 
 OpFoldResult IndexOp::fold(FoldAdaptor adaptor) {
-  auto linalgOp = cast<LinalgOp>((*this)->getParentOp());
+  auto linalgOp = dyn_cast_or_null<LinalgOp>((*this)->getParentOp());
+  // Early exit if parent op is not linalgOp.
+  if (!linalgOp)
+    return OpFoldResult{};
 
   // Index of unit dims is always 0.
   SmallVector<int64_t, 4> loopBounds = linalgOp.getStaticLoopRanges();

[kuhar]

Thanks for fixing this! Can we add a test?

[raikonenfnu]

Thanks for fixing this! Can we add a test?

Was thinking about this but can't think of a good test for it, do you have some suggestions? :)

[kuhar]

Thanks for fixing this! Can we add a test?

Was thinking about this but can't think of a good test for it, do you have some suggestions? :)

Can we take the IR from troch-mlir that caused the original crash?

[raikonenfnu]

Thanks for fixing this! Can we add a test?

Was thinking about this but can't think of a good test for it, do you have some suggestions? :)

Can we take the IR from troch-mlir that caused the original crash?

Fantastic question! I have a repro IR here https://gist.github.com/raikonenfnu/5f9df8748712fada5b167abd696a13fa. However, it is failing not from an IR level, but rather it fails during a lowering, because it calls into createOrFold here https://github.com/llvm/torch-mlir/blob/80a3dfddd341c72ab9bd6c6688b872bf3a5e4ddb/lib/Conversion/TorchToLinalg/Pooling.cpp#L918-L921

so not so sure if I can recreate that in IR 😅

[kuhar]

Oh, I understand, so it's more that linalg.index doesn't have a proper parent yet at that point. In this case, could you clarify this in the code comment above the null check?

[raikonenfnu]

Oh, I understand, so it's more that linalg.index doesn't have a proper parent yet at that point. In this case, could you clarify this in the code comment above the null check?

Exactly! :)

[raikonenfnu]

Oh, I understand, so it's more that linalg.index doesn't have a proper parent yet at that point. In this case, could you clarify this in the code comment above the null check?

done!

[kuhar]

LGTM % one suggestion

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.