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

[Stanley Winata]

https://github.com/raikonenfnu updated https://github.com/llvm/llvm-project/pull/137427

>From ba3bf3ab088b369fc760d2518e37a2a1de4be6d5 Mon Sep 17 00:00:00 2001
From: Stanley Winata <stanley.winata at amd.com>
Date: Fri, 25 Apr 2025 17:37:12 -0700
Subject: [PATCH 1/3] [mlir][linalg] fix indexOp folder to work in genericOp
 build with createOrFold

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>
---
 mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

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();

>From b5cb807f8cb1662531ffd01a1854b167b0f4d169 Mon Sep 17 00:00:00 2001
From: Stanley Winata <stanley.winata at amd.com>
Date: Fri, 25 Apr 2025 18:23:15 -0700
Subject: [PATCH 2/3] clear up comment on when early exit is useful

Signed-off-by: Stanley Winata <stanley.winata at amd.com>
---
 mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index bd2430927fca2..06bf2762b855d 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -2285,7 +2285,8 @@ LogicalResult IndexOp::verify() {
 
 OpFoldResult IndexOp::fold(FoldAdaptor adaptor) {
   auto linalgOp = dyn_cast_or_null<LinalgOp>((*this)->getParentOp());
-  // Early exit if parent op is not linalgOp.
+  // Early exit if linalg.index does not have a proper parent yet
+  // at this point. (e.g createOrFold in a genericOp::build)
   if (!linalgOp)
     return OpFoldResult{};
 

>From 81676b502ab6da19494c7a00dda45ac5876e0b9f Mon Sep 17 00:00:00 2001
From: Stanley Winata <68087699+raikonenfnu at users.noreply.github.com>
Date: Fri, 25 Apr 2025 18:34:40 -0700
Subject: [PATCH 3/3] Update mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

Co-authored-by: Jakub Kuderski <kubakuderski at gmail.com>
---
 mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 06bf2762b855d..ddc0348a4b115 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -2285,8 +2285,8 @@ LogicalResult IndexOp::verify() {
 
 OpFoldResult IndexOp::fold(FoldAdaptor adaptor) {
   auto linalgOp = dyn_cast_or_null<LinalgOp>((*this)->getParentOp());
-  // Early exit if linalg.index does not have a proper parent yet
-  // at this point. (e.g createOrFold in a genericOp::build)
+  // 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{};