[MLIR][OpenMP] Add support for critical construct

[kiranchandramohan]

Summary

This patch adds the critical construct to the OpenMP dialect. The
implementation models the definition in 2.17.1 of the OpenMP 5 standard.
A name and hint can be specified. The name is a global entity or has
external linkage, it is modelled as a FlatSymbolRefAttr. Hint is
modelled as an integer enum attribute.

[ftynse]

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

373

structure?

378

Ultra-nit: please align this line vertically with something.

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

243

Nit: isn't there a getValueOr() ?

mlir/test/Dialect/OpenMP/invalid.mlir

301

Do we really want printed names this long? Can't we just have nonspeculative here?

[kiranchandramohan]

Address review comments. Changes include,
-> Using shorter names for hint clause values.
-> Formatting and corrections to description.
-> Using getValueOr in place of hasValue, getValue etc.

[kiranchandramohan]

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

243

Yes. Thanks for pointing this out. Fixed.

mlir/test/Dialect/OpenMP/invalid.mlir

301

I have switched to using nonspeculative instead of omp_sync_hint_nonspeculative here. I agree that from a pure readability and mlir perspective it is good to use short names.

Two reasons for using the long name in the first version,

1. Using the name in the standard might help avoid a conversion while lowering from another intermediate representation (like parse-tree in Flang).
2. Sometimes the clause values can clause clashes with reserved keywords in C++. For eg. for the default clause there were issues with using the private value. https://bugs.llvm.org/show_bug.cgi?id=47225

Using long names, can help avoid such issues.

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

def MasterOp : OpenMP_Op<"master"> {
   }];
 
   let regions = (region AnyRegion:$region);
 
   let assemblyFormat = "$region attr-dict";
 }
 
 //===----------------------------------------------------------------------===//
+// 2.17.1 critical Construct
+//===----------------------------------------------------------------------===//
+// TODO: Autogenerate this from OMP.td in llvm/include/Frontend
+def omp_sync_hint_none: I32EnumAttrCase<"omp_sync_hint_none", 0>;
+def omp_sync_hint_uncontended: I32EnumAttrCase<"omp_sync_hint_uncontended", 1>;
+def omp_sync_hint_contended: I32EnumAttrCase<"omp_sync_hint_contended", 2>;
+def omp_sync_hint_nonspeculative: I32EnumAttrCase<"omp_sync_hint_nonspeculative", 3>;
+def omp_sync_hint_speculative: I32EnumAttrCase<"omp_sync_hint_speculative", 4>;
+
+def SyncHintKind: I32EnumAttr<"SyncHintKind", "OpenMP Sync Hint Kind",
+    [omp_sync_hint_none, omp_sync_hint_uncontended, omp_sync_hint_contended,
+     omp_sync_hint_nonspeculative, omp_sync_hint_speculative]> {
+  let cppNamespace = "::mlir::omp";
+  let stringToSymbolFnName = "ConvertToEnum";
+  let symbolToStringFnName = "ConvertToString";
+}
+
+def CriticalOp : OpenMP_Op<"critical"> {
+  let summary = "critical construct";
+  let description = [{
+    The critical construct imposes a restriction on the associated structured to
+    be executed by only a single thread at a time.
+  }];
+
+  let arguments = (ins OptionalAttr<FlatSymbolRefAttr>:$name,
+             OptionalAttr<SyncHintKind>:$hint);
+
+  let regions = (region AnyRegion:$region);
+
+  let assemblyFormat = [{
+    (`(` $name^ `)`)? (`hint` `(` $hint^ `)`)? $region attr-dict
+  }];
+
+  let verifier = "return ::verifyCriticalOp(*this);";
+}
+
+//===----------------------------------------------------------------------===//
 // 2.17.2 barrier Construct
 //===----------------------------------------------------------------------===//
 
 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.

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

if (decl.getAccumulatorType() && decl.getAccumulatorType() != varType) {
              << ") to be the same type as reduction declaration ("
              << decl.getAccumulatorType() << ")";
     }
   }
 
   return success();
 }
 
+static LogicalResult verifyCriticalOp(CriticalOp op) {
+  if (!op.name().hasValue() && op.hint().hasValue() &&
+      (op.hint().getValue() != SyncHintKind::omp_sync_hint_none))
+    return op.emitOpError() << "must specify a name unless the effect is as if "
+                               "hint(omp_sync_hint_none) is specified";
+  return success();
+}
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOps.cpp.inc"

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

convertOmpMaster(Operation &opInst, llvm::IRBuilderBase &builder,
 
   llvm::OpenMPIRBuilder::LocationDescription ompLoc(
       builder.saveIP(), builder.getCurrentDebugLocation());
   builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createMaster(
       ompLoc, bodyGenCB, finiCB));
   return success();
 }
 
+/// Converts an OpenMP 'critical' operation into LLVM IR using OpenMPIRBuilder.
+static LogicalResult
+convertOmpCritical(Operation &opInst, llvm::IRBuilderBase &builder,
+                   LLVM::ModuleTranslation &moduleTranslation) {
+  using InsertPointTy = llvm::OpenMPIRBuilder::InsertPointTy;
+  auto criticalOp = cast<omp::CriticalOp>(opInst);
+  // TODO: support error propagation in OpenMPIRBuilder and use it instead of
+  // relying on captured variables.
+  LogicalResult bodyGenStatus = success();
+
+  auto bodyGenCB = [&](InsertPointTy allocaIP, InsertPointTy codeGenIP,
+                       llvm::BasicBlock &continuationBlock) {
+    // CriticalOp has only one region associated with it.
+    auto &region = cast<omp::CriticalOp>(opInst).getRegion();
+    convertOmpOpRegions(region, "omp.critical.region", *codeGenIP.getBlock(),
+                        continuationBlock, builder, moduleTranslation,
+                        bodyGenStatus);
+  };
+
+  // TODO: Perform finalization actions for variables. This has to be
+  // called for variables which have destructors/finalizers.
+  auto finiCB = [&](InsertPointTy codeGenIP) {};
+
+  llvm::OpenMPIRBuilder::LocationDescription ompLoc(
+      builder.saveIP(), builder.getCurrentDebugLocation());
+  llvm::LLVMContext &llvmContext = moduleTranslation.getLLVMContext();
+  llvm::Constant *hint = nullptr;
+  if (criticalOp.hint().hasValue()) {
+    hint =
+        llvm::ConstantInt::get(llvm::Type::getInt32Ty(llvmContext),
+                               static_cast<int>(criticalOp.hint().getValue()));
+  } else {
+    hint = llvm::ConstantInt::get(llvm::Type::getInt32Ty(llvmContext), 0);
+  }
+  builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createCritical(
+      ompLoc, bodyGenCB, finiCB,
+      criticalOp.name().hasValue() ? criticalOp.name().getValue() : "", hint));
+  return success();
+}
+
 /// 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 success();
       })
       .Case([&](omp::ParallelOp) {
         return convertOmpParallel(*op, builder, moduleTranslation);
       })
       .Case([&](omp::MasterOp) {
         return convertOmpMaster(*op, builder, moduleTranslation);
       })
+      .Case([&](omp::CriticalOp) {
+        return convertOmpCritical(*op, builder, moduleTranslation);
+      })
       .Case([&](omp::WsLoopOp) {
         return convertOmpWsLoop(*op, builder, moduleTranslation);
       })
       .Case<omp::YieldOp, omp::TerminatorOp>([](auto op) {
         // `yield` and `terminator` can be just omitted. The block structure was
         // created in the function that handles their parent operation.
         assert(op->getNumOperands() == 0 &&
                "unexpected OpenMP terminator with operands");

mlir/test/Dialect/OpenMP/invalid.mlir

func @foo(%lb : index, %ub : index, %step : index, %mem : memref<1xf32>) {
   omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step)
   reduction(@add_f32 -> %mem : memref<1xf32>) {
     %2 = constant 2.0 : f32
     omp.reduction %2, %mem : memref<1xf32>
     omp.yield
   }
   return
 }
+
+// -----
+
+func @omp_critical() -> () {
+  // expected-error @below {{must specify a name unless the effect is as if hint(omp_sync_hint_none) is specified}}
+  omp.critical hint(omp_sync_hint_nonspeculative) {
+    omp.terminator
+  }
+  return
+}

mlir/test/Dialect/OpenMP/ops.mlir

reduction(@add2_f32 -> %0 : memref<1xf32>) {
     %1 = constant 2.0 : f32
     // CHECK: omp.reduction
     omp.reduction %1, %0 : memref<1xf32>
     omp.yield
   }
   return
 }
 
+// CHECK-LABEL: omp_critical
+func @omp_critical() -> () {
+  omp.critical {
+    omp.terminator
+  }
+
+  omp.critical(@mutex) hint(omp_sync_hint_nonspeculative) {
+    omp.terminator
+  }
+  return
+}

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

omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(guided) {
   // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u
   // CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0
   // CHECK  br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}}
    llvm.call @body(%iv) : (i64) -> ()
    omp.yield
  }
  llvm.return
 }
+
+// CHECK-LABEL: @omp_critical
+llvm.func @omp_critical(%x : !llvm.ptr<i32>, %xval : i32) -> () {
+  // CHECK: call void @__kmpc_critical_with_hint({{.*}}critical_user_.var{{.*}}, i32 0)
+  // CHECK: br label %omp.critical.region
+  // CHECK: omp.critical.region
+  omp.critical {
+  // CHECK: store
+    llvm.store %xval, %x : !llvm.ptr<i32>
+    omp.terminator
+  }
+  // CHECK: call void @__kmpc_end_critical({{.*}}critical_user_.var{{.*}})
+
+  // CHECK: call void @__kmpc_critical_with_hint({{.*}}critical_user_mutex.var{{.*}}, i32 2)
+  // CHECK: br label %omp.critical.region
+  // CHECK: omp.critical.region
+  omp.critical(@mutex) hint(omp_sync_hint_contended) {
+  // CHECK: store
+    llvm.store %xval, %x : !llvm.ptr<i32>
+    omp.terminator
+  }
+  // CHECK: call void @__kmpc_end_critical({{.*}}critical_user_mutex.var{{.*}})
+
+  llvm.return
+}