[MLIR] Add support for declaring critical construct names

[kiranchandramohan]

Summary

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.

[jeanPerier]

Looks like a good solution to me.

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
 //===----------------------------------------------------------------------===//
+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>;
 def omp_sync_hint_contended: I32EnumAttrCase<"contended", 2>;
 def omp_sync_hint_nonspeculative: I32EnumAttrCase<"nonspeculative", 3>;
 def omp_sync_hint_speculative: I32EnumAttrCase<"speculative", 4>;
 
 def SyncHintKind: I32EnumAttr<"SyncHintKind", "OpenMP Sync Hint Kind",

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

static LogicalResult verifyWsLoopOp(WsLoopOp op) {
   return success();
 }
 
 static LogicalResult verifyCriticalOp(CriticalOp op) {
   if (!op.name().hasValue() && op.hint().hasValue() &&
       (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();
 }
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOps.cpp.inc"

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

return llvm::TypeSwitch<Operation *, LogicalResult>(op)
         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, 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) {
+        // `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.
       .Default([&](Operation *inst) {
         return inst->emitError("unsupported OpenMP operation: ")
                << inst->getName();
       });
 }
 
 void mlir::registerOpenMPDialectTranslation(DialectRegistry &registry) {
   registry.insert<omp::OpenMPDialect>();

mlir/test/Dialect/OpenMP/invalid.mlir

reduction(@add_f32 -> %mem : memref<1xf32>) {
     omp.reduction %2, %mem : memref<1xf32>
     omp.yield
   }
   return
 }
 
 // -----
 
-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

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

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

omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(guided) {
    llvm.call @body(%iv) : (i64) -> ()
    omp.yield
  }
  llvm.return
 }
 
 // -----
 
+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)
   // CHECK: br label %omp.critical.region
   // CHECK: omp.critical.region
   omp.critical {
   // CHECK: store
     llvm.store %xval, %x : !llvm.ptr<i32>