[Mlir-commits] [mlir] 27534d6 - [mlir][LLVM] erase call mappings in forgetMapping() (#84955)

[llvmlistbot at llvm.org]

Author: Tom Eccles
Date: 2024-03-20T10:03:26Z
New Revision: 27534d69e275eeb6669296c038bf29d8a9b163e5

URL: https://github.com/llvm/llvm-project/commit/27534d69e275eeb6669296c038bf29d8a9b163e5
DIFF: https://github.com/llvm/llvm-project/commit/27534d69e275eeb6669296c038bf29d8a9b163e5.diff

LOG: [mlir][LLVM] erase call mappings in forgetMapping() (#84955)

It looks like the mappings for call instructions were forgotten here.
This fixes a bug in OpenMP when in-lining a region containing call
operations multiple times.

OpenMP array reductions 4/6
Previous PR: https://github.com/llvm/llvm-project/pull/84954
Next PR: https://github.com/llvm/llvm-project/pull/84957

Added: 
    mlir/test/Target/LLVMIR/openmp-reduction-call.mlir

Modified: 
    mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index c00628a420a000..995544238e4a3c 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -716,6 +716,8 @@ void ModuleTranslation::forgetMapping(Region &region) {
           branchMapping.erase(&op);
         if (isa<LLVM::GlobalOp>(op))
           globalsMapping.erase(&op);
+        if (isa<LLVM::CallOp>(op))
+          callMapping.erase(&op);
         llvm::append_range(
             toProcess,
             llvm::map_range(op.getRegions(), [](Region &r) { return &r; }));

diff  --git a/mlir/test/Target/LLVMIR/openmp-reduction-call.mlir b/mlir/test/Target/LLVMIR/openmp-reduction-call.mlir
new file mode 100644
index 00000000000000..9aa9ac7e4b5a52
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-reduction-call.mlir
@@ -0,0 +1,37 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// Test that we don't crash when there is a call operation in the combiner
+
+omp.reduction.declare @add_f32 : f32
+init {
+^bb0(%arg: f32):
+  %0 = llvm.mlir.constant(0.0 : f32) : f32
+  omp.yield (%0 : f32)
+}
+combiner {
+^bb1(%arg0: f32, %arg1: f32):
+// test this call here:
+  llvm.call @test_call() : () -> ()
+  %1 = llvm.fadd %arg0, %arg1 : f32
+  omp.yield (%1 : f32)
+}
+
+llvm.func @simple_reduction(%lb : i64, %ub : i64, %step : i64) {
+  %c1 = llvm.mlir.constant(1 : i32) : i32
+  %0 = llvm.alloca %c1 x i32 : (i32) -> !llvm.ptr
+  omp.parallel reduction(@add_f32 %0 -> %prv : !llvm.ptr) {
+    %1 = llvm.mlir.constant(2.0 : f32) : f32
+    %2 = llvm.load %prv : !llvm.ptr -> f32
+    %3 = llvm.fadd %1, %2 : f32
+    llvm.store %3, %prv : f32, !llvm.ptr
+    omp.terminator
+  }
+  llvm.return
+}
+
+llvm.func @test_call() -> ()
+
+// Call to the troublesome function will have been inlined twice: once into
+// main and once into the outlined function
+// CHECK: call void @test_call()
+// CHECK: call void @test_call()