[flang][openacc] Add lowering support for multi-dimensional arrays reduction

[clementval]

Summary

Lower multi-dimensional arrays reduction for add and mul operator.

Depends on D153448

[razvanlupusoru]

Looks great also!

[jeanPerier]

flang/test/Lower/OpenACC/acc-reduction.f90

128–136

Probably not a correctness issue, but the loop nest is reversed from a memory access point of view (Fortran is column major).

If you want, HLFIR provides a tool to generate loop nest according to the list of extents:
https://github.com/llvm/llvm-project/blob/74adc3e0ebfb42a48f02d8d7094d6848a37a21f5/flang/include/flang/Optimizer/Builder/HLFIRTools.h#L382 (however, this is a one based loop nest, so not exactly what you have here).

[clementval]

flang/test/Lower/OpenACC/acc-reduction.f90

128–136

Good point! I'll update that.

flang/lib/Lower/OpenACC.cpp

 
 static mlir::Value genCombiner(fir::FirOpBuilder &builder, mlir::Location loc,
                                mlir::acc::ReductionOperator op, mlir::Type ty,
                                mlir::Value value1, mlir::Value value2) {
 
   // Handle combiner on arrays.
   if (auto refTy = mlir::dyn_cast<fir::ReferenceType>(ty)) {
     if (auto seqTy = mlir::dyn_cast<fir::SequenceType>(refTy.getEleTy())) {
-      if (seqTy.getShape().size() > 1)
-        TODO(loc, "OpenACC reduction on array with more than one dimension");
       if (seqTy.hasDynamicExtents())
         TODO(loc, "OpenACC reduction on array with dynamic extents");
       mlir::Type idxTy = builder.getIndexType();
       mlir::Type refTy = fir::ReferenceType::get(seqTy.getEleTy());
+
+      llvm::SmallVector<fir::DoLoopOp> loops;
+      llvm::SmallVector<mlir::Value> ivs;
+      for (auto ext : seqTy.getShape()) {
         auto lb = builder.create<mlir::arith::ConstantOp>(
             loc, idxTy, builder.getIntegerAttr(idxTy, 0));
         auto ub = builder.create<mlir::arith::ConstantOp>(
-          loc, idxTy, builder.getIntegerAttr(idxTy, seqTy.getShape()[0] - 1));
+            loc, idxTy, builder.getIntegerAttr(idxTy, ext - 1));
         auto step = builder.create<mlir::arith::ConstantOp>(
             loc, idxTy, builder.getIntegerAttr(idxTy, 1));
         auto loop = builder.create<fir::DoLoopOp>(loc, lb, ub, step,
                                                   /*unordered=*/false);
         builder.setInsertionPointToStart(loop.getBody());
-      auto addr1 = builder.create<fir::CoordinateOp>(
-          loc, refTy, value1, mlir::ValueRange{loop.getInductionVar()});
-      auto addr2 = builder.create<fir::CoordinateOp>(
-          loc, refTy, value2, mlir::ValueRange{loop.getInductionVar()});
+        loops.push_back(loop);
+        ivs.push_back(loop.getInductionVar());
+      }
+      auto addr1 = builder.create<fir::CoordinateOp>(loc, refTy, value1, ivs);
+      auto addr2 = builder.create<fir::CoordinateOp>(loc, refTy, value2, ivs);
       auto load1 = builder.create<fir::LoadOp>(loc, addr1);
       auto load2 = builder.create<fir::LoadOp>(loc, addr2);
       auto combined =
           genCombiner(builder, loc, op, seqTy.getEleTy(), load1, load2);
       builder.create<fir::StoreOp>(loc, combined, addr1);
-      builder.setInsertionPointAfter(loop);
+      builder.setInsertionPointAfter(loops[0]);
       return value1;
     }
   }
 
   if (op == mlir::acc::ReductionOperator::AccAdd) {
     if (ty.isIntOrIndex())
       return builder.create<mlir::arith::AddIOp>(loc, value1, value2);
     if (mlir::isa<mlir::FloatType>(ty))

flang/test/Lower/OpenACC/acc-reduction.f90

 ! CHECK:   %[[INIT:.*]] = arith.constant 0.000000e+00 : f32
 ! CHECK:   acc.yield %[[INIT]] : f32
 ! CHECK: } combiner {
 ! CHECK: ^bb0(%[[ARG0:.*]]: f32, %[[ARG1:.*]]: f32):
 ! CHECK:   %[[COMBINED:.*]] = arith.addf %[[ARG0]], %[[ARG1]] {{.*}} : f32
 ! CHECK:   acc.yield %[[COMBINED]] : f32
 ! CHECK: }
 
+! CHECK-LABEL: acc.reduction.recipe @reduction_add_ref_100x10x2xi32 : !fir.ref<!fir.array<100x10x2xi32>> reduction_operator <add> init {
+! CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.array<100x10x2xi32>>):
+! CHECK:   %[[CST:.*]] = arith.constant dense<0> : vector<100x10x2xi32>
+! CHECK:   acc.yield %[[CST]] : vector<100x10x2xi32>
+! CHECK: } combiner {
+! CHECK: ^bb0(%[[ARG0:.*]]: !fir.ref<!fir.array<100x10x2xi32>>, %[[ARG1:.*]]: !fir.ref<!fir.array<100x10x2xi32>>):
+! CHECK:   %[[LB0:.*]] = arith.constant 0 : index
+! CHECK:   %[[UB0:.*]] = arith.constant 99 : index
+! CHECK:   %[[STEP0:.*]] = arith.constant 1 : index
+! CHECK:   fir.do_loop %[[IV0:.*]] = %[[LB0]] to %[[UB0]] step %[[STEP0]] {
+! CHECK:     %[[LB1:.*]] = arith.constant 0 : index
+! CHECK:     %[[UB1:.*]] = arith.constant 9 : index
+! CHECK:     %[[STEP1:.*]] = arith.constant 1 : index
+! CHECK:     fir.do_loop %[[IV1:.*]] = %[[LB1]] to %[[UB1]] step %[[STEP1]] {
+! CHECK:       %[[LB2:.*]] = arith.constant 0 : index
+! CHECK:       %[[UB2:.*]] = arith.constant 1 : index
+! CHECK:       %[[STEP2:.*]] = arith.constant 1 : index
+! CHECK:       fir.do_loop %[[IV2:.*]] = %[[LB2]] to %[[UB2]] step %[[STEP2]] {
+! CHECK:         %[[COORD1:.*]] = fir.coordinate_of %[[ARG0]], %[[IV0]], %[[IV1]], %[[IV2]] : (!fir.ref<!fir.array<100x10x2xi32>>, index, index, index) -> !fir.ref<i32>
+! CHECK:         %[[COORD2:.*]] = fir.coordinate_of %[[ARG1]], %[[IV0]], %[[IV1]], %[[IV2]] : (!fir.ref<!fir.array<100x10x2xi32>>, index, index, index) -> !fir.ref<i32>
+! CHECK:         %[[LOAD1:.*]] = fir.load %[[COORD1]] : !fir.ref<i32>
+! CHECK:         %[[LOAD2:.*]] = fir.load %[[COORD2]] : !fir.ref<i32>
+! CHECK:         %[[COMBINED:.*]] = arith.addi %[[LOAD1]], %[[LOAD2]] : i32
+! CHECK:         fir.store %[[COMBINED]] to %[[COORD1]] : !fir.ref<i32>
+! CHECK:       }
+! CHECK:     }
+! CHECK:   }
+! CHECK:   acc.yield %[[ARG0]] : !fir.ref<!fir.array<100x10x2xi32>>
+! CHECK: }
+
+! CHECK-LABEL: acc.reduction.recipe @reduction_add_ref_100x10xi32 : !fir.ref<!fir.array<100x10xi32>> reduction_operator <add> init {
+! CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.array<100x10xi32>>):
+! CHECK: %[[CST:.*]] = arith.constant dense<0> : vector<100x10xi32>
+! CHECK: acc.yield %[[CST]] : vector<100x10xi32>
+! CHECK: } combiner {
+! CHECK: ^bb0(%[[ARG0:.*]]: !fir.ref<!fir.array<100x10xi32>>, %[[ARG1:.*]]: !fir.ref<!fir.array<100x10xi32>>):
+! CHECK:   %[[LB0:.*]] = arith.constant 0 : index
+! CHECK:   %[[UB0:.*]] = arith.constant 99 : index
+! CHECK:   %[[STEP0:.*]] = arith.constant 1 : index
+! CHECK:   fir.do_loop %[[IV0:.*]] = %[[LB0]] to %[[UB0]] step %[[STEP0]] {
+! CHECK:     %[[LB1:.*]] = arith.constant 0 : index
+! CHECK:     %[[UB1:.*]] = arith.constant 9 : index
+! CHECK:     %[[STEP1:.*]] = arith.constant 1 : index
+! CHECK:     fir.do_loop %[[IV1:.*]] = %[[LB1]] to %[[UB1]] step %[[STEP1]] {
+! CHECK:       %[[COORD1:.*]] = fir.coordinate_of %[[ARG0]], %[[IV0]], %[[IV1]] : (!fir.ref<!fir.array<100x10xi32>>, index, index) -> !fir.ref<i32>
+! CHECK:       %[[COORD2:.*]] = fir.coordinate_of %[[ARG1]], %[[IV0]], %[[IV1]] : (!fir.ref<!fir.array<100x10xi32>>, index, index) -> !fir.ref<i32>
+! CHECK:       %[[LOAD1]] = fir.load %[[COORD1]] : !fir.ref<i32>
+! CHECK:       %[[LOAD2]] = fir.load %[[COORD2]] : !fir.ref<i32>
+! CHECK:       %[[COMBINED:.*]] = arith.addi %[[LOAD1]], %[[LOAD2]] : i32
+! CHECK:       fir.store %[[COMBINED]] to %[[COORD1]] : !fir.ref<i32>
+! CHECK:     }
+! CHECK:   }
+! CHECK:   acc.yield %[[ARG0]] : !fir.ref<!fir.array<100x10xi32>>
+! CHECK: }
+
 ! CHECK-LABEL: acc.reduction.recipe @reduction_add_ref_100xi32 : !fir.ref<!fir.array<100xi32>> reduction_operator <add> init {
 ! CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.array<100xi32>>):
 ! CHECK:   %[[CST:.*]] = arith.constant dense<0> : vector<100xi32>
 ! CHECK:   acc.yield %[[CST]] : vector<100xi32>
 ! CHECK: } combiner {
 ! CHECK: ^bb0(%[[ARG0:.*]]: !fir.ref<!fir.array<100xi32>>, %[[ARG1:.*]]: !fir.ref<!fir.array<100xi32>>):
 ! CHECK:   %[[LB:.*]] = arith.constant 0 : index
 ! CHECK:   %[[UB:.*]] = arith.constant 99 : index
do i = 1, 100
     b(i) = b(i) + a(i)
   end do
 end subroutine
 
 ! CHECK-LABEL: func.func @_QPacc_reduction_add_int_array_1d(
 ! CHECK-SAME:  %{{.*}}: !fir.ref<!fir.array<100xi32>> {fir.bindc_name = "a"}, %[[B:.*]]: !fir.ref<!fir.array<100xi32>> {fir.bindc_name = "b"})
 ! CHECK:       acc.loop reduction(@reduction_add_ref_100xi32 -> %[[B]] : !fir.ref<!fir.array<100xi32>>)
 
+subroutine acc_reduction_add_int_array_2d(a, b)
+  integer :: a(100, 10), b(100, 10)
+  integer :: i, j
+
+  !$acc loop collapse(2) reduction(+:b)
+  do i = 1, 100
+    do j = 1, 10
+      b(i, j) = b(i, j) + a(i, j)
+    end do
+  end do
+end subroutine
+
+! CHECK-LABEL: func.func @_QPacc_reduction_add_int_array_2d(
+! CHECK-SAME:  %[[ARG0:.*]]: !fir.ref<!fir.array<100x10xi32>> {fir.bindc_name = "a"}, %[[ARG1:.*]]: !fir.ref<!fir.array<100x10xi32>> {fir.bindc_name = "b"}) {
+! CHECK:       acc.loop reduction(@reduction_add_ref_100x10xi32 -> %[[ARG1]] : !fir.ref<!fir.array<100x10xi32>>) {
+! CHECK: } attributes {collapse = 2 : i64}
+
+subroutine acc_reduction_add_int_array_3d(a, b)
+  integer :: a(100, 10, 2), b(100, 10, 2)
+  integer :: i, j, k
+
+  !$acc loop collapse(3) reduction(+:b)
+  do i = 1, 100
+    do j = 1, 10
+      do k = 1, 2
+        b(i, j, k) = b(i, j, k) + a(i, j, k)
+      end do
+    end do
+  end do
+end subroutine
+
+! CHECK-LABEL: func.func @_QPacc_reduction_add_int_array_3d(
+! CHECK-SAME: %{{.*}}: !fir.ref<!fir.array<100x10x2xi32>> {fir.bindc_name = "a"}, %[[ARG1:.*]]: !fir.ref<!fir.array<100x10x2xi32>> {fir.bindc_name = "b"})
+! CHECK: acc.loop reduction(@reduction_add_ref_100x10x2xi32 -> %[[ARG1]] : !fir.ref<!fir.array<100x10x2xi32>>)
+! CHECK: } attributes {collapse = 3 : i64}
+
 subroutine acc_reduction_add_float(a, b)
   real :: a(100), b
   integer :: i
 
   !$acc loop reduction(+:b)
   do i = 1, 100
     b = b + a(i)
   end do