[Mlir-commits] [mlir] [GPU] ASAN failure due to use of erased operation (PR #116816)

[Prashant Kumar]

https://github.com/pashu123 created https://github.com/llvm/llvm-project/pull/116816

If the target op is an scf.forall op which is erased by `rewriteOneForallCommonImpl` method it leads to asan failure since it's used later by `replaceUnitMappingIdsHelper` fn.

Follows the same philosophy as done here:
https://github.com/llvm/llvm-project/blob/aff98e4be05a1060e489ce62a88ee0ff365e571a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp#L599 and later passed to the same method here:
https://github.com/llvm/llvm-project/blob/aff98e4be05a1060e489ce62a88ee0ff365e571a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp#L629

>From 740af37300ec20d92be878af2ef0f0013c5a3636 Mon Sep 17 00:00:00 2001
From: Prashant Kumar <pk5561 at gmail.com>
Date: Tue, 19 Nov 2024 20:17:47 +0530
Subject: [PATCH] [GPU] ASAN failure due to use of erased operation

If the target op is an scf.forall op which is erased by
`rewriteOneForallCommonImpl` method it leads to asan failure since it's
used later by `replaceUnitMappingIdsHelper` fn.

Follows the same philosophy as done here:
https://github.com/llvm/llvm-project/blob/aff98e4be05a1060e489ce62a88ee0ff365e571a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp#L599
and later passed to the same method here:
https://github.com/llvm/llvm-project/blob/aff98e4be05a1060e489ce62a88ee0ff365e571a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp#L629
---
 mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
index 1528da914d546b..e836820c960339 100644
--- a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
+++ b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
@@ -853,6 +853,8 @@ DiagnosedSilenceableFailure mlir::transform::gpu::mapNestedForallToThreadsImpl(
                                  "requires size-3 thread mapping");
   }
 
+  Block *parentBlock = target->getBlock();
+
   // Create an early zero index value for replacements.
   Location loc = target->getLoc();
   Value zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
@@ -872,7 +874,7 @@ DiagnosedSilenceableFailure mlir::transform::gpu::mapNestedForallToThreadsImpl(
 
   // Replace ids of dimensions known to be 1 by 0 to simplify the IR.
   // Here, the result of mapping determines the available mapping sizes.
-  replaceUnitMappingIdsHelper<ThreadIdOp>(rewriter, loc, target, zero,
+  replaceUnitMappingIdsHelper<ThreadIdOp>(rewriter, loc, parentBlock, zero,
                                           blockDims);
 
   return DiagnosedSilenceableFailure::success();