[MLIR] Support for flush operation, and translating the same to LLVM IR

[kiranktp]

Summary

This patch adds support for flush operation in OpenMP dialect and translation of this construct to LLVM IR.
The OpenMP IRBuilder is used for this translation.
The patch includes code changes and testcase modifications.

[ftynse]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

320

Let's not use Fortran syntax in C++ code

321

This comment doesn't seem to belong here

mlir/test/Dialect/OpenMP/ops.mlir

27

This doesn't actually check that you print back operands

28

Could we rather use the custom format?

[kiranchandramohan]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

322

Is this a TODO?
Do we require another version of CreateFlush in the OpenMP Builder to support the list?

mlir/test/Target/openmp-llvm.mlir

4

Shall we have more functions to test? Maybe one for those with arguments and one for those without? Also, rename this function.

20

Is it possible to check that the arguments are also passed correctly for this version of omp.flush?

[kiranktp]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

320

I will remove the Fortran syntax

321

This function is specific to openmp Op convertion and also this comment is very generic to flush construct (both C/Fortran). So I placed it here.

322

Openmp runtime will discard the list even if its passed. This is the same behavior with clang and legacy flang as well. The list will be discarded for flush construct.
Same has been documented in the OpenMP standard as well.
From OpenMP 4.5 Standard Doc:
"An implementation may implement a flush with a list by ignoring the list, and treating it the same as a flush without a list."
"Note – Use of a flush construct with a list is extremely error prone and users are strongly discouraged from attempting it."
I think this is the reason, there is only one version of CreateFlush without this list.
If the decision is to process the list and pass it on to IR builder, another version of CreateFlush has to be implemented.
In spite of this, the list will be discarded and a call to __kmpc_flush will be inserted without the list in IR builder.

mlir/test/Target/openmp-llvm.mlir

4

I will add a separate function with arguments and give a meaning full name.

[kiranktp]

Hi @ftynse, @kiranchandramohan,
I have incorporated the review comments.

[kiranktp]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

321

I have added the comment around flush construct handling.

mlir/test/Dialect/OpenMP/ops.mlir

27

I have updated the test case to validate the operands as well.

28

The current changes looks fine?

mlir/test/Target/openmp-llvm.mlir

20

Anyway the argument list is being discarded, we need not check for correctness of the arguments.

[ftynse]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

339

Please comply with the linter (run git clang-format on your patch)

[kiranchandramohan]

LGTM.

[kiranktp]

Fixed formatting issue.

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

def BarrierOp : OpenMP_Op<"barrier"> {
   let description = [{
     The barrier construct specifies an explicit barrier at the point at which
     the construct appears.
   }];
 
   let assemblyFormat = "attr-dict";
 }
 
+//===----------------------------------------------------------------------===//
+// 2.13.7 flush Construct
+//===----------------------------------------------------------------------===//
+def FlushOp : OpenMP_Op<"flush"> {
+  let summary = "flush construct";
+  let description = [{
+    The flush construct executes the OpenMP flush operation. This operation
+    makes a thread’s temporary view of memory consistent with memory and
+    enforces an order on the memory operations of the variables explicitly
+    specified or implied.
+  }];
+
+  let arguments = (ins Variadic<AnyType>:$varList);
+
+  let assemblyFormat = "attr-dict ($varList^ `:` type($varList))?";
+}
+
 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";

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

 /// (including OpenMP runtime calls).
 LogicalResult
 ModuleTranslation::convertOmpOperation(Operation &opInst,
                                        llvm::IRBuilder<> &builder) {
   if (!ompBuilder) {
     ompBuilder = std::make_unique<llvm::OpenMPIRBuilder>(*llvmModule);
     ompBuilder->initialize();
   }
   return llvm::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();
       })
+      .Case([&](omp::FlushOp) {
+        // No support in Openmp runtime funciton (__kmpc_flush) to accept
+        // the argument list.
+        // OpenMP standard states the following:
+        //  "An implementation may implement a flush with a list by ignoring
+        //   the list, and treating it the same as a flush without a list."
+        //
+        // The argument list is discarded so that, flush with a list is treated
+        // same as a flush without a list.
+        ompBuilder->CreateFlush(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

mlir/test/Dialect/OpenMP/ops.mlir

func @omp_taskwait() -> () {
   return
 }
 
 func @omp_taskyield() -> () {
   // CHECK: omp.taskyield
   omp.taskyield
   return
 }
+
+// CHECK-LABEL: func @omp_flush
+// CHECK-SAME: %[[ARG0:.*]]: !llvm.i32
+func @omp_flush(%arg0 : !llvm.i32) -> () {
+  // Test without data var
+  // CHECK: omp.flush
+  omp.flush
+
+  // Test with one data var
+  // CHECK: omp.flush %[[ARG0]] : !llvm.i32
+  "omp.flush"(%arg0) : (!llvm.i32) -> ()
+
+  // Test with two data var
+  // CHECK: omp.flush %[[ARG0]], %[[ARG0]] : !llvm.i32, !llvm.i32
+  "omp.flush"(%arg0, %arg0): (!llvm.i32, !llvm.i32) -> ()
+
+  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-LABEL: define void @test_stand_alone_directives()
+llvm.func @test_stand_alone_directives() {
   // 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]])
   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
 }
+
+// CHECK-LABEL: define void @test_flush_construct(i32 %0)
+llvm.func @test_flush_construct(%arg0: !llvm.i32) {
+  // CHECK: call void @__kmpc_flush(%struct.ident_t* @{{[0-9]+}}
+  omp.flush
+
+  // CHECK: call void @__kmpc_flush(%struct.ident_t* @{{[0-9]+}}
+  omp.flush %arg0 : !llvm.i32
+
+  // CHECK: call void @__kmpc_flush(%struct.ident_t* @{{[0-9]+}}
+  omp.flush %arg0, %arg0 : !llvm.i32, !llvm.i32
+
+  // CHECK-NEXT:    ret void
+  llvm.return
+}