[Mlir-commits] [mlir] [MLIR][OpenMP] - Fix translation of omp.target when private variables need cleaning up (PR #129205)

[Pranav Bhandarkar]

https://github.com/bhandarkar-pranav updated https://github.com/llvm/llvm-project/pull/129205

>From f7ec7560d72993c1bba7d1daac6bc8009b9056b3 Mon Sep 17 00:00:00 2001
From: Pranav Bhandarkar <pranav.bhandarkar at amd.com>
Date: Fri, 28 Feb 2025 01:10:53 -0600
Subject: [PATCH 1/2] [MLIR][OpenMP] - Fix translation of omp.target when
 private variables need cleaning up

This is a simple fix that ensures that the InsertPoint is properly fixed up
after we have translated the dealloc region of all privatized variables during
translation of omp.target from MLIR to LLVMIR.

Fix for https://github.com/llvm/llvm-project/issues/129202
---
 .../Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp   | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 5035551dd6023..32c7c501d03c3 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -4832,6 +4832,7 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
             "failed to inline `dealloc` region of `omp.private` "
             "op in the target region");
       }
+      return builder.saveIP();
     }
 
     return InsertPointTy(exitBlock.get(), exitBlock.get()->end());

>From 8c8b89c1e4f805bc4c3088d75502b18d9fd8cd40 Mon Sep 17 00:00:00 2001
From: Pranav Bhandarkar <pranav.bhandarkar at amd.com>
Date: Fri, 28 Feb 2025 11:55:39 -0600
Subject: [PATCH 2/2] Add a test

In fact, modify an existing test to check multi-block dealloc regions
---
 .../LLVMIR/openmp-target-private-allocatable.mlir     | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/mlir/test/Target/LLVMIR/openmp-target-private-allocatable.mlir b/mlir/test/Target/LLVMIR/openmp-target-private-allocatable.mlir
index a25f920287fe3..0ee9230b5af0e 100644
--- a/mlir/test/Target/LLVMIR/openmp-target-private-allocatable.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-target-private-allocatable.mlir
@@ -8,7 +8,15 @@ omp.private {type = private} @box.heap_privatizer : !llvm.struct<(ptr, i64, i32,
   llvm.call @alloc_foo_1(%arg0) : (!llvm.ptr) -> ()
   omp.yield(%arg1 : !llvm.ptr)
 } dealloc {
+// There is no reason for the dealloc region here to have two blocks.
+// But a multi-block test is useful in checking that the OMPIRBuilder
+// has updated the InsertPoint properly after translating the dealloc
+// region to LLVMIR.
+// See https://github.com/llvm/llvm-project/issues/129202 for more
+// context
 ^bb0(%arg0: !llvm.ptr):
+  llvm.br ^bb1
+^bb1:
   llvm.call @dealloc_foo_1(%arg0) : (!llvm.ptr) -> ()
   omp.yield
 }
@@ -60,3 +68,6 @@ llvm.func @_FortranAAssign(!llvm.ptr, !llvm.ptr, !llvm.ptr, i32) -> !llvm.struct
 
 // Now, check the deallocation of the private var.
 // CHECK:  call void @dealloc_foo_1(ptr %[[DESC_TO_DEALLOC]])
+// CHECK-NEXT: br label %[[CONT_BLOCK:.*]]
+// CHECK: [[CONT_BLOCK]]:
+// CHECK-NEXT: ret void