[MLIR][OpenMP] Add conversion support from FIR to LLVM Dialect for OMP Target Data with region

[TIFitis]

Summary

This enables conversion of OpenMP Target Data op with region from FIR Dialect to LLVM IR Dialect.

[TIFitis]

When trying to lower the following Fortran code to FIR and then LLVMIR I am getting an error in fir-opt which I'm not quite sure how to fix. Any hints/inputs would be greatly appreciated.

Fortran Code:

subroutine openmp_target_data
    integer :: a(1024), i
    !$omp target data map(tofrom: a)
        do i = 1, 1024
            a(i) = i
        end do
    !$omp end target data
end subroutine openmp_target_data

FIR: flang-new -fc1 -emit-fir -fopenmp test.f90 -o test.fir

func.func @_QPopenmp_target_data() {
  %0 = fir.alloca !fir.array<1024xi32> {bindc_name = "a", uniq_name = "_QFopenmp_target_dataEa"}
  %1 = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFopenmp_target_dataEi"}
  omp.target_data   map((tofrom -> %0 : !fir.ref<!fir.array<1024xi32>>)) {
    %c1_i32 = arith.constant 1 : i32
    %2 = fir.convert %c1_i32 : (i32) -> index
    %c1024_i32 = arith.constant 1024 : i32
    %3 = fir.convert %c1024_i32 : (i32) -> index
    %c1 = arith.constant 1 : index
    %4 = fir.convert %2 : (index) -> i32
    %5:2 = fir.do_loop %arg0 = %2 to %3 step %c1 iter_args(%arg1 = %4) -> (index, i32) {
      fir.store %arg1 to %1 : !fir.ref<i32>
      %6 = fir.load %1 : !fir.ref<i32>
      %7 = fir.load %1 : !fir.ref<i32>
      %8 = fir.convert %7 : (i32) -> i64
      %c1_i64 = arith.constant 1 : i64
      %9 = arith.subi %8, %c1_i64 : i64
      %10 = fir.coordinate_of %0, %9 : (!fir.ref<!fir.array<1024xi32>>, i64) -> !fir.ref<i32>
      fir.store %6 to %10 : !fir.ref<i32>
      %11 = arith.addi %arg0, %c1 : index
      %12 = fir.convert %c1 : (index) -> i32
      %13 = fir.load %1 : !fir.ref<i32>
      %14 = arith.addi %13, %12 : i32
      fir.result %11, %14 : index, i32
    }
    fir.store %5#1 to %1 : !fir.ref<i32>
    omp.terminator
  }
  return
}

LLVMIR: fir-opt --fir-to-llvm-ir="target=aarch64-unknown-linux-gnu" test.fir -o test.mlir

test.fir:12:14: error: failed to legalize operation 'fir.do_loop'
      %5:2 = fir.do_loop %arg0 = %2 to %3 step %c1 iter_args(%arg1 = %4) -> (index, i32) {
             ^
test.fir:12:14: note: see current operation: 
%18:2 = "fir.do_loop"(%9, %13, %15, %17) ({
^bb0(%arg0: index, %arg1: i32):
  "fir.store"(%arg1, %5) : (i32, !fir.ref<i32>) -> ()
  %19 = "fir.load"(%5) : (!fir.ref<i32>) -> i32
  %20 = "fir.load"(%5) : (!fir.ref<i32>) -> i32
  %21 = "fir.convert"(%20) : (i32) -> i64
  %22 = "arith.constant"() {value = 1 : i64} : () -> i64
  %23 = "arith.subi"(%21, %22) : (i64, i64) -> i64
  %24 = "fir.coordinate_of"(%2, %23) {baseType = !fir.ref<!fir.array<1024xi32>>} : (!fir.ref<!fir.array<1024xi32>>, i64) -> !fir.ref<i32>
  "fir.store"(%19, %24) : (i32, !fir.ref<i32>) -> ()
  %25 = "arith.addi"(%arg0, %15) : (index, index) -> index
  %26 = "fir.convert"(%15) : (index) -> i32
  %27 = "fir.load"(%5) : (!fir.ref<i32>) -> i32
  %28 = "arith.addi"(%27, %26) : (i32, i32) -> i32
  "fir.result"(%25, %28) : (index, i32) -> ()
}) {finalValue} : (index, index, index, i32) -> (index, i32)

Thanks,
Akash

[kiranchandramohan]

In D144380#4138622, @TIFitis wrote:

When trying to lower the following Fortran code to FIR and then LLVMIR I am getting an error in fir-opt which I'm not quite sure how to fix. Any hints/inputs would be greatly appreciated.

Fortran Code:

subroutine openmp_target_data
    integer :: a(1024), i
    !$omp target data map(tofrom: a)
        do i = 1, 1024
            a(i) = i
        end do
    !$omp end target data
end subroutine openmp_target_data

FIR: flang-new -fc1 -emit-fir -fopenmp test.f90 -o test.fir

func.func @_QPopenmp_target_data() {
  %0 = fir.alloca !fir.array<1024xi32> {bindc_name = "a", uniq_name = "_QFopenmp_target_dataEa"}
  %1 = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFopenmp_target_dataEi"}
  omp.target_data   map((tofrom -> %0 : !fir.ref<!fir.array<1024xi32>>)) {
    %c1_i32 = arith.constant 1 : i32
    %2 = fir.convert %c1_i32 : (i32) -> index
    %c1024_i32 = arith.constant 1024 : i32
    %3 = fir.convert %c1024_i32 : (i32) -> index
    %c1 = arith.constant 1 : index
    %4 = fir.convert %2 : (index) -> i32
    %5:2 = fir.do_loop %arg0 = %2 to %3 step %c1 iter_args(%arg1 = %4) -> (index, i32) {
      fir.store %arg1 to %1 : !fir.ref<i32>
      %6 = fir.load %1 : !fir.ref<i32>
      %7 = fir.load %1 : !fir.ref<i32>
      %8 = fir.convert %7 : (i32) -> i64
      %c1_i64 = arith.constant 1 : i64
      %9 = arith.subi %8, %c1_i64 : i64
      %10 = fir.coordinate_of %0, %9 : (!fir.ref<!fir.array<1024xi32>>, i64) -> !fir.ref<i32>
      fir.store %6 to %10 : !fir.ref<i32>
      %11 = arith.addi %arg0, %c1 : index
      %12 = fir.convert %c1 : (index) -> i32
      %13 = fir.load %1 : !fir.ref<i32>
      %14 = arith.addi %13, %12 : i32
      fir.result %11, %14 : index, i32
    }
    fir.store %5#1 to %1 : !fir.ref<i32>
    omp.terminator
  }
  return
}

LLVMIR: fir-opt --fir-to-llvm-ir="target=aarch64-unknown-linux-gnu" test.fir -o test.mlir

test.fir:12:14: error: failed to legalize operation 'fir.do_loop'
      %5:2 = fir.do_loop %arg0 = %2 to %3 step %c1 iter_args(%arg1 = %4) -> (index, i32) {
             ^
test.fir:12:14: note: see current operation: 
%18:2 = "fir.do_loop"(%9, %13, %15, %17) ({
^bb0(%arg0: index, %arg1: i32):
  "fir.store"(%arg1, %5) : (i32, !fir.ref<i32>) -> ()
  %19 = "fir.load"(%5) : (!fir.ref<i32>) -> i32
  %20 = "fir.load"(%5) : (!fir.ref<i32>) -> i32
  %21 = "fir.convert"(%20) : (i32) -> i64
  %22 = "arith.constant"() {value = 1 : i64} : () -> i64
  %23 = "arith.subi"(%21, %22) : (i64, i64) -> i64
  %24 = "fir.coordinate_of"(%2, %23) {baseType = !fir.ref<!fir.array<1024xi32>>} : (!fir.ref<!fir.array<1024xi32>>, i64) -> !fir.ref<i32>
  "fir.store"(%19, %24) : (i32, !fir.ref<i32>) -> ()
  %25 = "arith.addi"(%arg0, %15) : (index, index) -> index
  %26 = "fir.convert"(%15) : (index) -> i32
  %27 = "fir.load"(%5) : (!fir.ref<i32>) -> i32
  %28 = "arith.addi"(%27, %26) : (i32, i32) -> i32
  "fir.result"(%25, %28) : (index, i32) -> ()
}) {finalValue} : (index, index, index, i32) -> (index, i32)

Thanks,
Akash

FIR to LLVM IR conversion happens as a sequence of passes. All FIR operations cannot be converted to LLVM using the fir-to-llvm-ir Conversion pass. We have to run other passes before it. In this specific case only cfg-conversion need to be run.

fir-opt --cfg-conversion --fir-to-llvm-ir="target=aarch64-unknown-linux-gnu" test.fir -o test.mlir

[TIFitis]

Added tests.

[TIFitis]

Updated check regex in test.

[TIFitis]

Rebased

[kiranchandramohan]

mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp

141–149

Wouldn't this Operation require a region conversion as well?

[TIFitis]

mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp

141–149

Yes, its present in line 189.

mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp

 };
 } // namespace
 
 void mlir::configureOpenMPToLLVMConversionLegality(
     ConversionTarget &target, LLVMTypeConverter &typeConverter) {
   target.addDynamicallyLegalOp<
       mlir::omp::CriticalOp, mlir::omp::ParallelOp, mlir::omp::WsLoopOp,
       mlir::omp::SimdLoopOp, mlir::omp::MasterOp, mlir::omp::SectionsOp,
-      mlir::omp::SingleOp, mlir::omp::TaskOp>([&](Operation *op) {
+      mlir::omp::SingleOp, mlir::omp::TaskOp, mlir::omp::DataOp>(
+      [&](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, mlir::omp::DataOp,
-                               mlir::omp::EnterDataOp, mlir::omp::ExitDataOp>(
+  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());
           });
   target.addDynamicallyLegalOp<mlir::omp::ReductionOp>([&](Operation *op) {
     return typeConverter.isLegal(op->getOperandTypes());
   });
 }
 
 void mlir::populateOpenMPToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                                   RewritePatternSet &patterns) {
   patterns.add<
-      ReductionOpConversion, RegionOpConversion<omp::CriticalOp>,
-      RegionOpConversion<omp::MasterOp>, ReductionOpConversion,
-      RegionOpConversion<omp::MasterOp>, RegionOpConversion<omp::ParallelOp>,
-      RegionOpConversion<omp::WsLoopOp>, RegionOpConversion<omp::SectionsOp>,
-      RegionOpConversion<omp::SimdLoopOp>, RegionOpConversion<omp::SingleOp>,
-      RegionOpConversion<omp::TaskOp>,
+      LegalizeDataOpForLLVMTranslation<omp::DataOp>,
+      LegalizeDataOpForLLVMTranslation<omp::EnterDataOp>,
+      LegalizeDataOpForLLVMTranslation<omp::ExitDataOp>, ReductionOpConversion,
+      RegionOpConversion<omp::CriticalOp>, RegionOpConversion<omp::MasterOp>,
+      ReductionOpConversion, RegionOpConversion<omp::MasterOp>,
+      RegionOpConversion<omp::ParallelOp>, RegionOpConversion<omp::WsLoopOp>,
+      RegionOpConversion<omp::SectionsOp>, RegionOpConversion<omp::SimdLoopOp>,
+      RegionOpConversion<omp::SingleOp>, RegionOpConversion<omp::TaskOp>,
+      RegionOpConversion<omp::DataOp>,
       RegionLessOpWithVarOperandsConversion<omp::AtomicReadOp>,
       RegionLessOpWithVarOperandsConversion<omp::AtomicWriteOp>,
       RegionLessOpWithVarOperandsConversion<omp::FlushOp>,
-      RegionLessOpWithVarOperandsConversion<omp::ThreadprivateOp>,
-      LegalizeDataOpForLLVMTranslation<omp::DataOp>,
-      LegalizeDataOpForLLVMTranslation<omp::EnterDataOp>,
-      LegalizeDataOpForLLVMTranslation<omp::ExitDataOp>>(converter);
+      RegionLessOpWithVarOperandsConversion<omp::ThreadprivateOp>>(converter);
 }
 
 namespace {
 struct ConvertOpenMPToLLVMPass
     : public impl::ConvertOpenMPToLLVMPassBase<ConvertOpenMPToLLVMPass> {
   using Base::Base;
 
   void runOnOperation() override;