[Openmp-commits] [openmp] [mlir][OpenMP] Added translation for `omp.teams` to LLVM IR (PR #68042)

via Openmp-commits openmp-commits at lists.llvm.org
Tue Oct 3 11:06:38 PDT 2023


https://github.com/shraiysh updated https://github.com/llvm/llvm-project/pull/68042

>From c7c9e907d897ae667331761d8097ccb7852c5d93 Mon Sep 17 00:00:00 2001
From: Shraiysh Vaishay <shraiysh.vaishay at amd.com>
Date: Mon, 2 Oct 2023 16:43:13 -0500
Subject: [PATCH 1/2] [mlir][OpenMP] Added translation for `omp.teams` to LLVM
 IR

This patch adds translation from `omp.teams` operation to LLVM IR using
OpenMPIRBuilder.

The clauses are not handled in this patch.
---
 .../OpenMP/OpenMPToLLVMIRTranslation.cpp      |  21 +++
 mlir/test/Target/LLVMIR/openmp-teams.mlir     | 136 ++++++++++++++++++
 2 files changed, 157 insertions(+)
 create mode 100644 mlir/test/Target/LLVMIR/openmp-teams.mlir

diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 8f7f1963b3e5a4f..b9643be40e13c01 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -661,6 +661,24 @@ convertOmpSingle(omp::SingleOp &singleOp, llvm::IRBuilderBase &builder,
   return bodyGenStatus;
 }
 
+// Convert an OpenMP Teams construct to LLVM IR using OpenMPIRBuilder
+static LogicalResult 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()) {
+    return op.emitError("unhandled clauses for translation to LLVM IR");
+  }
+  auto bodyCB = [&](InsertPointTy allocaIP, InsertPointTy codegenIP){
+    LLVM::ModuleTranslation::SaveStack<OpenMPAllocaStackFrame> frame(moduleTranslation, allocaIP);
+    builder.restoreIP(codegenIP);
+    convertOmpOpRegions(op.getRegion(), "omp.teams.region", builder, moduleTranslation, bodyGenStatus);
+  };
+
+  llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
+  builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createTeams(ompLoc, bodyCB));
+  return bodyGenStatus;
+}
+
 /// Converts an OpenMP task construct into LLVM IR using OpenMPIRBuilder.
 static LogicalResult
 convertOmpTaskOp(omp::TaskOp taskOp, llvm::IRBuilderBase &builder,
@@ -2406,6 +2424,9 @@ LogicalResult OpenMPDialectLLVMIRTranslationInterface::convertOperation(
       .Case([&](omp::SingleOp op) {
         return convertOmpSingle(op, builder, moduleTranslation);
       })
+      .Case([&](omp::TeamsOp op) {
+        return convertOmpTeams(op, builder, moduleTranslation);
+      })
       .Case([&](omp::TaskOp op) {
         return convertOmpTaskOp(op, builder, moduleTranslation);
       })
diff --git a/mlir/test/Target/LLVMIR/openmp-teams.mlir b/mlir/test/Target/LLVMIR/openmp-teams.mlir
new file mode 100644
index 000000000000000..c9005fca94a7c20
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-teams.mlir
@@ -0,0 +1,136 @@
+// RUN: mlir-translate -mlir-to-llvmir -split-input-file %s | FileCheck %s
+
+llvm.func @foo()
+
+// CHECK-LABEL: @omp_teams_simple
+// CHECK: call void {{.*}} @__kmpc_fork_teams(ptr @{{.+}}, i32 0, ptr [[wrapperfn:.+]])
+// CHECK: ret void
+llvm.func @omp_teams_simple() {
+    omp.teams {
+        llvm.call @foo() : () -> ()
+        omp.terminator
+    }
+    llvm.return
+}
+
+// CHECK: define internal void @[[outlinedfn:.+]]()
+// CHECK:   call void @foo()
+// CHECK:   ret void
+// CHECK: define void [[wrapperfn]](ptr %[[global_tid:.+]], ptr %[[bound_tid:.+]])
+// CHECK:   call void @[[outlinedfn]]
+// CHECK:   ret void
+
+// -----
+
+llvm.func @foo(i32) -> ()
+
+// CHECK-LABEL: @omp_teams_shared_simple
+// CHECK-SAME: (i32 [[arg0:%.+]])
+// CHECK: [[structArg:%.+]] = alloca { i32 }
+// CHECK: br
+// CHECK: [[gep:%.+]] = getelementptr { i32 }, ptr [[structArg]], i32 0, i32 0
+// CHECK: store i32 [[arg0]], ptr [[gep]]
+// CHECK: call void {{.+}} @__kmpc_fork_teams(ptr @{{.+}}, i32 1, ptr [[wrapperfn:.+]], ptr [[structArg]])
+// CHECK: ret void
+llvm.func @omp_teams_shared_simple(%arg0: i32) {
+    omp.teams {
+        llvm.call @foo(%arg0) : (i32) -> ()
+        omp.terminator
+    }
+    llvm.return
+}
+
+// CHECK: define internal void [[outlinedfn:@.+]](ptr [[structArg:%.+]])
+// CHECK:   [[gep:%.+]] = getelementptr { i32 }, ptr [[structArg]], i32 0, i32 0
+// CHECK:   [[loadgep:%.+]] = load i32, ptr [[gep]]
+// CHECK:   call void @foo(i32 [[loadgep]])
+// CHECK:   ret void
+// CHECK: define void [[wrapperfn]](ptr [[global_tid:.+]], ptr [[bound_tid:.+]], ptr [[structArg:.+]])
+// CHECK:   call void [[outlinedfn]](ptr [[structArg]])
+// CHECK:   ret void
+
+// -----
+
+llvm.func @my_alloca_fn() -> !llvm.ptr<i32>
+llvm.func @foo(i32, f32, !llvm.ptr<i32>, f128, !llvm.ptr<i32>, i32) -> ()
+llvm.func @bar()
+
+// CHECK-LABEL: @omp_teams_branching_shared
+// CHECK-SAME: (i1 [[condition:%.+]], i32 [[arg0:%.+]], float [[arg1:%.+]], ptr [[arg2:%.+]], fp128 [[arg3:%.+]])
+
+// Checking that the allocation for struct argument happens in the alloca block.
+// CHECK: [[structArg:%.+]] = alloca { i1, i32, float, ptr, fp128, ptr, i32 }
+// CHECK: [[allocated:%.+]] = call ptr @my_alloca_fn()
+// CHECK: [[loaded:%.+]] = load i32, ptr [[allocated]]
+// CHECK: br label
+
+// Checking that the shared values are stored properly in the struct arg.
+// CHECK: [[conditionPtr:%.+]] = getelementptr {{.+}}, ptr [[structArg]]
+// CHECK: store i1 [[condition]], ptr [[conditionPtr]]
+// CHECK: [[arg0ptr:%.+]] = getelementptr {{.+}}, ptr [[structArg]], i32 0, i32 1
+// CHECK: store i32 [[arg0]], ptr [[arg0ptr]]
+// CHECK: [[arg1ptr:%.+]] = getelementptr {{.+}}, ptr [[structArg]], i32 0, i32 2
+// CHECK: store float [[arg1]], ptr [[arg1ptr]]
+// CHECK: [[arg2ptr:%.+]] = getelementptr {{.+}}, ptr [[structArg]], i32 0, i32 3
+// CHECK: store ptr [[arg2]], ptr [[arg2ptr]]
+// CHECK: [[arg3ptr:%.+]] = getelementptr {{.+}}, ptr [[structArg]], i32 0, i32 4
+// CHECK: store fp128 [[arg3]], ptr [[arg3ptr]]
+// CHECK: [[allocatedPtr:%.+]] = getelementptr {{.+}}, ptr [[structArg]], i32 0, i32 5
+// CHECK: store ptr [[allocated]], ptr [[allocatedPtr]]
+// CHECK: [[loadedPtr:%.+]] = getelementptr {{.+}}, ptr [[structArg]], i32 0, i32 6
+// CHECK: store i32 [[loaded]], ptr [[loadedPtr]]
+
+// Runtime call.
+// CHECK: call void {{.+}} @__kmpc_fork_teams(ptr @{{.+}}, i32 1, ptr [[wrapperfn:@.+]], ptr [[structArg]])
+// CHECK: br label
+// CHECK: call void @bar()
+// CHECK: ret void
+llvm.func @omp_teams_branching_shared(%condition: i1, %arg0: i32, %arg1: f32, %arg2: !llvm.ptr<i32>, %arg3: f128) {
+    %allocated = llvm.call @my_alloca_fn(): () -> !llvm.ptr<i32>
+    %loaded = llvm.load %allocated : !llvm.ptr<i32>
+    llvm.br ^codegenBlock
+^codegenBlock:
+    omp.teams {
+        llvm.cond_br %condition, ^true_block, ^false_block
+    ^true_block:
+        llvm.call @foo(%arg0, %arg1, %arg2, %arg3, %allocated, %loaded) : (i32, f32, !llvm.ptr<i32>, f128, !llvm.ptr<i32>, i32) -> ()
+        llvm.br ^exit
+    ^false_block:
+        llvm.br ^exit
+    ^exit:
+        omp.terminator
+    }
+    llvm.call @bar() : () -> ()
+    llvm.return
+}
+
+// Check the outlined function.
+// CHECK: define internal void [[outlinedfn:@.+]](ptr [[data:%.+]])
+// CHECK:   [[conditionPtr:%.+]] = getelementptr {{.+}}, ptr [[data]]
+// CHECK:   [[condition:%.+]] = load i1, ptr [[conditionPtr]]
+// CHECK:   [[arg0ptr:%.+]] = getelementptr {{.+}}, ptr [[data]], i32 0, i32 1
+// CHECK:   [[arg0:%.+]] = load i32, ptr [[arg0ptr]]
+// CHECK:   [[arg1ptr:%.+]] = getelementptr {{.+}}, ptr [[data]], i32 0, i32 2
+// CHECK:   [[arg1:%.+]] = load float, ptr [[arg1ptr]]
+// CHECK:   [[arg2ptr:%.+]] = getelementptr {{.+}}, ptr [[data]], i32 0, i32 3
+// CHECK:   [[arg2:%.+]] = load ptr, ptr [[arg2ptr]]
+// CHECK:   [[arg3ptr:%.+]] = getelementptr {{.+}}, ptr [[data]], i32 0, i32 4
+// CHECK:   [[arg3:%.+]] = load fp128, ptr [[arg3ptr]]
+// CHECK:   [[allocatedPtr:%.+]] = getelementptr {{.+}}, ptr [[data]], i32 0, i32 5
+// CHECK:   [[allocated:%.+]] = load ptr, ptr [[allocatedPtr]]
+// CHECK:   [[loadedPtr:%.+]] = getelementptr {{.+}}, ptr [[data]], i32 0, i32 6
+// CHECK:   [[loaded:%.+]] = load i32, ptr [[loadedPtr]]
+// CHECK:   br label
+
+// CHECK:   br i1 [[condition]], label %[[true:.+]], label %[[false:.+]]
+// CHECK: [[false]]:
+// CHECK-NEXT: br label
+// CHECK: [[true]]:
+// CHECK:   call void @foo(i32 [[arg0]], float [[arg1]], ptr [[arg2]], fp128 [[arg3]], ptr [[allocated]], i32 [[loaded]])
+// CHECK-NEXT: br label
+// CHECK: ret void
+
+// Check the wrapper function
+// CHECK: define void [[wrapperfn]](ptr [[globalTID:%.+]], ptr [[boundTID:%.+]], ptr [[data:%.+]])
+// CHECK:   call void [[outlinedfn]](ptr [[data]])
+// CHECK:   ret void

>From 4618cca227680ed3487961560f9419050890db55 Mon Sep 17 00:00:00 2001
From: Shraiysh Vaishay <shraiysh.vaishay at amd.com>
Date: Mon, 2 Oct 2023 20:48:45 -0500
Subject: [PATCH 2/2] Formatting

---
 .../OpenMP/OpenMPToLLVMIRTranslation.cpp      | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index b9643be40e13c01..9eefd1be8ad070e 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -662,20 +662,27 @@ convertOmpSingle(omp::SingleOp &singleOp, llvm::IRBuilderBase &builder,
 }
 
 // Convert an OpenMP Teams construct to LLVM IR using OpenMPIRBuilder
-static LogicalResult convertOmpTeams(omp::TeamsOp op, llvm::IRBuilderBase &builder, LLVM::ModuleTranslation &moduleTranslation) {
+static LogicalResult
+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.getNumTeamsLower() || op.getNumTeamsUpper() || op.getIfExpr() ||
+      op.getThreadLimit() || !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);
+  auto bodyCB = [&](InsertPointTy allocaIP, InsertPointTy codegenIP) {
+    LLVM::ModuleTranslation::SaveStack<OpenMPAllocaStackFrame> frame(
+        moduleTranslation, allocaIP);
     builder.restoreIP(codegenIP);
-    convertOmpOpRegions(op.getRegion(), "omp.teams.region", builder, moduleTranslation, bodyGenStatus);
+    convertOmpOpRegions(op.getRegion(), "omp.teams.region", builder,
+                        moduleTranslation, bodyGenStatus);
   };
 
   llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
-  builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createTeams(ompLoc, bodyCB));
+  builder.restoreIP(
+      moduleTranslation.getOpenMPBuilder()->createTeams(ompLoc, bodyCB));
   return bodyGenStatus;
 }
 



More information about the Openmp-commits mailing list