From 740af37300ec20d92be878af2ef0f0013c5a3636 Mon Sep 17 00:00:00 2001 From: Prashant Kumar 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 1528da914d546..e836820c96033 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(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(rewriter, loc, target, zero, + replaceUnitMappingIdsHelper(rewriter, loc, parentBlock, zero, blockDims); return DiagnosedSilenceableFailure::success();