[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

672

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

[kiranchandramohan]

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

687

Nit: spell out this auto.

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

2098

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?

2107

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

2146

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

[shraiysh]

Addressed comments.

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

2098

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

2107

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

2146

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 ↗

(On Diff #431883)

Nit: move "call." to line 1328.

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

679

Nit

681

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

2146

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

681

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

2146

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

2146

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

2146

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.

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

convertOmpOpRegions(singleOp.region(), "omp.single.region",
                         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,
+                    llvm::BasicBlock &continuationBB) {
+    convertOmpOpRegions(taskOp.region(), "omp.task.region",
+                        *codegenIP.getBlock(), continuationBB, builder,
+                        moduleTranslation, bodyGenStatus);
+  };
+  llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
+      findAllocaInsertPoint(builder, moduleTranslation);
+  llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
+  builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createTask(
+      ompLoc, allocaIP, bodyCB));
+  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:.+]], i32* %[[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 i8* @__kmpc_omp_task_alloc
+  // CHECK-SAME: (%{{.+}}* @1, i32 %[[omp_global_thread_num]], i32 1, i64 0,
+  // CHECK-SAME:  i64 0, i32 (i32, i8*)* bitcast (i32 (i32)* @[[wrapper_fn:.+]] to
+  // CHECK-SAME:  i32 (i32, i8*)*))
+  // CHECK: call i32 @__kmpc_omp_task(%{{.*}}, i32 %[[omp_global_thread_num]], i8* %[[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, i32* %[[alloca]]
+// CHECK:   %[[newval:.+]] = add i32 %[[val]], %[[val]]
+// CHECK:   store i32 %[[newval]], i32* %{{[^, ]+}}
+// 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:.+]], i32* %[[zaddr:.+]])
+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]], i32* %2
+  llvm.store %diff, %zaddr : !llvm.ptr<i32>
+  // CHECK: %[[omp_global_thread_num:.+]] = call i32 @__kmpc_global_thread_num({{.+}})
+  // CHECK: %[[task_data:.+]] = call i8* @__kmpc_omp_task_alloc
+  // CHECK-SAME: (%{{.+}}, i32 %[[omp_global_thread_num]], i32 1, i64 16, i64 0,
+  // CHECK-SAME: i32 (i32, i8*)* bitcast (i32 (i32, { i32, i32, i32* }*)*
+  // CHECK-SAME: @[[wrapper_fn:.+]] to i32 (i32, i8*)*))
+  // CHECK: call void @llvm.memcpy.p0i8.p0i8.i64(i8* {{.+}} %[[task_data]], i8* {{.+}}, i64 16, i1 false)
+  // CHECK: call i32 @__kmpc_omp_task(%{{.*}}, i32 %[[omp_global_thread_num]], i8* %[[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]], i32* %[[zaddr]]
+  llvm.store %b, %zaddr : !llvm.ptr<i32>
+  llvm.return
+}
+
+// CHECK: define internal void @[[outlined_fn:.+]]({ i32, i32, i32* }* %[[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]], i32* %{{.+}}
+// CHECK:   br label %[[exit_stub:[^, ]+]]
+// CHECK: [[exit_stub]]:
+// CHECK:   ret void
+
+
+// CHECK: define i32 @[[wrapper_fn]](i32 %{{.+}}, { i32, i32, i32* }* %[[task_data:.+]]) {
+// CHECK:   call void @[[outlined_fn]]({ i32, i32, i32* }* %[[task_data]])
+// CHECK:   ret i32 0
+// CHECK: }