[MLIR,OpenMP] : Add Conversion pattern for Critical Op

[kiranchandramohan]

Summary

The Conversion pattern enables conversion of Critical Op with block
arguments.

Fixes https://github.com/llvm/llvm-project/issues/56629

[kiranchandramohan]

There is an issue with the test. I will fix it before landing.

[kiranchandramohan]

Fix the test.

mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp

rewriter.replaceOpWithNewOp<T>(curOp, resTypes, convertedOperands,
                                    curOp->getAttrs());
     return success();
   }
 };
 } // namespace
 
 void mlir::configureOpenMPToLLVMConversionLegality(
     ConversionTarget &target, LLVMTypeConverter &typeConverter) {
-  target.addDynamicallyLegalOp<mlir::omp::ParallelOp, mlir::omp::WsLoopOp,
-                               mlir::omp::MasterOp, mlir::omp::SectionsOp,
-                               mlir::omp::SingleOp>([&](Operation *op) {
+  target.addDynamicallyLegalOp<mlir::omp::CriticalOp, mlir::omp::ParallelOp,
+                               mlir::omp::WsLoopOp, mlir::omp::MasterOp,
+                               mlir::omp::SectionsOp, mlir::omp::SingleOp>(
+      [&](Operation *op) {
         return typeConverter.isLegal(&op->getRegion(0)) &&
                typeConverter.isLegal(op->getOperandTypes()) &&
                typeConverter.isLegal(op->getResultTypes());
       });
   target
       .addDynamicallyLegalOp<mlir::omp::AtomicReadOp, mlir::omp::AtomicWriteOp,
                              mlir::omp::FlushOp, mlir::omp::ThreadprivateOp>(
           [&](Operation *op) {
             return typeConverter.isLegal(op->getOperandTypes()) &&
                    typeConverter.isLegal(op->getResultTypes());
           });
 }
 
 void mlir::populateOpenMPToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                                   RewritePatternSet &patterns) {
   patterns.add<
-      RegionOpConversion<omp::MasterOp>, RegionOpConversion<omp::ParallelOp>,
-      RegionOpConversion<omp::WsLoopOp>, RegionOpConversion<omp::SectionsOp>,
-      RegionOpConversion<omp::SingleOp>,
+      RegionOpConversion<omp::CriticalOp>, RegionOpConversion<omp::MasterOp>,
+      RegionOpConversion<omp::ParallelOp>, RegionOpConversion<omp::WsLoopOp>,
+      RegionOpConversion<omp::SectionsOp>, RegionOpConversion<omp::SingleOp>,
       RegionLessOpWithVarOperandsConversion<omp::AtomicReadOp>,
       RegionLessOpWithVarOperandsConversion<omp::AtomicWriteOp>,
       RegionLessOpWithVarOperandsConversion<omp::FlushOp>,
       RegionLessOpWithVarOperandsConversion<omp::ThreadprivateOp>>(converter);
 }
 
 namespace {
 struct ConvertOpenMPToLLVMPass

mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir

 // RUN: mlir-opt -convert-openmp-to-llvm -split-input-file %s | FileCheck %s
 
+// CHECK-LABEL: llvm.func @foo(i64, i64)
+func.func private @foo(index, index)
+
+// CHECK-LABEL: llvm.func @critical_block_arg
+func.func @critical_block_arg() {
+  // CHECK: omp.critical
+  omp.critical {
+  // CHECK-NEXT: ^[[BB0:.*]](%[[ARG1:.*]]: i64, %[[ARG2:.*]]: i64):
+  ^bb0(%arg1: index, %arg2: index):
+    // CHECK-NEXT: llvm.call @foo(%[[ARG1]], %[[ARG2]]) : (i64, i64) -> ()
+    func.call @foo(%arg1, %arg2) : (index, index) -> ()
+    omp.terminator
+  }
+  return
+}
+
+// -----
+
 // CHECK-LABEL: llvm.func @master_block_arg
 func.func @master_block_arg() {
   // CHECK: omp.master
   omp.master {
   // CHECK-NEXT: ^[[BB0:.*]](%[[ARG1:.*]]: i64, %[[ARG2:.*]]: i64):
   ^bb0(%arg1: index, %arg2: index):
     // CHECK-DAG: %[[CAST_ARG1:.*]] = builtin.unrealized_conversion_cast %[[ARG1]] : i64 to index
     // CHECK-DAG: %[[CAST_ARG2:.*]] = builtin.unrealized_conversion_cast %[[ARG2]] : i64 to index