[llvm-branch-commits] [mlir] [mlir][LLVM] erase call mappings in forgetMapping() (PR #84955)

[Tom Eccles via llvm-branch-commits]

https://github.com/tblah updated https://github.com/llvm/llvm-project/pull/84955

>From c62b31262bc619145866a304e10925a35462f5bf Mon Sep 17 00:00:00 2001
From: Tom Eccles <tom.eccles at arm.com>
Date: Wed, 21 Feb 2024 14:22:39 +0000
Subject: [PATCH 1/2] [mlir][LLVM] erase call mappings in forgetMapping()

It looks like the mappings for call instructions were forgotten here.
This fixes a bug in OpenMP when inlining a region containing call
operations multiple times.
---
 mlir/lib/Target/LLVMIR/ModuleTranslation.cpp | 2 ++
 1 file changed, 2 insertions(+)

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; }));

>From 83b87ec9057dc9fbe529c652b5020452fb9585bb Mon Sep 17 00:00:00 2001
From: Tom Eccles <tom.eccles at arm.com>
Date: Fri, 15 Mar 2024 14:05:53 +0000
Subject: [PATCH 2/2] Add test

---
 .../Target/LLVMIR/openmp-reduction-call.mlir  | 47 +++++++++++++++++++
 1 file changed, 47 insertions(+)
 create mode 100644 mlir/test/Target/LLVMIR/openmp-reduction-call.mlir

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..60419a85f66aa2
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-reduction-call.mlir
@@ -0,0 +1,47 @@
+// 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)
+}
+atomic {
+^bb2(%arg2: !llvm.ptr, %arg3: !llvm.ptr):
+  %2 = llvm.load %arg3 : !llvm.ptr -> f32
+  llvm.atomicrmw fadd %arg2, %2 monotonic : !llvm.ptr, f32
+  omp.yield
+}
+
+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 {
+    omp.wsloop reduction(@add_f32 %0 -> %prv : !llvm.ptr)
+    for (%iv) : i64 = (%lb) to (%ub) step (%step) {
+      %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.yield
+    }
+    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()