[Flang,MLIR,OpenMP] Fix a few tests that were not converting to LLVM

A few OpenMP tests were retaining the FIR operands even after running
the LLVM conversion pass. To fix these tests the legality checkes for
OpenMP conversion are made stricter to include operands and results.
The Flush, Single and Sections operations are added to conversions or
legality checks. The RegionLessOpConversion is appropriately renamed
to clarify that it works only for operations with Variable operands.
The operands of the flush operation are changed to match those of
Variable Operands.

Fix for an OpenMP issue mentioned in
https://github.com/llvm/llvm-project/issues/55210.

Reviewed By: shraiysh, peixin, awarzynski

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

1 files changed, 24 insertions, 12 deletions

diff --git a/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp b/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp
index 3c8416b..bb8f523 100644
--- a/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp
+++ b/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp
@@ -47,7 +47,8 @@ struct RegionOpConversion : public ConvertOpToLLVMPattern<OpType> {
 };
 
 template <typename T>
-struct RegionLessOpConversion : public ConvertOpToLLVMPattern<T> {
+struct RegionLessOpWithVarOperandsConversion
+    : public ConvertOpToLLVMPattern<T> {
   using ConvertOpToLLVMPattern<T>::ConvertOpToLLVMPattern;
   LogicalResult
   matchAndRewrite(T curOp, typename T::Adaptor adaptor,
@@ -57,6 +58,9 @@ struct RegionLessOpConversion : public ConvertOpToLLVMPattern<T> {
     if (failed(converter->convertTypes(curOp->getResultTypes(), resTypes)))
       return failure();
     SmallVector<Value> convertedOperands;
+    assert(curOp.getNumVariableOperands() ==
+               curOp.getOperation()->getNumOperands() &&
+           "unexpected non-variable operands");
     for (unsigned idx = 0; idx < curOp.getNumVariableOperands(); ++idx) {
       Value originalVariableOperand = curOp.getVariableOperand(idx);
       if (!originalVariableOperand)
@@ -78,23 +82,31 @@ struct RegionLessOpConversion : public ConvertOpToLLVMPattern<T> {
 void mlir::configureOpenMPToLLVMConversionLegality(
     ConversionTarget &target, LLVMTypeConverter &typeConverter) {
   target.addDynamicallyLegalOp<mlir::omp::ParallelOp, mlir::omp::WsLoopOp,
-                               mlir::omp::MasterOp>(
-      [&](Operation *op) { return typeConverter.isLegal(&op->getRegion(0)); });
+                               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::ThreadprivateOp>([&](Operation *op) {
-        return typeConverter.isLegal(op->getOperandTypes());
-      });
+                             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>,
-               RegionLessOpConversion<omp::AtomicReadOp>,
-               RegionLessOpConversion<omp::AtomicWriteOp>,
-               RegionLessOpConversion<omp::ThreadprivateOp>>(converter);
+  patterns.add<
+      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 {