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

[kiranchandramohan]

Summary

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.

[shraiysh]

LGTM. Thank you!

[awarzynski]

flang/test/Lower/OpenMP/flush.f90

12–13

[nit] I'd be inclined to change this to:

!LLVMIRDialect-NOT: ((.*}}!fir{{.*}}

Basically, the actual operation is irrelevant. What's important is the fhat there are no fir operations. And once you keep it generic, you will never have to worry about the type of the operation (might be handy when fir.ref is replaced with something else).

[peixin]

LGTM

flang/test/Lower/OpenMP/flush.f90

4

Nit: Should this be moved to test/Fir since the tests under Lower/OpenMP are targeted for lowering test? It's also OK to do this when you upstream and refactor all the OpenMP lowering tests.

[kiranchandramohan]

Fix MLIR tests.

[kiranchandramohan]

flang/test/Lower/OpenMP/flush.f90

4

Yes, will do this as a separate patch.

12–13

As @peixin suggests above, we will move the conversion test to a different file. This test will only retain the FIR+OpenMP lowering and in that case, retaining the FIR type is important.

[awarzynski]

flang/test/Lower/OpenMP/flush.f90

12–13

So you'll be removing the fir-to-llvm-ir conversion from here?

Also, after re-rereading this, I'm finally realizing what was meant here. Below is my interpretation:

!FIRDialect: omp.flush(%{{.*}}, %{{.*}}, %{{.*}} : !fir.ref<i32>, !fir.ref<i32>, !fir.ref<i32>)
!LLVMIRDialect: omp.flush(%{{.*}}, %{{.*}}, %{{.*}} : !llvm.ptr<i32>, !llvm.ptr<i32>, !llvm.ptr<i32>)

I probably would prefer the above (it's more explicit). For the following, it is hard to guess that !OMPDialect and !FIRDialect/!LLVMDialect are to be matched on the same line:

!OMPDialect: omp.flush(%{{.*}}, %{{.*}}, %{{.*}} :
!FIRDialect: !fir.ref<i32>, !fir.ref<i32>, !fir.ref<i32>)
!LLVMIRDialect: !llvm.ptr<i32>, !llvm.ptr<i32>, !llvm.ptr<i32>)

I think that the real problem here is that lowering and conversion are tested in one file, but that's going to change anyway, right?

[kiranchandramohan]

flang/test/Lower/OpenMP/flush.f90

12–13

Yes, we will be moving fir-to-llvm-conversion to a separate test file that tests the conversion to LLVM.

Yes, you are correct that the real problem is the twin testing in this file. Once it moves then we will be replacing all these with just CHECK lines and the operation and the operands will hence be in the same line.

Given the above, is it OK to proceed with this patch as it is?

[awarzynski]

LGTM!

flang/test/Lower/OpenMP/flush.f90

 ! This test checks lowering of OpenMP Flush Directive.
 
 !RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | FileCheck %s --check-prefixes="FIRDialect,OMPDialect"
-!RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | fir-opt --cfg-conversion | fir-opt --fir-to-llvm-ir | FileCheck %s --check-prefixes="OMPDialect"
+!RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | fir-opt --cfg-conversion | fir-opt --fir-to-llvm-ir | FileCheck %s --check-prefixes="LLVMIRDialect,OMPDialect"
 
 subroutine flush_standalone(a, b, c)
     integer, intent(inout) :: a, b, c
 
 !$omp flush(a,b,c)
 !$omp flush
-!OMPDialect: omp.flush(%{{.*}}, %{{.*}}, %{{.*}} : !fir.ref<i32>, !fir.ref<i32>, !fir.ref<i32>)
+!OMPDialect: omp.flush(%{{.*}}, %{{.*}}, %{{.*}} :
+!FIRDialect: !fir.ref<i32>, !fir.ref<i32>, !fir.ref<i32>)
+!LLVMIRDialect: !llvm.ptr<i32>, !llvm.ptr<i32>, !llvm.ptr<i32>)
 !OMPDialect: omp.flush
 
 end subroutine flush_standalone
 
 subroutine flush_parallel(a, b, c)
     integer, intent(inout) :: a, b, c
 
 !$omp parallel
 !OMPDialect:  omp.parallel {
 
-!OMPDialect: omp.flush(%{{.*}}, %{{.*}}, %{{.*}} : !fir.ref<i32>, !fir.ref<i32>, !fir.ref<i32>)
+!OMPDialect: omp.flush(%{{.*}}, %{{.*}}, %{{.*}} :
+!FIRDialect: !fir.ref<i32>, !fir.ref<i32>, !fir.ref<i32>)
+!LLVMIRDialect: !llvm.ptr<i32>, !llvm.ptr<i32>, !llvm.ptr<i32>)
 !OMPDialect: omp.flush
 !$omp flush(a,b,c)
 !$omp flush
 
 !FIRDialect: %{{.*}} = fir.load %{{.*}} : !fir.ref<i32>
 !FIRDialect: %{{.*}} = fir.load %{{.*}} : !fir.ref<i32>
 !FIRDialect: %{{.*}} = arith.addi %{{.*}}, %{{.*}} : i32
 !FIRDialect: fir.store %{{.*}} to %{{.*}} : !fir.ref<i32>

flang/test/Lower/OpenMP/parallel-sections.f90

 
 !FIRDialect: func @_QPomp_parallel_sections
 subroutine omp_parallel_sections_allocate(x, y)
   use omp_lib
   integer, intent(inout) :: x, y
   !FIRDialect: %[[allocator:.*]] = arith.constant 1 : i32
   !LLVMDialect: %[[allocator:.*]] = llvm.mlir.constant(1 : i32) : i32
   !OMPDialect: omp.parallel {
-  !OMPDialect: omp.sections allocate(%[[allocator]] : i32 -> %{{.*}} : !fir.ref<i32>) {
+  !OMPDialect: omp.sections allocate(
+  !FIRDialect: %[[allocator]] : i32 -> %{{.*}} : !fir.ref<i32>) {
+  !LLVMDialect: %[[allocator]] : i32 -> %{{.*}} : !llvm.ptr<i32>) {
   !$omp parallel sections allocate(omp_high_bw_mem_alloc: x)
     !OMPDialect: omp.section {
     !$omp section
       x = x + 12
       !OMPDialect: omp.terminator
     !OMPDialect: omp.section {
     !$omp section
       y = y + 5
       !OMPDialect: omp.terminator
   !OMPDialect: omp.terminator
   !OMPDialect: omp.terminator
   !$omp end parallel sections
 end subroutine omp_parallel_sections_allocate

flang/test/Lower/OpenMP/parallel.f90

 !RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | FileCheck %s --check-prefixes="FIRDialect,OMPDialect"
-!RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | fir-opt --fir-to-llvm-ir | FileCheck %s --check-prefixes="OMPDialect"
+!RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | fir-opt --fir-to-llvm-ir | FileCheck %s --check-prefixes="LLVMDialect,OMPDialect"
 
 !FIRDialect-LABEL: func @_QPparallel_simple
 subroutine parallel_simple()
    !OMPDialect: omp.parallel
 !$omp parallel
    !FIRDialect: fir.call
    call f1()
 !$omp end parallel

 !===============================================================================
 ! `allocate` clause
 !===============================================================================
 
 !FIRDialect-LABEL: func @_QPparallel_allocate
 subroutine parallel_allocate()
    use omp_lib
    integer :: x
-   !OMPDialect: omp.parallel allocate(%{{.+}} : i32 -> %{{.+}} : !fir.ref<i32>) {
+   !OMPDialect: omp.parallel allocate(
+   !FIRDialect: %{{.+}} : i32 -> %{{.+}} : !fir.ref<i32>
+   !LLVMDialect: %{{.+}} : i32 -> %{{.+}} : !llvm.ptr<i32>
+   !OMPDialect: ) {
    !$omp parallel allocate(omp_high_bw_mem_alloc: x) private(x)
    !FIRDialect: arith.addi
    x = x + 12
    !OMPDialect: omp.terminator
    !$omp end parallel
 end subroutine parallel_allocate
 
 !===============================================================================
subroutine parallel_multiple_clauses(alpha, num_threads)
 
    !OMPDialect: omp.parallel if({{.*}} : i1) num_threads({{.*}} : i32) {
    !$omp parallel num_threads(num_threads) if(alpha .le. 0)
    !FIRDialect: fir.call
    call f3()
    !OMPDialect: omp.terminator
    !$omp end parallel
 
-   !OMPDialect: omp.parallel if({{.*}} : i1) num_threads({{.*}} : i32) allocate(%{{.+}} : i32 -> %{{.+}} : !fir.ref<i32>) {
+   !OMPDialect: omp.parallel if({{.*}} : i1) num_threads({{.*}} : i32) allocate(
+   !FIRDialect: %{{.+}} : i32 -> %{{.+}} : !fir.ref<i32>
+   !LLVMDialect: %{{.+}} : i32 -> %{{.+}} : !llvm.ptr<i32>
+   !OMPDialect: ) {
    !$omp parallel num_threads(num_threads) if(alpha .le. 0) allocate(omp_high_bw_mem_alloc: alpha) private(alpha)
    !FIRDialect: fir.call
    call f3()
    !FIRDialect: arith.addi
    alpha = alpha + 12
    !OMPDialect: omp.terminator
    !$omp end parallel
 
 end subroutine parallel_multiple_clauses

flang/test/Lower/OpenMP/single.f90

 !RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | FileCheck %s --check-prefixes="FIRDialect,OMPDialect"
-!RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | fir-opt --cfg-conversion | fir-opt --fir-to-llvm-ir | FileCheck %s --check-prefixes="OMPDialect"
+!RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | fir-opt --cfg-conversion | fir-opt --fir-to-llvm-ir | FileCheck %s --check-prefixes="LLVMDialect,OMPDialect"
 
 !===============================================================================
 ! Single construct
 !===============================================================================
 
 !FIRDialect-LABEL: func @_QPomp_single
 !FIRDialect-SAME: (%[[x:.*]]: !fir.ref<i32> {fir.bindc_name = "x"})
 subroutine omp_single(x)

 !===============================================================================
 
 !FIRDialect-LABEL: func @_QPsingle_allocate
 subroutine single_allocate()
   use omp_lib
   integer :: x
   !OMPDialect: omp.parallel {
   !$omp parallel
-  !OMPDialect: omp.single allocate(%{{.+}} : i32 -> %{{.+}} : !fir.ref<i32>) {
+  !OMPDialect: omp.single allocate(
+  !FIRDialect: %{{.+}} : i32 -> %{{.+}} : !fir.ref<i32>
+  !LLVMDialect: %{{.+}} : i32 -> %{{.+}} : !llvm.ptr<i32>
+  !OMPDialect: ) {
   !$omp single allocate(omp_high_bw_mem_alloc: x) private(x)
   !FIRDialect: arith.addi
   x = x + 12
   !OMPDialect: omp.terminator
   !$omp end single
   !OMPDialect: omp.terminator
   !$omp end parallel
 end subroutine single_allocate

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

def FlushOp : OpenMP_Op<"flush"> {
   let summary = "flush construct";
   let description = [{
     The flush construct executes the OpenMP flush operation. This operation
     makes a thread’s temporary view of memory consistent with memory and
     enforces an order on the memory operations of the variables explicitly
     specified or implied.
   }];
 
-  let arguments = (ins Variadic<AnyType>:$varList);
+  let arguments = (ins Variadic<OpenMP_PointerLikeType>:$varList);
 
   let assemblyFormat = [{ ( `(` $varList^ `:` type($varList) `)` )? attr-dict}];
+  let extraClassDeclaration = [{
+    /// The number of variable operands.
+    unsigned getNumVariableOperands() {
+      return getOperation()->getNumOperands();
+    }
+    /// The i-th variable operand passed.
+    Value getVariableOperand(unsigned i) {
+      return getOperand(i);
+    }
+  }];
 }
 //===----------------------------------------------------------------------===//
 // 2.14.5 target construct
 //===----------------------------------------------------------------------===//
 
 def TargetOp : OpenMP_Op<"target",[AttrSizedOperandSegments]> {
   let summary = "target construct";
   let description = [{

mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp

if (failed(rewriter.convertRegionTypes(&newOp.region(),
       return failure();
 
     rewriter.eraseOp(curOp);
     return success();
   }
 };
 
 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,
                   ConversionPatternRewriter &rewriter) const override {
     TypeConverter *converter = ConvertToLLVMPattern::getTypeConverter();
     SmallVector<Type> resTypes;
     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)
         return failure();
       if (originalVariableOperand.getType().isa<MemRefType>()) {
         // TODO: Support memref type in variable operands
         return rewriter.notifyMatchFailure(curOp,
                                            "memref is not supported yet");
       }
       convertedOperands.emplace_back(adaptor.getOperands()[idx]);
     }
     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>(
-      [&](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 {
 struct ConvertOpenMPToLLVMPass
     : public ConvertOpenMPToLLVMBase<ConvertOpenMPToLLVMPass> {
   void runOnOperation() override;
 };
 } // namespace

mlir/test/Dialect/OpenMP/ops.mlir

 
 func.func @omp_taskyield() -> () {
   // CHECK: omp.taskyield
   omp.taskyield
   return
 }
 
 // CHECK-LABEL: func @omp_flush
-// CHECK-SAME: ([[ARG0:%.*]]: i32) {
-func.func @omp_flush(%arg0 : i32) -> () {
+// CHECK-SAME: ([[ARG0:%.*]]: memref<i32>) {
+func.func @omp_flush(%arg0 : memref<i32>) -> () {
   // Test without data var
   // CHECK: omp.flush
   omp.flush
 
   // Test with one data var
-  // CHECK: omp.flush([[ARG0]] : i32)
-  omp.flush(%arg0 : i32)
+  // CHECK: omp.flush([[ARG0]] : memref<i32>)
+  omp.flush(%arg0 : memref<i32>)
 
   // Test with two data var
-  // CHECK: omp.flush([[ARG0]], [[ARG0]] : i32, i32)
-  omp.flush(%arg0, %arg0: i32, i32)
+  // CHECK: omp.flush([[ARG0]], [[ARG0]] : memref<i32>, memref<i32>)
+  omp.flush(%arg0, %arg0: memref<i32>, memref<i32>)
 
   return
 }
 
 func.func @omp_terminator() -> () {
   // CHECK: omp.terminator
   omp.terminator
 }

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

llvm.func @test_stand_alone_directives() {
   // CHECK: [[OMP_THREAD2:%.*]] = call i32 @__kmpc_global_thread_num(ptr @{{[0-9]+}})
   // CHECK-NEXT:  [[RET_VAL:%.*]] = call i32 @__kmpc_omp_taskyield(ptr @{{[0-9]+}}, i32 [[OMP_THREAD2]], i32 0)
   omp.taskyield
 
   // CHECK-NEXT:    ret void
   llvm.return
 }
 
-// CHECK-LABEL: define void @test_flush_construct(i32 %0)
-llvm.func @test_flush_construct(%arg0: i32) {
+// CHECK-LABEL: define void @test_flush_construct(ptr %{{[0-9]+}})
+llvm.func @test_flush_construct(%arg0: !llvm.ptr<i32>) {
   // CHECK: call void @__kmpc_flush(ptr @{{[0-9]+}}
   omp.flush
 
   // CHECK: call void @__kmpc_flush(ptr @{{[0-9]+}}
-  omp.flush (%arg0 : i32)
+  omp.flush (%arg0 : !llvm.ptr<i32>)
 
   // CHECK: call void @__kmpc_flush(ptr @{{[0-9]+}}
-  omp.flush (%arg0, %arg0 : i32, i32)
+  omp.flush (%arg0, %arg0 : !llvm.ptr<i32>, !llvm.ptr<i32>)
 
   %0 = llvm.mlir.constant(1 : i64) : i64
   //  CHECK: alloca {{.*}} align 4
   %1 = llvm.alloca %0 x i32 {in_type = i32, name = "a"} : (i64) -> !llvm.ptr<i32>
   // CHECK: call void @__kmpc_flush(ptr @{{[0-9]+}}
   omp.flush
   //  CHECK: load i32, ptr
   %2 = llvm.load %1 : !llvm.ptr<i32>