[flang][openacc] Add lowering support for 1d array reduction for add/mul operator

[clementval]

Summary

Lower 1d array reduction for add and mul operator. Multi-dimensional arrays and
other operator will follow.

[razvanlupusoru]

Looks great! Thank you!

[clementval]

Fix test

flang/lib/Lower/OpenACC.cpp

getReductionOperator(const Fortran::parser::AccReductionOperator &op) {
   case Fortran::parser::AccReductionOperator::Operator::Eqv:
     return mlir::acc::ReductionOperator::AccEqv;
   case Fortran::parser::AccReductionOperator::Operator::Neqv:
     return mlir::acc::ReductionOperator::AccNeqv;
   }
   llvm_unreachable("unexpected reduction operator");
 }
 
+/// Get the correct DenseElementsAttr attribute for the given init value.
+/// The verifier on the DenseElementsAttr is strict about the init value passed
+/// to it so it must matched the type.
+static mlir::DenseElementsAttr getDenseAttr(mlir::ShapedType shTy,
+                                            int64_t value) {
+  if (shTy.getElementType().isIntOrIndex()) {
+    if (auto intTy = mlir::dyn_cast<mlir::IntegerType>(shTy.getElementType())) {
+      if (intTy.getIntOrFloatBitWidth() == 8)
+        return mlir::DenseElementsAttr::get(shTy, static_cast<int8_t>(value));
+      if (intTy.getIntOrFloatBitWidth() == 16)
+        return mlir::DenseElementsAttr::get(shTy, static_cast<int16_t>(value));
+      if (intTy.getIntOrFloatBitWidth() == 32)
+        return mlir::DenseElementsAttr::get(shTy, static_cast<int32_t>(value));
+      if (intTy.getIntOrFloatBitWidth() == 64)
+        return mlir::DenseElementsAttr::get(shTy, value);
+    }
+  }
+
+  if (mlir::isa<mlir::FloatType>(shTy.getElementType())) {
+    if (auto intTy = mlir::dyn_cast<mlir::FloatType>(shTy.getElementType())) {
+      if (intTy.getIntOrFloatBitWidth() == 16)
+        return mlir::DenseElementsAttr::get(shTy, static_cast<float>(value));
+      if (intTy.getIntOrFloatBitWidth() == 32)
+        return mlir::DenseElementsAttr::get(shTy, static_cast<float>(value));
+      if (intTy.getIntOrFloatBitWidth() == 64)
+        return mlir::DenseElementsAttr::get(shTy, static_cast<double>(value));
+      if (intTy.getIntOrFloatBitWidth() == 128)
+        return mlir::DenseElementsAttr::get(shTy,
+                                            static_cast<long double>(value));
+    }
+  }
+
+  llvm_unreachable("unsupported dense attribute type");
+}
+
 static mlir::Value genReductionInitValue(fir::FirOpBuilder &builder,
                                          mlir::Location loc, mlir::Type ty,
                                          mlir::acc::ReductionOperator op) {
   if (op != mlir::acc::ReductionOperator::AccAdd &&
       op != mlir::acc::ReductionOperator::AccMul &&
       op != mlir::acc::ReductionOperator::AccMin &&
       op != mlir::acc::ReductionOperator::AccMax)
     TODO(loc, "reduction operator");
if (op == mlir::acc::ReductionOperator::AccMin) {
     // 1 for *
     int64_t initValue = op == mlir::acc::ReductionOperator::AccMul ? 1 : 0;
     if (ty.isIntOrIndex())
       return builder.create<mlir::arith::ConstantOp>(
           loc, ty, builder.getIntegerAttr(ty, initValue));
     if (mlir::isa<mlir::FloatType>(ty))
       return builder.create<mlir::arith::ConstantOp>(
           loc, ty, builder.getFloatAttr(ty, initValue));
+    if (auto refTy = mlir::dyn_cast<fir::ReferenceType>(ty)) {
+      if (auto seqTy = mlir::dyn_cast<fir::SequenceType>(refTy.getEleTy())) {
+        mlir::Type vecType =
+            mlir::VectorType::get(seqTy.getShape(), seqTy.getEleTy());
+        mlir::DenseElementsAttr denseAttr =
+            getDenseAttr(vecType.cast<mlir::ShapedType>(), initValue);
+        return builder.create<mlir::arith::ConstantOp>(loc, vecType, denseAttr);
+      }
+    }
   }
 
   TODO(loc, "reduction type");
 }
 
 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());
+      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));
+      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()});
+      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);
+      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))
       return builder.create<mlir::arith::AddFOp>(loc, value1, value2);
     TODO(loc, "reduction add type");
   }
 
genReductions(const Fortran::parser::AccObjectListWithReduction &objectList,
   for (const auto &accObject : objects.v) {
     llvm::SmallVector<mlir::Value> bounds;
     std::stringstream asFortran;
     mlir::Location operandLocation = genOperandLocation(converter, accObject);
     mlir::Value baseAddr = gatherDataOperandAddrAndBounds(
         converter, builder, semanticsContext, stmtCtx, accObject,
         operandLocation, asFortran, bounds);
 
-    if (!fir::isa_trivial(fir::unwrapRefType(baseAddr.getType())))
+    mlir::Type reductionTy = fir::unwrapRefType(baseAddr.getType());
+    if (auto seqTy = mlir::dyn_cast<fir::SequenceType>(reductionTy))
+      reductionTy = seqTy.getEleTy();
+
+    if (!fir::isa_trivial(reductionTy))
       TODO(operandLocation, "reduction with unsupported type");
 
     mlir::Type ty = fir::unwrapRefType(baseAddr.getType());
+    if (!fir::isa_trivial(ty))
+      ty = baseAddr.getType();
     std::string recipeName = fir::getTypeAsString(
         ty, converter.getKindMap(),
         ("reduction_" + stringifyReductionOperator(mlirOp)).str());
     mlir::acc::ReductionRecipeOp recipe =
         Fortran::lower::createOrGetReductionRecipe(builder, recipeName,
                                                    operandLocation, ty, mlirOp);
     reductionRecipes.push_back(mlir::SymbolRefAttr::get(
         builder.getContext(), recipe.getSymName().str()));

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

 ! CHECK:   %[[INIT:.*]] = arith.constant 1.000000e+00 : f32
 ! CHECK:   acc.yield %[[INIT]] : f32
 ! CHECK: } combiner {
 ! CHECK: ^bb0(%[[ARG0:.*]]: f32, %[[ARG1:.*]]: f32):
 ! CHECK:   %[[COMBINED:.*]] = arith.mulf %[[ARG0]], %[[ARG1]] {{.*}} : f32
 ! CHECK:   acc.yield %[[COMBINED]] : f32
 ! CHECK: }
 
+! CHECK-LABEL: acc.reduction.recipe @reduction_mul_ref_100xi32 : !fir.ref<!fir.array<100xi32>> reduction_operator <mul> init {
+! CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.array<100xi32>>):
+! CHECK:   %[[CST:.*]] = arith.constant dense<1> : 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
+! CHECK:   %[[STEP:.*]] = arith.constant 1 : index
+! CHECK:   fir.do_loop %[[IV:.*]] = %[[LB]] to %[[UB]] step %[[STEP]] {
+! CHECK:     %[[COORD1:.*]] = fir.coordinate_of %[[ARG0]], %[[IV]] : (!fir.ref<!fir.array<100xi32>>, index) -> !fir.ref<i32>
+! CHECK:     %[[COORD2:.*]] = fir.coordinate_of %[[ARG1]], %[[IV]] : (!fir.ref<!fir.array<100xi32>>, index) -> !fir.ref<i32>
+! CHECK:     %[[LOAD1:.*]] = fir.load %[[COORD1]] : !fir.ref<i32>
+! CHECK:     %[[LOAD2:.*]] = fir.load %[[COORD2]] : !fir.ref<i32>
+! CHECK:     %[[COMBINED:.*]] = arith.muli %[[LOAD1]], %[[LOAD2]] : i32
+! CHECK:     fir.store %[[COMBINED]] to %[[COORD1]] : !fir.ref<i32>
+! CHECK:   }
+! CHECK:   acc.yield %[[ARG0]] : !fir.ref<!fir.array<100xi32>>
+! CHECK: }
+
 ! CHECK-LABEL: acc.reduction.recipe @reduction_mul_i32 : i32 reduction_operator <mul> init {
 ! CHECK: ^bb0(%{{.*}}: i32):
 ! CHECK:   %[[INIT:.*]] = arith.constant 1 : i32
 ! CHECK:   acc.yield %[[INIT]] : i32
 ! CHECK: } combiner {
 ! CHECK: ^bb0(%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32):
 ! CHECK:   %[[COMBINED:.*]] = arith.muli %[[ARG0]], %[[ARG1]] : i32
 ! CHECK:   acc.yield %[[COMBINED]] : i32
 ! CHECK: }
 
+! CHECK-LABEL: acc.reduction.recipe @reduction_add_ref_100xf32 : !fir.ref<!fir.array<100xf32>> reduction_operator <add> init {
+! CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.array<100xf32>>):
+! CHECK:   %[[CST:.*]] = arith.constant dense<0.000000e+00> : vector<100xf32>
+! CHECK:   acc.yield %[[CST]] : vector<100xf32>
+! CHECK: } combiner {
+! CHECK: ^bb0(%[[ARG0:.*]]: !fir.ref<!fir.array<100xf32>>, %[[ARG1:.*]]: !fir.ref<!fir.array<100xf32>>):
+! CHECK:   %[[LB:.*]] = arith.constant 0 : index
+! CHECK:   %[[UB:.*]] = arith.constant 99 : index
+! CHECK:   %[[STEP:.*]] = arith.constant 1 : index
+! CHECK:   fir.do_loop %[[IV:.*]] = %[[LB]] to %[[UB]] step %[[STEP]] {
+! CHECK:   %[[COORD1:.*]] = fir.coordinate_of %[[ARG0]], %[[IV]] : (!fir.ref<!fir.array<100xf32>>, index) -> !fir.ref<f32>
+! CHECK:   %[[COORD2:.*]] = fir.coordinate_of %[[ARG1]], %[[IV]] : (!fir.ref<!fir.array<100xf32>>, index) -> !fir.ref<f32>
+! CHECK:   %[[LOAD1:.*]] = fir.load %[[COORD1]] : !fir.ref<f32>
+! CHECK:   %[[LOAD2:.*]] = fir.load %[[COORD2]] : !fir.ref<f32>
+! CHECK:   %[[COMBINED:.*]] = arith.addf %[[LOAD1]], %[[LOAD2]] fastmath<contract> : f32
+! CHECK:   fir.store %[[COMBINED]] to %[[COORD1]] : !fir.ref<f32>
+! CHECK:   }
+! CHECK:   acc.yield %[[ARG0]] : !fir.ref<!fir.array<100xf32>>
+! CHECK: }
+
 ! CHECK-LABEL: acc.reduction.recipe @reduction_add_f32 : f32 reduction_operator <add> init {
 ! CHECK: ^bb0(%{{.*}}: f32):
 ! 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_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
+! CHECK:   %[[STEP:.*]] = arith.constant 1 : index
+! CHECK:   fir.do_loop %[[IV:.*]] = %[[LB]] to %[[UB]] step %[[STEP]] {
+! CHECK:     %[[COORD1:.*]] = fir.coordinate_of %[[ARG0]], %[[IV]] : (!fir.ref<!fir.array<100xi32>>, index) -> !fir.ref<i32>
+! CHECK:     %[[COORD2:.*]] = fir.coordinate_of %[[ARG1]], %[[IV]] : (!fir.ref<!fir.array<100xi32>>, 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:   acc.yield %[[ARG0]] : !fir.ref<!fir.array<100xi32>>
+! CHECK: }
+
 ! CHECK-LABEL: acc.reduction.recipe @reduction_add_i32 : i32 reduction_operator <add> init {
 ! CHECK: ^bb0(%{{.*}}: i32):
 ! CHECK:   %[[INIT:.*]] = arith.constant 0 : i32
 ! CHECK:   acc.yield %[[INIT]] : i32
 ! CHECK: } combiner {
 ! CHECK: ^bb0(%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32):
 ! CHECK:   %[[COMBINED:.*]] = arith.addi %[[ARG0]], %[[ARG1]] : i32
 ! CHECK:   acc.yield %[[COMBINED]] : i32
 ! CHECK: }
 
 subroutine acc_reduction_add_int(a, b)
   integer :: a(100)
   integer :: i, b
 
   !$acc loop reduction(+:b)
   do i = 1, 100
     b = b + a(i)
   end do
 end subroutine
 
 ! CHECK-LABEL: func.func @_QPacc_reduction_add_int(
 ! CHECK-SAME:  %{{.*}}: !fir.ref<!fir.array<100xi32>> {fir.bindc_name = "a"}, %[[B:.*]]: !fir.ref<i32> {fir.bindc_name = "b"})
 ! CHECK:       acc.loop reduction(@reduction_add_i32 -> %[[B]] : !fir.ref<i32>)
 
+subroutine acc_reduction_add_int_array_1d(a, b)
+  integer :: a(100)
+  integer :: i, b(100)
+
+  !$acc loop reduction(+:b)
+  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_float(a, b)
   real :: a(100), b
   integer :: i
 
   !$acc loop reduction(+:b)
   do i = 1, 100
     b = b + a(i)
   end do
 end subroutine
 
 ! CHECK-LABEL: func.func @_QPacc_reduction_add_float(
 ! CHECK-SAME:  %{{.*}}: !fir.ref<!fir.array<100xf32>> {fir.bindc_name = "a"}, %[[B:.*]]: !fir.ref<f32> {fir.bindc_name = "b"})
 ! CHECK:       acc.loop reduction(@reduction_add_f32 -> %[[B]] : !fir.ref<f32>)
 
+subroutine acc_reduction_add_float_array_1d(a, b)
+  real :: a(100), b(100)
+  integer :: i
+
+  !$acc loop reduction(+:b)
+  do i = 1, 100
+    b(i) = b(i) + a(i)
+  end do
+end subroutine
+
+! CHECK-LABEL: func.func @_QPacc_reduction_add_float_array_1d(
+! CHECK-SAME:  %{{.*}}: !fir.ref<!fir.array<100xf32>> {fir.bindc_name = "a"}, %[[B:.*]]: !fir.ref<!fir.array<100xf32>> {fir.bindc_name = "b"})
+! CHECK:       acc.loop reduction(@reduction_add_ref_100xf32 -> %[[B]] : !fir.ref<!fir.array<100xf32>>)
+
 subroutine acc_reduction_mul_int(a, b)
   integer :: a(100)
   integer :: i, b
 
   !$acc loop reduction(*:b)
   do i = 1, 100
     b = b * a(i)
   end do
 end subroutine
 
 ! CHECK-LABEL: func.func @_QPacc_reduction_mul_int(
 ! CHECK-SAME:  %{{.*}}: !fir.ref<!fir.array<100xi32>> {fir.bindc_name = "a"}, %[[B:.*]]: !fir.ref<i32> {fir.bindc_name = "b"})
 ! CHECK:       acc.loop reduction(@reduction_mul_i32 -> %[[B]] : !fir.ref<i32>)
 
+subroutine acc_reduction_mul_int_array_1d(a, b)
+  integer :: a(100)
+  integer :: i, b(100)
+
+  !$acc loop reduction(*:b)
+  do i = 1, 100
+    b(i) = b(i) * a(i)
+  end do
+end subroutine
+
+! CHECK-LABEL: func.func @_QPacc_reduction_mul_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_mul_ref_100xi32 -> %[[B]] : !fir.ref<!fir.array<100xi32>>)
+
 subroutine acc_reduction_mul_float(a, b)
   real :: a(100), b
   integer :: i
 
   !$acc loop reduction(*:b)
   do i = 1, 100
     b = b * a(i)
   end do
 end subroutine
 
 ! CHECK-LABEL: func.func @_QPacc_reduction_mul_float(
 ! CHECK-SAME:  %{{.*}}: !fir.ref<!fir.array<100xf32>> {fir.bindc_name = "a"}, %[[B:.*]]: !fir.ref<f32> {fir.bindc_name = "b"})
 ! CHECK:       acc.loop reduction(@reduction_mul_f32 -> %[[B]] : !fir.ref<f32>)
 
+subroutine acc_reduction_mul_float_array_1d(a, b)
+  real :: a(100), b(100)
+  integer :: i
+
+  !$acc loop reduction(*:b)
+  do i = 1, 100
+    b(i) = b(i) * a(i)
+  end do
+end subroutine
+
+! CHECK-LABEL: func.func @_QPacc_reduction_mul_float_array_1d(
+! CHECK-SAME:  %{{.*}}: !fir.ref<!fir.array<100xf32>> {fir.bindc_name = "a"}, %[[B:.*]]: !fir.ref<!fir.array<100xf32>> {fir.bindc_name = "b"})
+! CHECK:       acc.loop reduction(@reduction_mul_ref_100xf32 -> %[[B]] : !fir.ref<!fir.array<100xf32>>)
 
 subroutine acc_reduction_min_int(a, b)
   integer :: a(100)
   integer :: i, b
 
   !$acc loop reduction(min:b)
   do i = 1, 100
     b = min(b, a(i))

mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp

 
 //===----------------------------------------------------------------------===//
 // ReductionRecipeOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult acc::ReductionRecipeOp::verifyRegions() {
   if (failed(verifyInitLikeSingleArgRegion(*this, getInitRegion(), "reduction",
                                            "init", getType(),
-                                           /*verifyYield=*/true)))
+                                           /*verifyYield=*/false)))
     return failure();
 
   if (getCombinerRegion().empty())
     return emitOpError() << "expects non-empty combiner region";
 
   Block &reductionBlock = getCombinerRegion().front();
   if (reductionBlock.getNumArguments() != 2 ||
       reductionBlock.getArgument(0).getType() != getType() ||

mlir/test/Dialect/OpenACC/invalid.mlir

 acc.reduction.recipe @reduction_i64 : i64 reduction_operator<add> init {
 ^bb0(%0: i32):
   %1 = arith.constant 0 : i64
   acc.yield %1 : i64
 } combiner {}
 
 // -----
 
-// expected-error@+1 {{expects init region to yield a value of the reduction type}}
-acc.reduction.recipe @reduction_i64 : i64 reduction_operator<add> init {
-^bb0(%0: i64):
-  %1 = arith.constant 0 : i32
-  acc.yield %1 : i32
-} combiner {}
-
-// -----
-
 // expected-error@+1 {{expects non-empty combiner region}}
 acc.reduction.recipe @reduction_i64 : i64 reduction_operator<add> init {
 ^bb0(%0: i64):
   %1 = arith.constant 0 : i64
   acc.yield %1 : i64
 } combiner {}
 
 // -----