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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-gpu

Author: Prashant Kumar (pashu123)

<details>
<summary>Changes</summary>

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

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


1 Files Affected:

- (modified) mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp (+3-1) 


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

``````````

</details>


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