[MLIR] Support for taskwait and taskyield operations, and translating the same to LLVM IR

[kiranktp]

Summary

This patch adds support for taskwait and taskyield operations in OpenMP dialect and translation of the these constructs to LLVM IR. The OpenMP IRBuilder is used for this translation.
The patch includes code changes and a testcase modifications.

[mehdi_amini]

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

75

The declarative assembly syntax may be shorter here :)

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

331

The sequence may be representable with a TypeSwitch here?

[kiranchandramohan]

Thanks @kiranktp for working on this. Looks OK to me.

Just pinging @rogfer01 since he wrote the portion of the builder for taskwait and taskyield.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

310–311

There might be more than one LLVM IR operation corresponding to an OpenMP MLIR operation.

Would the following be better?
Given an OpenMP MLIR operation, create the corresponding LLVM IR (including OpenMP runtime calls).

mlir/test/Target/openmp-llvm.mlir

4–7

Should we move all these checks closer to the operation? What is the MLIR/LLVM style for this?

[SouraVX]

mlir/test/Target/openmp-llvm.mlir

4–7

They can be moved closer to operations https://llvm.org/docs/CommandGuide/FileCheck.html, it's also better for better readibility and cleanly extending test case further in future.

[kiranktp]

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

75

Did you mean, this declarative syntax is too short here. Could you please elaborate a bit on this?

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

310–311

But this function will be called for single OpenMP MLIR operation.
Nevertheless, your wording for comment is neutral and more clear. I will incorporate them.

331

I was not aware of TypeSwitch feature. Thanks for the suggestion Mehdi. I will modify as you suggested.

mlir/test/Target/openmp-llvm.mlir

4–7

Thanks for the review Kiran.
Yes, we can move these checks closure to the operation. I will do this modification.

[mehdi_amini]

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

75

Sorry I was terse, I meant to use a different mechanism here possibly, look the section about « Declarative Assembly Format» here : https://mlir.llvm.org/docs/OpDefinitions/

[kiranktp]

Incorporated all review comments.

[kiranktp]

Hi Mehdi, Thanks for the review. I do not have commit access. Could you please merge this code on my behalf.?

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

 
 def BarrierOp : OpenMP_Op<"barrier"> {
   let summary = "barrier construct";
   let description = [{
     The barrier construct specifies an explicit barrier at the point at which
     the construct appears.
   }];
 
-  let parser = [{ return success(); }];
-  let printer = [{ p << getOperationName(); }];
+  let assemblyFormat = "attr-dict";
+}
+
+def TaskwaitOp : OpenMP_Op<"taskwait"> {
+  let summary = "taskwait construct";
+  let description = [{
+    The taskwait construct specifies a wait on the completion of child tasks
+    of the current task.
+  }];
+
+  let assemblyFormat = "attr-dict";
+}
+
+def TaskyieldOp : OpenMP_Op<"taskyield"> {
+  let summary = "taskyield construct";
+  let description = [{
+    The taskyield construct specifies that the current task can be suspended
+    in favor of execution of a different task.
+  }];
+
+  let assemblyFormat = "attr-dict";
 }
 
 #endif // OPENMP_OPS

mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h

protected:
   /// LLVM IR module. The MLIR LLVM IR dialect holds a pointer to an
   /// LLVMContext, the LLVM IR module will be created in that context.
   ModuleTranslation(Operation *module,
                     std::unique_ptr<llvm::Module> llvmModule);
   virtual ~ModuleTranslation();
 
   virtual LogicalResult convertOperation(Operation &op,
                                          llvm::IRBuilder<> &builder);
+  virtual LogicalResult convertOmpOperation(Operation &op,
+                                            llvm::IRBuilder<> &builder);
   static std::unique_ptr<llvm::Module> prepareLLVMModule(Operation *m);
 
   /// A helper to look up remapped operands in the value remapping table.
   SmallVector<llvm::Value *, 8> lookupValues(ValueRange values);
 
 private:
   /// Check whether the module contains only supported ops directly in its body.
   static LogicalResult checkSupportedModuleOps(Operation *m);

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

 //===- ModuleTranslation.cpp - MLIR to LLVM conversion --------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the translation between an MLIR LLVM dialect module and
 // the corresponding LLVMIR module. It only handles core LLVM IR operations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Target/LLVMIR/ModuleTranslation.h"
 
 #include "DebugTranslation.h"
+#include "mlir/ADT/TypeSwitch.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/Support/LLVM.h"
 
 #include "llvm/ADT/SetVector.h"
: mlirModule(module), llvmModule(std::move(llvmModule)),
           std::make_unique<DebugTranslation>(module, *this->llvmModule)),
       ompDialect(
           module->getContext()->getRegisteredDialect<omp::OpenMPDialect>()) {
   assert(satisfiesLLVMModule(mlirModule) &&
          "mlirModule should honor LLVM's module semantics.");
 }
 ModuleTranslation::~ModuleTranslation() {}
 
+/// Given an OpenMP MLIR operation, create the corresponding LLVM IR
+/// (including OpenMP runtime calls).
+LogicalResult
+ModuleTranslation::convertOmpOperation(Operation &opInst,
+                                       llvm::IRBuilder<> &builder) {
+  if (!ompBuilder) {
+    ompBuilder = std::make_unique<llvm::OpenMPIRBuilder>(*llvmModule);
+    ompBuilder->initialize();
+  }
+  return mlir::TypeSwitch<Operation *, LogicalResult>(&opInst)
+      .Case([&](omp::BarrierOp) {
+        ompBuilder->CreateBarrier(builder.saveIP(), llvm::omp::OMPD_barrier);
+        return success();
+      })
+      .Case([&](omp::TaskwaitOp) {
+        ompBuilder->CreateTaskwait(builder.saveIP());
+        return success();
+      })
+      .Case([&](omp::TaskyieldOp) {
+        ompBuilder->CreateTaskyield(builder.saveIP());
+        return success();
+      })
+      .Default([&](Operation *inst) {
+        return inst->emitError("unsupported OpenMP operation: ")
+               << inst->getName();
+      });
+}
+
 /// Given a single MLIR operation, create the corresponding LLVM IR operation
 /// using the `builder`.  LLVM IR Builder does not have a generic interface so
 /// this has to be a long chain of `if`s calling different functions with a
 /// different number of arguments.
 LogicalResult ModuleTranslation::convertOperation(Operation &opInst,
                                                   llvm::IRBuilder<> &builder) {
   auto extractPosition = [](ArrayAttr attr) {
     SmallVector<unsigned, 4> position;
if (auto addressOfOp = dyn_cast<LLVM::AddressOfOp>(opInst)) {
     // The verifier should not have allowed this.
     assert(global && "referencing an undefined global");
 
     valueMapping[addressOfOp.getResult()] = globalsMapping.lookup(global);
     return success();
   }
 
   if (opInst.getDialect() == ompDialect) {
-    if (!ompBuilder) {
-      ompBuilder = std::make_unique<llvm::OpenMPIRBuilder>(*llvmModule);
-      ompBuilder->initialize();
-    }
-
-    if (isa<omp::BarrierOp>(opInst)) {
-      ompBuilder->CreateBarrier(builder.saveIP(), llvm::omp::OMPD_barrier);
-      return success();
-    }
-    return opInst.emitError("unsupported OpenMP operation: ")
-           << opInst.getName();
+    return convertOmpOperation(opInst, builder);
   }
 
   return opInst.emitError("unsupported or non-LLVM operation: ")
          << opInst.getName();
 }
 
 /// Convert block to LLVM IR.  Unless `ignoreArguments` is set, emit PHI nodes
 /// to define values corresponding to the MLIR block arguments.  These nodes

mlir/test/Dialect/OpenMP/ops.mlir

 // RUN: mlir-opt -verify-diagnostics %s | FileCheck %s
 
 func @omp_barrier() -> () {
   // CHECK: omp.barrier
   omp.barrier
   return
 }
+
+func @omp_taskwait() -> () {
+  // CHECK: omp.taskwait
+  omp.taskwait
+  return
+}
+
+func @omp_taskyield() -> () {
+  // CHECK: omp.taskyield
+  omp.taskyield
+  return
+}

mlir/test/Target/openmp-llvm.mlir

 // RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
 
 // CHECK-LABEL: define void @empty()
+llvm.func @empty() {
   // CHECK: [[OMP_THREAD:%.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @{{[0-9]+}})
   // CHECK-NEXT:  call void @__kmpc_barrier(%struct.ident_t* @{{[0-9]+}}, i32 [[OMP_THREAD]])
-// CHECK-NEXT:    ret void
-llvm.func @empty() {
   omp.barrier
+
+  // CHECK: [[OMP_THREAD1:%.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @{{[0-9]+}})
+  // CHECK-NEXT:  [[RET_VAL:%.*]] = call i32 @__kmpc_omp_taskwait(%struct.ident_t* @{{[0-9]+}}, i32 [[OMP_THREAD1]])
+  omp.taskwait
+
+  // CHECK: [[OMP_THREAD2:%.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @{{[0-9]+}})
+  // CHECK-NEXT:  [[RET_VAL:%.*]] = call i32 @__kmpc_omp_taskyield(%struct.ident_t* @{{[0-9]+}}, i32 [[OMP_THREAD2]], i32 0)
+  omp.taskyield
+
+// CHECK-NEXT:    ret void
   llvm.return
 }