[llvm] b437014 - [OpenMPIRBuilder] Do not call host runtime for GPU teams codegen (#79984)

[via llvm-commits]

Author: Dominik Adamski
Date: 2024-01-31T12:16:35+01:00
New Revision: b4370140b4467eddc42e7b8075959c692daaf3e8

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

LOG: [OpenMPIRBuilder] Do not call host runtime for GPU teams codegen (#79984)

Patch ensures that host runtime functions are not called for handling
OpenMP teams clause on the device.

GPU code for pragma `omp target teams distribute parallel do` will
require only one call to OpenMP loop-worksharing GPU runtime. Support
for it will be added later.

This patch does not include changes required for handling `omp target
teams` for the host side.

Added: 
    mlir/test/Target/LLVMIR/omptarget-teams-llvm.mlir

Modified: 
    llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 211281452de22..8eb8a13686dd3 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -6251,8 +6251,10 @@ OpenMPIRBuilder::createTeams(const LocationDescription &Loc,
   BasicBlock *AllocaBB =
       splitBB(Builder, /*CreateBranch=*/true, "teams.alloca");
 
+  bool SubClausesPresent =
+      (NumTeamsLower || NumTeamsUpper || ThreadLimit || IfExpr);
   // Push num_teams
-  if (NumTeamsLower || NumTeamsUpper || ThreadLimit || IfExpr) {
+  if (!Config.isTargetDevice() && SubClausesPresent) {
     assert((NumTeamsLower == nullptr || NumTeamsUpper != nullptr) &&
            "if lowerbound is non-null, then upperbound must also be non-null "
            "for bounds on num_teams");
@@ -6305,7 +6307,8 @@ OpenMPIRBuilder::createTeams(const LocationDescription &Loc,
   OI.ExcludeArgsFromAggregate.push_back(createFakeIntVal(
       Builder, OuterAllocaIP, ToBeDeleted, AllocaIP, "tid", true));
 
-  OI.PostOutlineCB = [this, Ident, ToBeDeleted](Function &OutlinedFn) mutable {
+  auto HostPostOutlineCB = [this, Ident,
+                            ToBeDeleted](Function &OutlinedFn) mutable {
     // The stale call instruction will be replaced with a new call instruction
     // for runtime call with the outlined function.
 
@@ -6342,6 +6345,9 @@ OpenMPIRBuilder::createTeams(const LocationDescription &Loc,
     }
   };
 
+  if (!Config.isTargetDevice())
+    OI.PostOutlineCB = HostPostOutlineCB;
+
   addOutlineInfo(std::move(OI));
 
   Builder.SetInsertPoint(ExitBB, ExitBB->begin());

diff  --git a/mlir/test/Target/LLVMIR/omptarget-teams-llvm.mlir b/mlir/test/Target/LLVMIR/omptarget-teams-llvm.mlir
new file mode 100644
index 0000000000000..96cced7a1d584
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/omptarget-teams-llvm.mlir
@@ -0,0 +1,25 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// The aim of the test is to check the LLVM IR codegen for the device
+// for omp teams construct
+
+module attributes {omp.is_target_device = true} {
+  llvm.func @foo(i32)
+  llvm.func @omp_target_teams_shared_simple(%arg0 : i32)  {
+    omp.teams {
+      llvm.call @foo(%arg0) : (i32) -> ()
+      omp.terminator
+    }
+  llvm.return
+  }
+}
+
+// CHECK-LABEL: @omp_target_teams_shared_simple
+// CHECK-SAME: (i32 [[ARG0:%.+]])
+// CHECK: call void @[[OUTLINED_FN:.*]](
+// CHECK-NOT: call {{.+}} @__kmpc_fork_teams
+// CHECK: ret void
+
+//CHECK: define internal void @[[OUTLINED_FN]](
+//CHECK: call void @foo(i32 %[[FOO_ARG:.*]])
+//CHECK: ret void