[MLIR][OpenMP] Add support for declaring critical construct names

Add an operation omp.critical.declare to declare names/symbols of
critical sections. Named omp.critical operations should use symbols
declared by omp.critical.declare. Having a declare operation ensures
that the names of critical sections are global and unique. In the
lowering flow to LLVM IR, the OpenMP IRBuilder creates unique names
for critical sections.

Reviewed By: ftynse, jeanPerier

Differential Revision: https://reviews.llvm.org/D108713

Author: kiranchandramohan

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

@@ -337,6 +337,21 @@ def MasterOp : OpenMP_Op<"master"> {
 //===----------------------------------------------------------------------===//
 // 2.17.1 critical Construct
 //===----------------------------------------------------------------------===//
+def CriticalDeclareOp : OpenMP_Op<"critical.declare", [Symbol]> {
+  let summary = "declares a named critical section.";
+
+  let description = [{
+    Declares a named critical section.
+
+    The name can be used in critical constructs in the dialect.
+  }];
+
+  let arguments = (ins SymbolNameAttr:$sym_name);
+
+  let assemblyFormat = "$sym_name attr-dict";
+}
+
+
 // TODO: Autogenerate this from OMP.td in llvm/include/Frontend
 def omp_sync_hint_none: I32EnumAttrCase<"none", 0>;
 def omp_sync_hint_uncontended: I32EnumAttrCase<"uncontended", 1>;

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

@@ -1007,6 +1007,17 @@ static LogicalResult verifyCriticalOp(CriticalOp op) {
       (op.hint().getValue() != SyncHintKind::none))
     return op.emitOpError() << "must specify a name unless the effect is as if "
                                "hint(none) is specified";
+
+  if (op.nameAttr()) {
+    auto symbolRef = op.nameAttr().cast<SymbolRefAttr>();
+    auto decl =
+        SymbolTable::lookupNearestSymbolFrom<CriticalDeclareOp>(op, symbolRef);
+    if (!decl) {
+      return op.emitOpError() << "expected symbol reference " << symbolRef
+                              << " to point to a critical declaration";
+    }
+  }
+
   return success();
 }
 

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

@@ -833,14 +833,18 @@ LogicalResult OpenMPDialectLLVMIRTranslationInterface::convertOperation(
       .Case([&](omp::WsLoopOp) {
         return convertOmpWsLoop(*op, builder, moduleTranslation);
       })
-      .Case<omp::YieldOp, omp::TerminatorOp, omp::ReductionDeclareOp>(
-          [](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.
-            return success();
-          })
+      .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
+        // ignored for lowering. The OpenMP IRBuilder will create unique
+        // name for critical section names.
+        return success();
+      })
       .Default([&](Operation *inst) {
         return inst->emitError("unsupported OpenMP operation: ")
                << inst->getName();

mlir/test/Dialect/OpenMP/invalid.mlir

@@ -296,10 +296,20 @@ func @foo(%lb : index, %ub : index, %step : index, %mem : memref<1xf32>) {
 
 // -----
 
-func @omp_critical() -> () {
+func @omp_critical1() -> () {
   // expected-error @below {{must specify a name unless the effect is as if hint(none) is specified}}
   omp.critical hint(nonspeculative) {
     omp.terminator
   }
   return
 }
+
+// -----
+
+func @omp_critical2() -> () {
+  // expected-error @below {{expected symbol reference @excl to point to a critical declaration}}
+  omp.critical(@excl) hint(speculative) {
+    omp.terminator
+  }
+  return
+}

mlir/test/Dialect/OpenMP/ops.mlir

@@ -367,6 +367,10 @@ func @reduction2(%lb : index, %ub : index, %step : index) {
   return
 }
 
+// CHECK: omp.critical.declare
+// CHECK-LABEL: @mutex
+omp.critical.declare @mutex
+
 // CHECK-LABEL: omp_critical
 func @omp_critical() -> () {
   omp.critical {

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

@@ -541,6 +541,8 @@ llvm.func @collapse_wsloop(
 
 // -----
 
+omp.critical.declare @mutex
+
 // 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)