[flang][openacc] Support static slice in reduction lowering

[clementval]

Summary

Lower static array slices reduction with the correct
type.

[razvanlupusoru]

Looks good to me!

flang/include/flang/Lower/OpenACC.h

 //===----------------------------------------------------------------------===//
 
 #ifndef FORTRAN_LOWER_OPENACC_H
 #define FORTRAN_LOWER_OPENACC_H
 
 #include "mlir/Dialect/OpenACC/OpenACC.h"
 
 namespace llvm {
+template <typename T, unsigned N>
+class SmallVector;
 class StringRef;
 }
 
 namespace mlir {
 class Location;
 class Type;
 class OpBuilder;
+class Value;
 } // namespace mlir
 
 namespace fir {
 class FirOpBuilder;
 }
 
 namespace Fortran {
 namespace parser {

 mlir::acc::PrivateRecipeOp createOrGetPrivateRecipe(mlir::OpBuilder &,
                                                     llvm::StringRef,
                                                     mlir::Location, mlir::Type);
 
 /// Get a acc.reduction.recipe op for the given type or create it if it does not
 /// exist yet.
 mlir::acc::ReductionRecipeOp
 createOrGetReductionRecipe(fir::FirOpBuilder &, llvm::StringRef, mlir::Location,
-                           mlir::Type, mlir::acc::ReductionOperator);
+                           mlir::Type, mlir::acc::ReductionOperator,
+                           llvm::SmallVector<mlir::Value> &);
 
 /// Get a acc.firstprivate.recipe op for the given type or create it if it does
 /// not exist yet.
 mlir::acc::FirstprivateRecipeOp createOrGetFirstprivateRecipe(mlir::OpBuilder &,
                                                               llvm::StringRef,
                                                               mlir::Location,
                                                               mlir::Type);
 
 } // namespace lower
 } // namespace Fortran
 
 #endif // FORTRAN_LOWER_OPENACC_H

flang/lib/Lower/OpenACC.cpp

if (op == mlir::acc::ReductionOperator::AccMin) {
     }
 
     if constexpr (std::is_same_v<R, int64_t>)
       return value;
   }
   llvm_unreachable("OpenACC reduction unsupported type");
 }
 
-static mlir::Value genReductionInitValue(fir::FirOpBuilder &builder,
-                                         mlir::Location loc, mlir::Type ty,
-                                         mlir::acc::ReductionOperator op) {
+/// Determine if the bounds represent a dynamic shape.
+bool hasDynamicShape(llvm::SmallVector<mlir::Value> &bounds) {
+  if (bounds.empty())
+    return false;
+  for (auto b : bounds) {
+    auto op = mlir::dyn_cast<mlir::acc::DataBoundsOp>(b.getDefiningOp());
+    if (((op.getLowerbound() && !fir::getIntIfConstant(op.getLowerbound())) ||
+         (op.getUpperbound() && !fir::getIntIfConstant(op.getUpperbound()))) &&
+        op.getExtent() && !fir::getIntIfConstant(op.getExtent()))
+      return true;
+  }
+  return false;
+}
+
+static mlir::Value
+genReductionInitValue(fir::FirOpBuilder &builder, mlir::Location loc,
+                      mlir::Type ty, mlir::acc::ReductionOperator op,
+                      llvm::SmallVector<mlir::Value> &bounds) {
   if (op == mlir::acc::ReductionOperator::AccLand ||
       op == mlir::acc::ReductionOperator::AccLor ||
       op == mlir::acc::ReductionOperator::AccEqv ||
       op == mlir::acc::ReductionOperator::AccNeqv) {
     assert(mlir::isa<fir::LogicalType>(ty) && "expect fir.logical type");
     bool value = true; // .true. for .and. and .eqv.
     if (op == mlir::acc::ReductionOperator::AccLor ||
         op == mlir::acc::ReductionOperator::AccNeqv)
if (op == mlir::acc::ReductionOperator::AccNeqv)
     return genComparisonCombiner(builder, loc, mlir::arith::CmpIPredicate::ne,
                                  value1, value2);
 
   TODO(loc, "reduction operator");
 }
 
 mlir::acc::ReductionRecipeOp Fortran::lower::createOrGetReductionRecipe(
     fir::FirOpBuilder &builder, llvm::StringRef recipeName, mlir::Location loc,
-    mlir::Type ty, mlir::acc::ReductionOperator op) {
+    mlir::Type ty, mlir::acc::ReductionOperator op,
+    llvm::SmallVector<mlir::Value> &bounds) {
   mlir::ModuleOp mod =
       builder.getBlock()->getParent()->getParentOfType<mlir::ModuleOp>();
   if (auto recipe = mod.lookupSymbol<mlir::acc::ReductionRecipeOp>(recipeName))
     return recipe;
 
   auto crtPos = builder.saveInsertionPoint();
   mlir::OpBuilder modBuilder(mod.getBodyRegion());
   auto recipe =
       modBuilder.create<mlir::acc::ReductionRecipeOp>(loc, recipeName, ty, op);
   builder.createBlock(&recipe.getInitRegion(), recipe.getInitRegion().end(),
                       {ty}, {loc});
   builder.setInsertionPointToEnd(&recipe.getInitRegion().back());
-  mlir::Value initValue = genReductionInitValue(builder, loc, ty, op);
+  mlir::Value initValue = genReductionInitValue(builder, loc, ty, op, bounds);
   builder.create<mlir::acc::YieldOp>(loc, initValue);
 
   builder.createBlock(&recipe.getCombinerRegion(),
                       recipe.getCombinerRegion().end(), {ty, ty}, {loc, loc});
   builder.setInsertionPointToEnd(&recipe.getCombinerRegion().back());
   mlir::Value v1 = recipe.getCombinerRegion().front().getArgument(0);
   mlir::Value v2 = recipe.getCombinerRegion().front().getArgument(1);
   mlir::Value combinedValue = genCombiner(builder, loc, op, ty, v1, v2);
   builder.create<mlir::acc::YieldOp>(loc, combinedValue);
   builder.restoreInsertionPoint(crtPos);
   return recipe;
 }
 
-/// Determine if the bounds represent a dynamic shape.
-bool hasDynamicShape(llvm::SmallVector<mlir::Value> &bounds) {
-  if (bounds.empty())
-    return false;
-  for (auto b : bounds) {
-    auto op = mlir::dyn_cast<mlir::acc::DataBoundsOp>(b.getDefiningOp());
-    if (((op.getLowerbound() && !fir::getIntIfConstant(op.getLowerbound())) ||
-         (op.getUpperbound() && !fir::getIntIfConstant(op.getUpperbound()))) &&
-        op.getExtent() && !fir::getIntIfConstant(op.getExtent()))
-      return true;
-  }
-  return false;
-}
-
 static void
 genReductions(const Fortran::parser::AccObjectListWithReduction &objectList,
               Fortran::lower::AbstractConverter &converter,
               Fortran::semantics::SemanticsContext &semanticsContext,
               Fortran::lower::StatementContext &stmtCtx,
               llvm::SmallVectorImpl<mlir::Value> &reductionOperands,
               llvm::SmallVector<mlir::Attribute> &reductionRecipes) {
   fir::FirOpBuilder &builder = converter.getFirOpBuilder();
if (auto seqTy = mlir::dyn_cast<fir::SequenceType>(reductionTy))
       reductionTy = seqTy.getEleTy();
 
     if (!fir::isa_trivial(reductionTy) &&
         ((fir::isAllocatableType(reductionTy) ||
           fir::isPointerType(reductionTy)) &&
          !bounds.empty()))
       TODO(operandLocation, "reduction with unsupported type");
 
+    mlir::Type retTy = getTypeFromBounds(bounds, baseAddr.getType());
     auto op = createDataEntryOp<mlir::acc::ReductionOp>(
         builder, operandLocation, baseAddr, asFortran, bounds,
-        /*structured=*/true, mlir::acc::DataClause::acc_reduction,
-        baseAddr.getType());
+        /*structured=*/true, mlir::acc::DataClause::acc_reduction, retTy);
     mlir::Type ty = fir::unwrapRefType(op.getAccPtr().getType());
     if (!fir::isa_trivial(ty))
-      ty = baseAddr.getType();
+      ty = retTy;
     std::string recipeName = fir::getTypeAsString(
         ty, converter.getKindMap(),
         ("reduction_" + stringifyReductionOperator(mlirOp)).str());
     mlir::acc::ReductionRecipeOp recipe =
-        Fortran::lower::createOrGetReductionRecipe(builder, recipeName,
-                                                   operandLocation, ty, mlirOp);
+        Fortran::lower::createOrGetReductionRecipe(
+            builder, recipeName, operandLocation, ty, mlirOp, bounds);
     reductionRecipes.push_back(mlir::SymbolRefAttr::get(
         builder.getContext(), recipe.getSymName().str()));
     reductionOperands.push_back(op.getAccPtr());
   }
 }
 
 static void
 addOperands(llvm::SmallVectorImpl<mlir::Value> &operands,

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

 ! This test checks lowering of OpenACC reduction clause.
 
 ! RUN: bbc -fopenacc -emit-fir %s -o - | FileCheck %s
 
+! CHECK-LABEL: acc.reduction.recipe @reduction_add_ref_10xi32 : !fir.ref<!fir.array<10xi32>> reduction_operator <add> init {
+! CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.array<10xi32>>):
+! CHECK:   %[[CST:.*]] = arith.constant dense<0> : vector<10xi32>
+! CHECK:   acc.yield %[[CST]] : vector<10xi32>
+! CHECK: } combiner {
+! CHECK: ^bb0(%[[ARG0:.*]]: !fir.ref<!fir.array<10xi32>>, %[[ARG1:.*]]: !fir.ref<!fir.array<10xi32>>):
+! CHECK:   %[[LB:.*]] = arith.constant 0 : index
+! CHECK:   %[[UB:.*]] = arith.constant 9 : 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<10xi32>>, index) -> !fir.ref<i32>
+! CHECK:     %[[COORD2:.*]] = fir.coordinate_of %[[ARG1]], %[[IV]] : (!fir.ref<!fir.array<10xi32>>, 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<10xi32>>
+! CHECK: }
+
 ! CHECK-LABEL: acc.reduction.recipe @reduction_mul_z32 : !fir.complex<4> reduction_operator <mul> init {
 ! CHECK: ^bb0(%{{.*}}: !fir.complex<4>):
 ! CHECK:   %[[REAL:.*]] = arith.constant 1.000000e+00 : f32
 ! CHECK:   %[[IMAG:.*]] = arith.constant 0.000000e+00 : f32
 ! CHECK:   %[[UNDEF:.*]] = fir.undefined !fir.complex<4>
 ! CHECK:   %[[UNDEF1:.*]] = fir.insert_value %[[UNDEF]], %[[REAL]], [0 : index] : (!fir.complex<4>, f32) -> !fir.complex<4>
 ! CHECK:   %[[UNDEF2:.*]] = fir.insert_value %[[UNDEF1]], %[[IMAG]], [1 : index] : (!fir.complex<4>, f32) -> !fir.complex<4>
 ! CHECK:   acc.yield %[[UNDEF2]] : !fir.complex<4>

 end subroutine
 
 ! CHECK-LABEL: func.func @_QPacc_reduction_add_pointer(
 ! CHECK-SAME: %[[ARG0:.*]]: !fir.ref<!fir.box<!fir.ptr<i32>>> {fir.bindc_name = "i"})
 ! CHECK: %[[LOAD:.*]] = fir.load %[[ARG0]] : !fir.ref<!fir.box<!fir.ptr<i32>>>
 ! CHECK: %[[BOX_ADDR:.*]] = fir.box_addr %[[LOAD]] : (!fir.box<!fir.ptr<i32>>) -> !fir.ptr<i32>
 ! CHECK: %[[RED:.*]] = acc.reduction varPtr(%[[BOX_ADDR]] : !fir.ptr<i32>) -> !fir.ptr<i32> {name = "i"}
 ! CHECK: acc.parallel reduction(@reduction_add_i32 -> %[[RED]] : !fir.ptr<i32>)
+
+subroutine acc_reduction_add_static_slice(a)
+  integer :: a(100)
+  !$acc parallel reduction(+:a(11:20))
+  !$acc end parallel
+end subroutine
+
+! CHECK-LABEL: func.func @_QPacc_reduction_add_static_slice(
+! CHECK-SAME: %[[ARG0:.*]]: !fir.ref<!fir.array<100xi32>> {fir.bindc_name = "a"})
+! CHECK: %[[C1:.*]] = arith.constant 1 : index
+! CHECK: %[[LB:.*]] = arith.constant 10 : index
+! CHECK: %[[UB:.*]] = arith.constant 19 : index
+! CHECK: %[[BOUND:.*]] = acc.bounds lowerbound(%[[LB]] : index) upperbound(%[[UB]] : index) stride(%[[C1]] : index) startIdx(%[[C1]] : index)
+! CHECK: %[[RED:.*]] = acc.reduction varPtr(%[[ARG0]] : !fir.ref<!fir.array<100xi32>>) bounds(%[[BOUND]]) -> !fir.ref<!fir.array<10xi32>> {name = "a(11:20)"}
+! CHECK: acc.parallel reduction(@reduction_add_ref_10xi32 -> %[[RED]] : !fir.ref<!fir.array<10xi32>>)