From ba3bf3ab088b369fc760d2518e37a2a1de4be6d5 Mon Sep 17 00:00:00 2001 From: Stanley Winata Date: Fri, 25 Apr 2025 17:37:12 -0700 Subject: [PATCH 1/4] [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 --- 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((*this)->getParentOp()); + auto linalgOp = dyn_cast_or_null((*this)->getParentOp()); + // Early exit if parent op is not linalgOp. + if (!linalgOp) + return OpFoldResult{}; // Index of unit dims is always 0. SmallVector loopBounds = linalgOp.getStaticLoopRanges(); From b5cb807f8cb1662531ffd01a1854b167b0f4d169 Mon Sep 17 00:00:00 2001 From: Stanley Winata Date: Fri, 25 Apr 2025 18:23:15 -0700 Subject: [PATCH 2/4] clear up comment on when early exit is useful Signed-off-by: Stanley Winata --- 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((*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@users.noreply.github.com> Date: Fri, 25 Apr 2025 18:34:40 -0700 Subject: [PATCH 3/4] Update mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp Co-authored-by: Jakub Kuderski --- 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((*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{}; From 4bf836cba9950b131eaabf5f225e123c84f25a92 Mon Sep 17 00:00:00 2001 From: Stanley Winata Date: Fri, 25 Apr 2025 18:38:11 -0700 Subject: [PATCH 4/4] clang formatting Signed-off-by: Stanley Winata --- 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 ddc0348a4b115..089ccc6680e48 100644 --- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp +++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp @@ -2286,7 +2286,8 @@ LogicalResult IndexOp::verify() { OpFoldResult IndexOp::fold(FoldAdaptor adaptor) { auto linalgOp = dyn_cast_or_null((*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`. + // point, e.g., when calling `createOrFold` during IR construction in + // `genericOp::build`. if (!linalgOp) return OpFoldResult{};