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

[TIFitis]

Summary

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

Depends on D147247

[TIFitis]

Added previous patch dependency.

[kiranchandramohan]

The test is failing.

mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp

155–156

Looking back at this, not sure why we needed this pattern. Doesn't RegionLessOpConversion cover this case?

216

I think this is not required.

[TIFitis]

Addressed reviewer comments.
Fixed op conversion for target data ops.

[TIFitis]

In D147439#4244280, @kiranchandramohan wrote:

The test is failing.

The test was failing because of the missing dependency on target map clause support. On rerun it is passing.

Sorry, you are right about the conversion. I've updated the target op along with the target data ops to use the existing region conversion patters.

[TIFitis]

ping for reviews. Thanks :)

[kiranchandramohan]

LG.

flang/test/Fir/convert-to-llvm-openmp-and-fir.fir

 // CHECK:             llvm.store %[[VAL_13]], %[[VAL_3]] : !llvm.ptr<i32>
 // CHECK:             omp.terminator
 // CHECK:           }
 // CHECK:           llvm.return
 // CHECK:         }
 
 // -----
 
+func.func @_QPomp_target() {
+  %0 = fir.alloca !fir.array<512xi32> {bindc_name = "a", uniq_name = "_QFomp_targetEa"}
+  %c64_i32 = arith.constant 64 : i32
+  omp.target   thread_limit(%c64_i32 : i32) map((tofrom -> %0 : !fir.ref<!fir.array<512xi32>>)) {
+    %c10_i32 = arith.constant 10 : i32
+    %c1_i64 = arith.constant 1 : i64
+    %c1_i64_0 = arith.constant 1 : i64
+    %1 = arith.subi %c1_i64, %c1_i64_0 : i64
+    %2 = fir.coordinate_of %0, %1 : (!fir.ref<!fir.array<512xi32>>, i64) -> !fir.ref<i32>
+    fir.store %c10_i32 to %2 : !fir.ref<i32>
+    omp.terminator
+  }
+  return
+}
+
+// CHECK-LABEL:   llvm.func @_QPomp_target() {
+// CHECK:           %[[VAL_0:.*]] = llvm.mlir.constant(1 : i64) : i64
+// CHECK:           %[[VAL_1:.*]] = llvm.alloca %[[VAL_0]] x !llvm.array<512 x i32> {bindc_name = "a", in_type = !fir.array<512xi32>, operand_segment_sizes = array<i32: 0, 0>, uniq_name = "_QFomp_targetEa"} : (i64) -> !llvm.ptr<array<512 x i32>>
+// CHECK:           %[[VAL_2:.*]] = llvm.mlir.constant(64 : i32) : i32
+// CHECK:           omp.target   thread_limit(%[[VAL_2]] : i32) map((tofrom -> %[[VAL_1]] : !llvm.ptr<array<512 x i32>>)) {
+// CHECK:             %[[VAL_3:.*]] = llvm.mlir.constant(10 : i32) : i32
+// CHECK:             %[[VAL_4:.*]] = llvm.mlir.constant(1 : i64) : i64
+// CHECK:             %[[VAL_5:.*]] = llvm.mlir.constant(1 : i64) : i64
+// CHECK:             %[[VAL_6:.*]] = llvm.mlir.constant(0 : i64) : i64
+// CHECK:             %[[VAL_7:.*]] = llvm.getelementptr %[[VAL_1]][0, %[[VAL_6]]] : (!llvm.ptr<array<512 x i32>>, i64) -> !llvm.ptr<i32>
+// CHECK:             llvm.store %[[VAL_3]], %[[VAL_7]] : !llvm.ptr<i32>
+// CHECK:             omp.terminator
+// CHECK:           }
+// CHECK:           llvm.return
+// CHECK:         }
+
+// -----
+
 func.func @_QPsimdloop_with_nested_loop() {
   %0 = fir.alloca i32 {adapt.valuebyref}
   %1 = fir.alloca !fir.array<10xi32> {bindc_name = "a", uniq_name = "_QFsimdloop_with_nested_loopEa"}
   %2 = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFsimdloop_with_nested_loopEi"}
   %3 = fir.alloca i32 {bindc_name = "j", uniq_name = "_QFsimdloop_with_nested_loopEj"}
   %c1_i32 = arith.constant 1 : i32
   %c10_i32 = arith.constant 10 : i32
   %c1_i32_0 = arith.constant 1 : i32

 // CHECK:  ^bb0(%[[ARG_1:.*]]: i32, %[[ARG_2:.*]]: i32):
 // CHECK:    %[[ZERO_1:.*]] = llvm.mlir.constant(0 : i64) : i32
 // CHECK:    %[[ARGVAL_1:.*]] = llvm.icmp "ne" %[[ARG_1]], %[[ZERO_1]] : i32
 // CHECK:    %[[ZERO_2:.*]] = llvm.mlir.constant(0 : i64) : i32
 // CHECK:    %[[ARGVAL_2:.*]] = llvm.icmp "ne" %[[ARG_2]], %[[ZERO_2]] : i32
 // CHECK:    %[[RES:.*]] = llvm.icmp "eq" %[[ARGVAL_1]], %[[ARGVAL_2]] : i1
 // CHECK:    %[[RES_EXT:.*]] = llvm.zext %[[RES]] : i1 to i32
 // CHECK:    omp.yield(%[[RES_EXT]] : i32)
-// CHECK:  } 
+// CHECK:  }
 // CHECK-LABEL:  @_QPsimple_reduction
 // CHECK-SAME: %[[ARRAY_REF:.*]]: !llvm.ptr<array<100 x i32>>
 // CHECK:    %[[RED_ACCUMULATOR:.*]] = llvm.alloca %2 x i32 {bindc_name = "x", in_type = !fir.logical<4>, operand_segment_sizes = array<i32: 0, 0>, uniq_name = "_QFsimple_reductionEx"} : (i64) -> !llvm.ptr<i32>
 // CHECK:    omp.parallel   {
 // CHECK:      omp.wsloop   reduction(@[[EQV_REDUCTION]] -> %[[RED_ACCUMULATOR]] : !llvm.ptr<i32>) for
 // CHECK:        %[[ARRAY_ELEM_REF:.*]] = llvm.getelementptr %[[ARRAY_REF]][0, %{{.*}}] : (!llvm.ptr<array<100 x i32>>, i64) -> !llvm.ptr<i32>
 // CHECK:        %[[ARRAY_ELEM:.*]] = llvm.load %[[ARRAY_ELEM_REF]] : !llvm.ptr<i32>
 // CHECK:        omp.reduction %[[ARRAY_ELEM]], %[[RED_ACCUMULATOR]] : i32, !llvm.ptr<i32>
 // CHECK:        omp.yield
 // CHECK:      omp.terminator
 // CHECK:    llvm.return
 
 omp.reduction.declare @eqv_reduction : !fir.logical<4> init {
 ^bb0(%arg0: !fir.logical<4>):
   %true = arith.constant true
   %0 = fir.convert %true : (i1) -> !fir.logical<4>
   omp.yield(%0 : !fir.logical<4>)
 } combiner {
 ^bb0(%arg0: !fir.logical<4>, %arg1: !fir.logical<4>):
   %0 = fir.convert %arg0 : (!fir.logical<4>) -> i1
   %1 = fir.convert %arg1 : (!fir.logical<4>) -> i1
   %2 = arith.cmpi eq, %0, %1 : i1
   %3 = fir.convert %2 : (i1) -> !fir.logical<4>
   omp.yield(%3 : !fir.logical<4>)
-} 
+}
 func.func @_QPsimple_reduction(%arg0: !fir.ref<!fir.array<100x!fir.logical<4>>> {fir.bindc_name = "y"}) {
   %0 = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFsimple_reductionEi"}
   %1 = fir.alloca !fir.logical<4> {bindc_name = "x", uniq_name = "_QFsimple_reductionEx"}
   %true = arith.constant true
   %2 = fir.convert %true : (i1) -> !fir.logical<4>
   fir.store %2 to %1 : !fir.ref<!fir.logical<4>>
   omp.parallel   {
     %3 = fir.alloca i32 {adapt.valuebyref, pinned}

mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp

if (curOp.getAccumulator().getType().isa<MemRefType>()) {
       // TODO: Support memref type in variable operands
       return rewriter.notifyMatchFailure(curOp, "memref is not supported yet");
     }
     rewriter.replaceOpWithNewOp<omp::ReductionOp>(
         curOp, TypeRange(), adaptor.getOperands(), curOp->getAttrs());
     return success();
   }
 };
 
-template <typename Op>
-struct LegalizeDataOpForLLVMTranslation : public ConvertOpToLLVMPattern<Op> {
-  using ConvertOpToLLVMPattern<Op>::ConvertOpToLLVMPattern;
-  LogicalResult
-  matchAndRewrite(Op curOp, typename Op::Adaptor adaptor,
-                  ConversionPatternRewriter &rewriter) const override {
-    rewriter.replaceOpWithNewOp<Op>(curOp, TypeRange(), adaptor.getOperands(),
-                                    curOp.getOperation()->getAttrs());
-    return success();
-  }
-};
-
 struct ReductionDeclareOpConversion
     : public ConvertOpToLLVMPattern<omp::ReductionDeclareOp> {
   using ConvertOpToLLVMPattern<omp::ReductionDeclareOp>::ConvertOpToLLVMPattern;
   LogicalResult
   matchAndRewrite(omp::ReductionDeclareOp curOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto newOp = rewriter.create<omp::ReductionDeclareOp>(
         curOp.getLoc(), TypeRange(), curOp.getSymNameAttr(),
         TypeAttr::get(this->getTypeConverter()->convertType(
matchAndRewrite(omp::ReductionDeclareOp curOp, OpAdaptor adaptor,
     return success();
   }
 };
 } // namespace
 
 void mlir::configureOpenMPToLLVMConversionLegality(
     ConversionTarget &target, LLVMTypeConverter &typeConverter) {
   target.addDynamicallyLegalOp<
-      mlir::omp::AtomicUpdateOp, mlir::omp::CriticalOp, mlir::omp::DataOp,
-      mlir::omp::ParallelOp, mlir::omp::WsLoopOp, mlir::omp::SimdLoopOp,
-      mlir::omp::MasterOp, mlir::omp::SectionOp, mlir::omp::SectionsOp,
-      mlir::omp::SingleOp, mlir::omp::TaskOp>([&](Operation *op) {
+      mlir::omp::AtomicUpdateOp, mlir::omp::CriticalOp, mlir::omp::TargetOp,
+      mlir::omp::DataOp, mlir::omp::ParallelOp, mlir::omp::WsLoopOp,
+      mlir::omp::SimdLoopOp, mlir::omp::MasterOp, mlir::omp::SectionOp,
+      mlir::omp::SectionsOp, mlir::omp::SingleOp, mlir::omp::TaskOp>(
+      [&](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::YieldOp,
                                mlir::omp::EnterDataOp, mlir::omp::ExitDataOp>(
       [&](Operation *op) {
         return typeConverter.isLegal(op->getOperandTypes()) &&
                typeConverter.isLegal(op->getResultTypes());
       });
   target.addDynamicallyLegalOp<mlir::omp::ReductionOp>([&](Operation *op) {
     return typeConverter.isLegal(op->getOperandTypes());
   });
   target.addDynamicallyLegalOp<mlir::omp::ReductionDeclareOp>(
       [&](Operation *op) {
         return typeConverter.isLegal(&op->getRegion(0)) &&
                typeConverter.isLegal(&op->getRegion(1)) &&
                typeConverter.isLegal(&op->getRegion(2)) &&
                typeConverter.isLegal(op->getOperandTypes()) &&
                typeConverter.isLegal(op->getResultTypes());
       });
 }
 
 void mlir::populateOpenMPToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                                   RewritePatternSet &patterns) {
   patterns.add<
-      LegalizeDataOpForLLVMTranslation<omp::DataOp>,
-      LegalizeDataOpForLLVMTranslation<omp::EnterDataOp>,
-      LegalizeDataOpForLLVMTranslation<omp::ExitDataOp>, ReductionOpConversion,
-      ReductionDeclareOpConversion, RegionOpConversion<omp::CriticalOp>,
-      RegionOpConversion<omp::MasterOp>, ReductionOpConversion,
-      RegionOpConversion<omp::ParallelOp>, RegionOpConversion<omp::WsLoopOp>,
-      RegionOpConversion<omp::SectionsOp>, RegionOpConversion<omp::SectionOp>,
-      RegionOpConversion<omp::SimdLoopOp>, RegionOpConversion<omp::SingleOp>,
-      RegionOpConversion<omp::TaskOp>, RegionOpConversion<omp::DataOp>,
+      ReductionOpConversion, ReductionDeclareOpConversion,
+      RegionOpConversion<omp::CriticalOp>, RegionOpConversion<omp::MasterOp>,
+      ReductionOpConversion, RegionOpConversion<omp::ParallelOp>,
+      RegionOpConversion<omp::WsLoopOp>, RegionOpConversion<omp::SectionsOp>,
+      RegionOpConversion<omp::SectionOp>, RegionOpConversion<omp::SimdLoopOp>,
+      RegionOpConversion<omp::SingleOp>, RegionOpConversion<omp::TaskOp>,
+      RegionOpConversion<omp::DataOp>, RegionOpConversion<omp::TargetOp>,
       RegionLessOpWithVarOperandsConversion<omp::AtomicReadOp>,
       RegionLessOpWithVarOperandsConversion<omp::AtomicWriteOp>,
       RegionOpWithVarOperandsConversion<omp::AtomicUpdateOp>,
       RegionLessOpWithVarOperandsConversion<omp::FlushOp>,
       RegionLessOpWithVarOperandsConversion<omp::ThreadprivateOp>,
-      RegionLessOpConversion<omp::YieldOp>>(converter);
+      RegionLessOpConversion<omp::YieldOp>,
+      RegionLessOpConversion<omp::EnterDataOp>,
+      RegionLessOpConversion<omp::ExitDataOp>>(converter);
 }
 
 namespace {
 struct ConvertOpenMPToLLVMPass
     : public impl::ConvertOpenMPToLLVMPassBase<ConvertOpenMPToLLVMPass> {
   using Base::Base;
 
   void runOnOperation() override;

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

omp.target_data   map((tofrom -> %a : !llvm.ptr<array<1024 x i32>>)) {
     llvm.store %1, %i : !llvm.ptr<i32>
     omp.terminator
   }
   llvm.return
 }
 
 // -----
 
+// CHECK-LABEL:   llvm.func @_QPomp_target(
+// CHECK:                             %[[ARG_0:.*]]: !llvm.ptr<array<1024 x i32>>,
+// CHECK:                             %[[ARG_1:.*]]: !llvm.ptr<i32>) {
+// CHECK:           %[[VAL_0:.*]] = llvm.mlir.constant(64 : i32) : i32
+// CHECK:           omp.target   thread_limit(%[[VAL_0]] : i32) map((tofrom -> %[[ARG_0]] : !llvm.ptr<array<1024 x i32>>)) {
+// CHECK:             %[[VAL_1:.*]] = llvm.mlir.constant(10 : i32) : i32
+// CHECK:             llvm.store %[[VAL_1]], %[[ARG_1]] : !llvm.ptr<i32>
+// CHECK:             omp.terminator
+// CHECK:           }
+// CHECK:           llvm.return
+// CHECK:         }
+
+llvm.func @_QPomp_target(%a : !llvm.ptr<array<1024 x i32>>, %i : !llvm.ptr<i32>) {
+  %0 = llvm.mlir.constant(64 : i32) : i32
+  omp.target   thread_limit(%0 : i32) map((tofrom -> %a : !llvm.ptr<array<1024 x i32>>)) {
+    %1 = llvm.mlir.constant(10 : i32) : i32
+    llvm.store %1, %i : !llvm.ptr<i32>
+    omp.terminator
+  }
+  llvm.return
+}
+
+// -----
+
 // CHECK-LABEL: @_QPsb
 // CHECK: omp.sections
 // CHECK: omp.section
 // CHECK: llvm.br
 // CHECK: llvm.icmp
 // CHECK: llvm.cond_br
 // CHECK: llvm.br
 // CHECK: omp.terminator