[flang-commits] [flang] [OpenMP][mlir] Added `num_teams`, `thread_limit` translation to LLVM IR (PR #68821)
via flang-commits
flang-commits at lists.llvm.org
Sun Oct 15 11:17:29 PDT 2023
https://github.com/shraiysh updated https://github.com/llvm/llvm-project/pull/68821
>From fff04e51e5d799fd9f39e0d58fc3d1b854325981 Mon Sep 17 00:00:00 2001
From: Shraiysh Vaishay <shraiysh.vaishay at amd.com>
Date: Wed, 11 Oct 2023 11:51:19 -0500
Subject: [PATCH 1/2] [OpenMP][mlir] Added `num_teams`, `thread_limit`
translation to LLVM IR
This patch adds translation to LLVM IR for `num_teams` and
`thread_limit` in for `omp.teams` operation.
---
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 22 +++-
mlir/test/Target/LLVMIR/openmp-teams.mlir | 111 ++++++++++++++++++
2 files changed, 127 insertions(+), 6 deletions(-)
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 1ec3bb8e7562a9e..ae974c14fac41a6 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -667,11 +667,9 @@ convertOmpTeams(omp::TeamsOp op, llvm::IRBuilderBase &builder,
LLVM::ModuleTranslation &moduleTranslation) {
using InsertPointTy = llvm::OpenMPIRBuilder::InsertPointTy;
LogicalResult bodyGenStatus = success();
- if (op.getNumTeamsLower() || op.getNumTeamsUpper() || op.getIfExpr() ||
- op.getThreadLimit() || !op.getAllocatorsVars().empty() ||
- op.getReductions()) {
+ if (op.getIfExpr() || !op.getAllocatorsVars().empty() || op.getReductions())
return op.emitError("unhandled clauses for translation to LLVM IR");
- }
+
auto bodyCB = [&](InsertPointTy allocaIP, InsertPointTy codegenIP) {
LLVM::ModuleTranslation::SaveStack<OpenMPAllocaStackFrame> frame(
moduleTranslation, allocaIP);
@@ -680,9 +678,21 @@ convertOmpTeams(omp::TeamsOp op, llvm::IRBuilderBase &builder,
moduleTranslation, bodyGenStatus);
};
+ llvm::Value *numTeamsLower = nullptr;
+ if (auto numTeamsLowerVar = op.getNumTeamsLower())
+ numTeamsLower = moduleTranslation.lookupValue(numTeamsLowerVar);
+
+ llvm::Value *numTeamsUpper = nullptr;
+ if (auto numTeamsUpperVar = op.getNumTeamsUpper())
+ numTeamsUpper = moduleTranslation.lookupValue(numTeamsUpperVar);
+
+ llvm::Value *threadLimit = nullptr;
+ if (auto threadLimitVar = op.getThreadLimit())
+ threadLimit = moduleTranslation.lookupValue(threadLimitVar);
+
llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
- builder.restoreIP(
- moduleTranslation.getOpenMPBuilder()->createTeams(ompLoc, bodyCB));
+ builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createTeams(
+ ompLoc, bodyCB, numTeamsLower, numTeamsUpper, threadLimit));
return bodyGenStatus;
}
diff --git a/mlir/test/Target/LLVMIR/openmp-teams.mlir b/mlir/test/Target/LLVMIR/openmp-teams.mlir
index 18fc2bb5a3c61b2..87ef90223ed704a 100644
--- a/mlir/test/Target/LLVMIR/openmp-teams.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-teams.mlir
@@ -124,3 +124,114 @@ llvm.func @omp_teams_branching_shared(%condition: i1, %arg0: i32, %arg1: f32, %a
// CHECK-NEXT: br label
// CHECK: ret void
+// -----
+
+llvm.func @beforeTeams()
+llvm.func @duringTeams()
+llvm.func @afterTeams()
+
+// CHECK-LABEL: @omp_teams_thread_limit
+// CHECK-SAME: (i32 [[THREAD_LIMIT:.+]])
+llvm.func @omp_teams_thread_limit(%threadLimit: i32) {
+ // CHECK-NEXT: call void @beforeTeams()
+ llvm.call @beforeTeams() : () -> ()
+ // CHECK: [[THREAD_NUM:%.+]] = call i32 @__kmpc_global_thread_num
+ // CHECK-NEXT: call void @__kmpc_push_num_teams_51({{.+}}, i32 [[THREAD_NUM]], i32 0, i32 0, i32 [[THREAD_LIMIT]])
+ // CHECK: call void (ptr, i32, ptr, ...) @__kmpc_fork_teams(ptr @1, i32 0, ptr [[OUTLINED_FN:.+]])
+ omp.teams thread_limit(%threadLimit : i32) {
+ llvm.call @duringTeams() : () -> ()
+ omp.terminator
+ }
+ // CHECK: call void @afterTeams
+ llvm.call @afterTeams() : () -> ()
+ // CHECK: ret void
+ llvm.return
+}
+
+// CHECK: define internal void [[OUTLINED_FN]](ptr {{.+}}, ptr {{.+}})
+// CHECK: call void @duringTeams()
+// CHECK: ret void
+
+// -----
+
+llvm.func @beforeTeams()
+llvm.func @duringTeams()
+llvm.func @afterTeams()
+
+// CHECK-LABEL: @omp_teams_num_teams_upper
+// CHECK-SAME: (i32 [[NUM_TEAMS_UPPER:.+]])
+llvm.func @omp_teams_num_teams_upper(%numTeamsUpper: i32) {
+ // CHECK-NEXT: call void @beforeTeams()
+ llvm.call @beforeTeams() : () -> ()
+ // CHECK: [[THREAD_NUM:%.+]] = call i32 @__kmpc_global_thread_num
+ // CHECK-NEXT: call void @__kmpc_push_num_teams_51({{.+}}, i32 [[THREAD_NUM]], i32 [[NUM_TEAMS_UPPER]], i32 [[NUM_TEAMS_UPPER]], i32 0)
+ // CHECK: call void (ptr, i32, ptr, ...) @__kmpc_fork_teams(ptr @1, i32 0, ptr [[OUTLINED_FN:.+]])
+ omp.teams num_teams(to %numTeamsUpper : i32) {
+ llvm.call @duringTeams() : () -> ()
+ omp.terminator
+ }
+ // CHECK: call void @afterTeams
+ llvm.call @afterTeams() : () -> ()
+ // CHECK: ret void
+ llvm.return
+}
+
+// CHECK: define internal void [[OUTLINED_FN]](ptr {{.+}}, ptr {{.+}})
+// CHECK: call void @duringTeams()
+// CHECK: ret void
+
+// -----
+
+llvm.func @beforeTeams()
+llvm.func @duringTeams()
+llvm.func @afterTeams()
+
+// CHECK-LABEL: @omp_teams_num_teams_lower_and_upper
+// CHECK-SAME: (i32 [[NUM_TEAMS_LOWER:.+]], i32 [[NUM_TEAMS_UPPER:.+]])
+llvm.func @omp_teams_num_teams_lower_and_upper(%numTeamsLower: i32, %numTeamsUpper: i32) {
+ // CHECK-NEXT: call void @beforeTeams()
+ llvm.call @beforeTeams() : () -> ()
+ // CHECK: [[THREAD_NUM:%.+]] = call i32 @__kmpc_global_thread_num
+ // CHECK-NEXT: call void @__kmpc_push_num_teams_51({{.+}}, i32 [[THREAD_NUM]], i32 [[NUM_TEAMS_LOWER]], i32 [[NUM_TEAMS_UPPER]], i32 0)
+ // CHECK: call void (ptr, i32, ptr, ...) @__kmpc_fork_teams(ptr @1, i32 0, ptr [[OUTLINED_FN:.+]])
+ omp.teams num_teams(%numTeamsLower : i32 to %numTeamsUpper: i32) {
+ llvm.call @duringTeams() : () -> ()
+ omp.terminator
+ }
+ // CHECK: call void @afterTeams
+ llvm.call @afterTeams() : () -> ()
+ // CHECK: ret void
+ llvm.return
+}
+
+// CHECK: define internal void [[OUTLINED_FN]](ptr {{.+}}, ptr {{.+}})
+// CHECK: call void @duringTeams()
+// CHECK: ret void
+
+// -----
+
+llvm.func @beforeTeams()
+llvm.func @duringTeams()
+llvm.func @afterTeams()
+
+// CHECK-LABEL: @omp_teams_num_teams_and_thread_limit
+// CHECK-SAME: (i32 [[NUM_TEAMS_LOWER:.+]], i32 [[NUM_TEAMS_UPPER:.+]], i32 [[THREAD_LIMIT:.+]])
+llvm.func @omp_teams_num_teams_and_thread_limit(%numTeamsLower: i32, %numTeamsUpper: i32, %threadLimit: i32) {
+ // CHECK-NEXT: call void @beforeTeams()
+ llvm.call @beforeTeams() : () -> ()
+ // CHECK: [[THREAD_NUM:%.+]] = call i32 @__kmpc_global_thread_num
+ // CHECK-NEXT: call void @__kmpc_push_num_teams_51({{.+}}, i32 [[THREAD_NUM]], i32 [[NUM_TEAMS_LOWER]], i32 [[NUM_TEAMS_UPPER]], i32 [[THREAD_LIMIT]])
+ // CHECK: call void (ptr, i32, ptr, ...) @__kmpc_fork_teams(ptr @1, i32 0, ptr [[OUTLINED_FN:.+]])
+ omp.teams num_teams(%numTeamsLower : i32 to %numTeamsUpper: i32) thread_limit(%threadLimit: i32) {
+ llvm.call @duringTeams() : () -> ()
+ omp.terminator
+ }
+ // CHECK: call void @afterTeams
+ llvm.call @afterTeams() : () -> ()
+ // CHECK: ret void
+ llvm.return
+}
+
+// CHECK: define internal void [[OUTLINED_FN]](ptr {{.+}}, ptr {{.+}})
+// CHECK: call void @duringTeams()
+// CHECK: ret void
>From 6811ac6058aca8d81259ab6c905d0b3e91a4c3e3 Mon Sep 17 00:00:00 2001
From: Shraiysh Vaishay <shraiysh.vaishay at amd.com>
Date: Sun, 15 Oct 2023 13:16:57 -0500
Subject: [PATCH 2/2] Address comments
---
.../LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 076a5dd4638d756..e3dc68a1b8b7d72 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -678,15 +678,15 @@ convertOmpTeams(omp::TeamsOp op, llvm::IRBuilderBase &builder,
};
llvm::Value *numTeamsLower = nullptr;
- if (auto numTeamsLowerVar = op.getNumTeamsLower())
+ if (Value numTeamsLowerVar = op.getNumTeamsLower())
numTeamsLower = moduleTranslation.lookupValue(numTeamsLowerVar);
llvm::Value *numTeamsUpper = nullptr;
- if (auto numTeamsUpperVar = op.getNumTeamsUpper())
+ if (Value numTeamsUpperVar = op.getNumTeamsUpper())
numTeamsUpper = moduleTranslation.lookupValue(numTeamsUpperVar);
llvm::Value *threadLimit = nullptr;
- if (auto threadLimitVar = op.getThreadLimit())
+ if (Value threadLimitVar = op.getThreadLimit())
threadLimit = moduleTranslation.lookupValue(threadLimitVar);
llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
More information about the flang-commits
mailing list