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

[TIFitis]

Summary

This enables conversion of OpenMP Target Data, Enter Data and Exit Data from FIR Dialect to LLVM IR Dialect.

[kiranchandramohan]

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

157–161

%a, %b etc should be references, I guess. Can you switch to using the !llvm.ptr type for this?

[TIFitis]

Addressed reviewer comments.

[clementval]

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

153

You should not have the regex otherwise you redefine the variable instead of checking it with the captured one. Same below.

[TIFitis]

Removed regex from //CHECKs.

[kiranchandramohan]

LGTM.

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

 // CHECK:   llvm.store %[[I1]], %[[ARR_I_REF]] : !llvm.ptr<i32>
 // CHECK: omp.yield
 // CHECK: }
 // CHECK: omp.terminator
 // CHECK: }
 // CHECK: llvm.return
 // CHECK: }
 
+func.func @_QPomp_target_data() {
+  %0 = fir.alloca !fir.array<1024xi32> {bindc_name = "a", uniq_name = "_QFomp_target_dataEa"}
+  %1 = fir.alloca !fir.array<1024xi32> {bindc_name = "b", uniq_name = "_QFomp_target_dataEb"}
+  %2 = fir.alloca !fir.array<1024xi32> {bindc_name = "c", uniq_name = "_QFomp_target_dataEc"}
+  %3 = fir.alloca !fir.array<1024xi32> {bindc_name = "d", uniq_name = "_QFomp_target_dataEd"}
+  omp.target_enter_data   map((to -> %0 : !fir.ref<!fir.array<1024xi32>>), (to -> %1 : !fir.ref<!fir.array<1024xi32>>), (always, alloc -> %2 : !fir.ref<!fir.array<1024xi32>>))
+  omp.target_exit_data   map((from -> %0 : !fir.ref<!fir.array<1024xi32>>), (from -> %1 : !fir.ref<!fir.array<1024xi32>>), (release -> %2 : !fir.ref<!fir.array<1024xi32>>), (always, delete -> %3 : !fir.ref<!fir.array<1024xi32>>))
+  return
+}
+
+// CHECK-LABEL:   llvm.func @_QPomp_target_data() {
+// CHECK:           %[[VAL_0:.*]] = llvm.mlir.constant(1 : i64) : i64
+// CHECK:           %[[VAL_1:.*]] = llvm.alloca %[[VAL_0]] x !llvm.array<1024 x i32> {bindc_name = "a", in_type = !fir.array<1024xi32>, operand_segment_sizes = array<i32: 0, 0>, uniq_name = "_QFomp_target_dataEa"} : (i64) -> !llvm.ptr<array<1024 x i32>>
+// CHECK:           %[[VAL_2:.*]] = llvm.mlir.constant(1 : i64) : i64
+// CHECK:           %[[VAL_3:.*]] = llvm.alloca %[[VAL_2]] x !llvm.array<1024 x i32> {bindc_name = "b", in_type = !fir.array<1024xi32>, operand_segment_sizes = array<i32: 0, 0>, uniq_name = "_QFomp_target_dataEb"} : (i64) -> !llvm.ptr<array<1024 x i32>>
+// CHECK:           %[[VAL_4:.*]] = llvm.mlir.constant(1 : i64) : i64
+// CHECK:           %[[VAL_5:.*]] = llvm.alloca %[[VAL_4]] x !llvm.array<1024 x i32> {bindc_name = "c", in_type = !fir.array<1024xi32>, operand_segment_sizes = array<i32: 0, 0>, uniq_name = "_QFomp_target_dataEc"} : (i64) -> !llvm.ptr<array<1024 x i32>>
+// CHECK:           %[[VAL_6:.*]] = llvm.mlir.constant(1 : i64) : i64
+// CHECK:           %[[VAL_7:.*]] = llvm.alloca %[[VAL_6]] x !llvm.array<1024 x i32> {bindc_name = "d", in_type = !fir.array<1024xi32>, operand_segment_sizes = array<i32: 0, 0>, uniq_name = "_QFomp_target_dataEd"} : (i64) -> !llvm.ptr<array<1024 x i32>>
+// CHECK:           omp.target_enter_data   map((to -> %[[VAL_1]] : !llvm.ptr<array<1024 x i32>>), (to -> %[[VAL_3]] : !llvm.ptr<array<1024 x i32>>), (always, alloc -> %[[VAL_5]] : !llvm.ptr<array<1024 x i32>>))
+// CHECK:           omp.target_exit_data   map((from -> %[[VAL_1]] : !llvm.ptr<array<1024 x i32>>), (from -> %[[VAL_3]] : !llvm.ptr<array<1024 x i32>>), (release -> %[[VAL_5]] : !llvm.ptr<array<1024 x i32>>), (always, delete -> %[[VAL_7]] : !llvm.ptr<array<1024 x i32>>))
+// 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

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();
+  }
+};
 } // 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>([&](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());
-          });
+  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::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>,
       RegionLessOpWithVarOperandsConversion<omp::AtomicReadOp>,
       RegionLessOpWithVarOperandsConversion<omp::AtomicWriteOp>,
       RegionLessOpWithVarOperandsConversion<omp::FlushOp>,
-      RegionLessOpWithVarOperandsConversion<omp::ThreadprivateOp>>(converter);
+      RegionLessOpWithVarOperandsConversion<omp::ThreadprivateOp>,
+      LegalizeDataOpForLLVMTranslation<omp::DataOp>,
+      LegalizeDataOpForLLVMTranslation<omp::EnterDataOp>,
+      LegalizeDataOpForLLVMTranslation<omp::ExitDataOp>>(converter);
 }
 
 namespace {
 struct ConvertOpenMPToLLVMPass
     : public impl::ConvertOpenMPToLLVMBase<ConvertOpenMPToLLVMPass> {
   void runOnOperation() override;
 };
 } // namespace

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

func.func @simdloop_block_arg(%val : i32, %ub : i32, %i : index) {
   ^bb2:
     %2 = arith.addi %0, %i : index
     cf.br ^bb1(%2 : index)
   ^bb3:
     omp.yield
   }
   return
 }
+
+// -----
+
+// CHECK-LABEL: @_QPomp_target_data
+// CHECK: (%[[ARG0:.*]]: !llvm.ptr<i32>, %[[ARG1:.*]]: !llvm.ptr<i32>, %[[ARG2:.*]]: !llvm.ptr<i32>, %[[ARG3:.*]]: !llvm.ptr<i32>)
+// CHECK:         omp.target_enter_data   map((to -> %[[ARG0]] : !llvm.ptr<i32>), (to -> %[[ARG1]] : !llvm.ptr<i32>), (always, alloc -> %[[ARG2]] : !llvm.ptr<i32>))
+// CHECK:         omp.target_exit_data   map((from -> %[[ARG0]] : !llvm.ptr<i32>), (from -> %[[ARG1]] : !llvm.ptr<i32>), (release -> %[[ARG2]] : !llvm.ptr<i32>), (always, delete -> %[[ARG3]] : !llvm.ptr<i32>))
+// CHECK:         llvm.return
+
+llvm.func @_QPomp_target_data(%a : !llvm.ptr<i32>, %b : !llvm.ptr<i32>, %c : !llvm.ptr<i32>, %d : !llvm.ptr<i32>) {
+  omp.target_enter_data   map((to -> %a : !llvm.ptr<i32>), (to -> %b : !llvm.ptr<i32>), (always, alloc -> %c : !llvm.ptr<i32>))
+  omp.target_exit_data   map((from -> %a : !llvm.ptr<i32>), (from -> %b : !llvm.ptr<i32>), (release -> %c : !llvm.ptr<i32>), (always, delete -> %d : !llvm.ptr<i32>))
+  llvm.return
+}