[mlir][OpenMP] omp.task translation to LLVM IR

[shraiysh]

Summary

This patch adds translation for omp.task from OpenMPDialect to LLVM IR
Dialect and adds tests for the same.

Depends on D71989

[ftynse]

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

682

Nit: there is no need to pass MLIR ops by reference.

[kiranchandramohan]

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

697

Nit: spell out this auto.

mlir/test/Target/LLVMIR/openmp-llvm.mlir

2178

Nit: are x, y, zaddr used to check that there are no arguments for the wrapper function and hence the task data size is 0?

2187

Nit: Is the alloca significant for your test? Do you want to check that it is in the outlined region?

2226

Why is the task size 16 and not 12 (for the three liveins)?

[shraiysh]

Addressed comments.

mlir/test/Target/LLVMIR/openmp-llvm.mlir

2178

Yes, nothing from outer scope should be counted for tasksize, I will change it to just %x and %zaddr if you'd like.

2187

Yes, added check for it. Thanks for pointing this out.

2226

x and y are 4-byte each, and the pointer zaddr is of 8 bytes (64-bit). So, the total is 16.

[shraiysh]

Rebase with main. Requesting review again as it has been a while since this was accepted.

[Meinersbur]

A bit-ness question

[shraiysh]

In D123919#3542687, @Meinersbur wrote:

A bit-ness question

I didn’t understand this comment. Is this a typo or incomplete comment?

[shraiysh]

Ping for review!

[kiranchandramohan]

LG. Please wait a day before submitting.

[peixin]

LGTM

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

1329–1330

Nit: move "call." to line 1328.

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

689

Nit

691

Is allocaIP used in this?

[Meinersbur]

In D123919#3542895, @shraiysh wrote:

In D123919#3542687, @Meinersbur wrote:

A bit-ness question

I didn’t understand this comment. Is this a typo or incomplete comment?

Sorry, I added an inline-comment but Phabricator refused to submit without also adding something in the main comment box. Now, I don't see the inline-comment so that may be the reason Phabricator complained.

mlir/test/Target/LLVMIR/openmp-llvm.mlir

2226

A pointer is 8 bytes only on 64-bit platforms. Does this mean that this test will fail on platforms that have pointers of different sizes? I also don't see the target triple fixed by mlir-translate, so will this test fail on non-64 hosts?

[shraiysh]

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

691

It is not used in this instance, but it is required in OpenMPIRBuilder for the most general case and so this is required.

mlir/test/Target/LLVMIR/openmp-llvm.mlir

2226

It could. I don't know who should set the target-triple/data-layout. Should it be a command line option for mlir-translate? A similar test is there in OpenMP IRBuilder here, which could cause issues (it hasn't caused any so far) and I should probably rectify that too.

[Meinersbur]

mlir/test/Target/LLVMIR/openmp-llvm.mlir

2226

The OpenMPIRBuilder tests use an empty TargetTriple ("") which represents an 64-bit architecture with byte-aligment. What does mlir-translate do?

[shraiysh]

Added target triple to make sure that the test doesn't fail on non-x86 machines.

[shraiysh]

mlir/test/Target/LLVMIR/openmp-llvm.mlir

2226

I have added the target triple to the testcase now. This should pass it on other machines. Does this look okay now?

[shraiysh]

@Meinersbur does this patch look okay now? I wanted to merge this, if there are no issues.

[shraiysh]

If there are no further issues, I will merge this tomorrow. It would be nice to have E2E task support.

[Meinersbur]

LGTM, thank you.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

 
 OpenMPIRBuilder::InsertPointTy
 OpenMPIRBuilder::createTask(const LocationDescription &Loc,
                             InsertPointTy AllocaIP, BodyGenCallbackTy BodyGenCB,
                             bool Tied, Value *Final) {
   if (!updateToLocation(Loc))
     return InsertPointTy();
 
+  uint32_t SrcLocStrSize;
+  Constant *SrcLocStr = getOrCreateSrcLocStr(Loc, SrcLocStrSize);
+  Value *Ident = getOrCreateIdent(SrcLocStr, SrcLocStrSize);
   // The current basic block is split into four basic blocks. After outlining,
   // they will be mapped as follows:
   // ```
   // def current_fn() {
   //   current_basic_block:
   //     br label %task.exit
   //   task.exit:
   //     ; instructions after task
OpenMPIRBuilder::createTask(const LocationDescription &Loc,
   BasicBlock *TaskBodyBB = splitBB(Builder, /*CreateBranch=*/true, "task.body");
   BasicBlock *TaskAllocaBB =
       splitBB(Builder, /*CreateBranch=*/true, "task.alloca");
 
   OutlineInfo OI;
   OI.EntryBB = TaskAllocaBB;
   OI.OuterAllocaBB = AllocaIP.getBlock();
   OI.ExitBB = TaskExitBB;
-  OI.PostOutlineCB = [this, &Loc, Tied, Final](Function &OutlinedFn) {
+  OI.PostOutlineCB = [this, Ident, Tied, Final](Function &OutlinedFn) {
     // The input IR here looks like the following-
     // ```
     // func @current_fn() {
     //   outlined_fn(%args)
     // }
     // func @outlined_fn(%args) { ... }
     // ```
     //
OI.PostOutlineCB = [this, Ident, Tied, Final](Function &OutlinedFn) {
     Builder.SetInsertPoint(StaleCI);
 
     // Gather the arguments for emitting the runtime call for
     // @__kmpc_omp_task_alloc
     Function *TaskAllocFn =
         getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_omp_task_alloc);
 
     // Arguments - `loc_ref` (Ident) and `gtid` (ThreadID)
     // call.
-    uint32_t SrcLocStrSize;
-    Constant *SrcLocStr = getOrCreateSrcLocStr(Loc, SrcLocStrSize);
-    Value *Ident = getOrCreateIdent(SrcLocStr, SrcLocStrSize);
     Value *ThreadID = getOrCreateThreadID(Ident);
 
     // Argument - `flags`
     // Task is tied iff (Flags & 1) == 1.
     // Task is untied iff (Flags & 1) == 0.
     // Task is final iff (Flags & 2) == 2.
     // Task is not final iff (Flags & 2) == 0.
     // TODO: Handle the other flags.
     Value *Flags = Builder.getInt32(Tied);

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

convertOmpOpRegions(singleOp.region(), "omp.single.region", builder,
                         moduleTranslation, bodyGenStatus);
   };
   auto finiCB = [&](InsertPointTy codeGenIP) {};
   builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createSingle(
       ompLoc, bodyCB, finiCB, singleOp.nowait(), /*DidIt=*/nullptr));
   return bodyGenStatus;
 }
 
+/// Converts an OpenMP task construct into LLVM IR using OpenMPIRBuilder.
+static LogicalResult
+convertOmpTaskOp(omp::TaskOp taskOp, llvm::IRBuilderBase &builder,
+                 LLVM::ModuleTranslation &moduleTranslation) {
+  using InsertPointTy = llvm::OpenMPIRBuilder::InsertPointTy;
+  LogicalResult bodyGenStatus = success();
+  if (taskOp.if_expr() || taskOp.final_expr() || taskOp.untiedAttr() ||
+      taskOp.mergeableAttr() || taskOp.in_reductions() || taskOp.priority() ||
+      !taskOp.allocate_vars().empty()) {
+    return taskOp.emitError("unhandled clauses for translation to LLVM IR");
+  }
+  auto bodyCB = [&](InsertPointTy allocaIP, InsertPointTy codegenIP) {
+    builder.restoreIP(codegenIP);
+    convertOmpOpRegions(taskOp.region(), "omp.task.region", builder,
+                        moduleTranslation, bodyGenStatus);
+  };
+  llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
+      findAllocaInsertPoint(builder, moduleTranslation);
+  llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
+  builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createTask(
+      ompLoc, allocaIP, bodyCB, !taskOp.untied()));
+  return bodyGenStatus;
+}
+
 /// Converts an OpenMP workshare loop into LLVM IR using OpenMPIRBuilder.
 static LogicalResult
 convertOmpWsLoop(Operation &opInst, llvm::IRBuilderBase &builder,
                  LLVM::ModuleTranslation &moduleTranslation) {
   auto loop = cast<omp::WsLoopOp>(opInst);
   // TODO: this should be in the op verifier instead.
   if (loop.lowerBound().empty())
     return failure();
return llvm::TypeSwitch<Operation *, LogicalResult>(op)
         return convertOmpAtomicCapture(op, builder, moduleTranslation);
       })
       .Case([&](omp::SectionsOp) {
         return convertOmpSections(*op, builder, moduleTranslation);
       })
       .Case([&](omp::SingleOp op) {
         return convertOmpSingle(op, builder, moduleTranslation);
       })
+      .Case([&](omp::TaskOp op) {
+        return convertOmpTaskOp(op, builder, moduleTranslation);
+      })
       .Case<omp::YieldOp, omp::TerminatorOp, omp::ReductionDeclareOp,
             omp::CriticalDeclareOp>([](auto op) {
         // `yield` and `terminator` can be just omitted. The block structure
         // was created in the region that handles their parent operation.
         // `reduction.declare` will be used by reductions and is not
         // converted directly, skip it.
         // `critical.declare` is only used to declare names of critical
         // sections which will be used by `critical` ops and hence can be

mlir/test/Target/LLVMIR/openmp-llvm.mlir

omp.parallel  {
     omp.terminator
   }
 
   llvm.store %2, %4 : !llvm.ptr<i32>
   llvm.return
 }
 
 llvm.mlir.global internal @_QFsubEx() : i32
+
+// -----
+
+// CHECK-LABEL: define void @omp_task
+// CHECK-SAME: (i32 %[[x:.+]], i32 %[[y:.+]], ptr %[[zaddr:.+]])
+llvm.func @omp_task(%x: i32, %y: i32, %zaddr: !llvm.ptr<i32>) {
+  // CHECK: %[[omp_global_thread_num:.+]] = call i32 @__kmpc_global_thread_num({{.+}})
+  // CHECK: %[[task_data:.+]] = call ptr @__kmpc_omp_task_alloc
+  // CHECK-SAME: (ptr @{{.+}}, i32 %[[omp_global_thread_num]], i32 1, i64 0,
+  // CHECK-SAME:  i64 0, ptr @[[wrapper_fn:.+]])
+  // CHECK: call i32 @__kmpc_omp_task(ptr @{{.+}}, i32 %[[omp_global_thread_num]], ptr %[[task_data]])
+  omp.task {
+    %n = llvm.mlir.constant(1 : i64) : i64
+    %valaddr = llvm.alloca %n x i32 : (i64) -> !llvm.ptr<i32>
+    %val = llvm.load %valaddr : !llvm.ptr<i32>
+    %double = llvm.add %val, %val : i32
+    llvm.store %double, %valaddr : !llvm.ptr<i32>
+    omp.terminator
+  }
+  llvm.return
+}
+
+// CHECK: define internal void @[[outlined_fn:.+]]()
+// CHECK: task.alloca{{.*}}:
+// CHECK:   br label %[[task_body:[^, ]+]]
+// CHECK: [[task_body]]:
+// CHECK:   br label %[[task_region:[^, ]+]]
+// CHECK: [[task_region]]:
+// CHECK:   %[[alloca:.+]] = alloca i32, i64 1
+// CHECK:   %[[val:.+]] = load i32, ptr %[[alloca]]
+// CHECK:   %[[newval:.+]] = add i32 %[[val]], %[[val]]
+// CHECK:   store i32 %[[newval]], ptr %{{[^, ]+}}
+// CHECK:   br label %[[exit_stub:[^, ]+]]
+// CHECK: [[exit_stub]]:
+// CHECK:   ret void
+
+
+// CHECK: define i32 @[[wrapper_fn]](i32 %{{.+}}) {
+// CHECK:   call void @[[outlined_fn]]()
+// CHECK:   ret i32 0
+// CHECK: }
+
+// -----
+
+// CHECK-LABEL: define void @omp_task
+// CHECK-SAME: (i32 %[[x:.+]], i32 %[[y:.+]], ptr %[[zaddr:.+]])
+module attributes {llvm.target_triple = "x86_64-unknown-linux-gnu"} {
+  llvm.func @omp_task(%x: i32, %y: i32, %zaddr: !llvm.ptr<i32>) {
+    // CHECK: %[[diff:.+]] = sub i32 %[[x]], %[[y]],
+    %diff = llvm.sub %x, %y : i32
+    // CHECK: store i32 %[[diff]], ptr %2
+    llvm.store %diff, %zaddr : !llvm.ptr<i32>
+    // CHECK: %[[omp_global_thread_num:.+]] = call i32 @__kmpc_global_thread_num({{.+}})
+    // CHECK: %[[task_data:.+]] = call ptr @__kmpc_omp_task_alloc
+    // CHECK-SAME: (ptr @{{.+}}, i32 %[[omp_global_thread_num]], i32 1, i64 16, i64 0,
+    // CHECK-SAME: ptr @[[wrapper_fn:.+]])
+    // CHECK: call void @llvm.memcpy.p0.p0.i64(ptr {{.+}} %[[task_data]], ptr {{.+}}, i64 16, i1 false)
+    // CHECK: call i32 @__kmpc_omp_task(ptr @{{.+}}, i32 %[[omp_global_thread_num]], ptr %[[task_data]])
+    omp.task {
+      %z = llvm.add %x, %y : i32
+      llvm.store %z, %zaddr : !llvm.ptr<i32>
+      omp.terminator
+    }
+    // CHECK: %[[prod:.+]] = mul i32 %[[x]], %[[y]]
+    %b = llvm.mul %x, %y : i32
+    // CHECK: store i32 %[[prod]], ptr %[[zaddr]]
+    llvm.store %b, %zaddr : !llvm.ptr<i32>
+    llvm.return
+  }
+}
+
+// CHECK: define internal void @[[outlined_fn:.+]](ptr %[[task_data:.+]])
+// CHECK: task.alloca{{.*}}:
+// CHECK:   br label %[[task_body:[^, ]+]]
+// CHECK: [[task_body]]:
+// CHECK:   br label %[[task_region:[^, ]+]]
+// CHECK: [[task_region]]:
+// CHECK:   %[[sum:.+]] = add i32 %{{.+}}, %{{.+}}
+// CHECK:   store i32 %[[sum]], ptr %{{.+}}
+// CHECK:   br label %[[exit_stub:[^, ]+]]
+// CHECK: [[exit_stub]]:
+// CHECK:   ret void
+
+
+// CHECK: define i32 @[[wrapper_fn]](i32 %{{.+}}, ptr %[[task_data:.+]]) {
+// CHECK:   call void @[[outlined_fn]](ptr %[[task_data]])
+// CHECK:   ret i32 0
+// CHECK: }