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

[clementval]

Summary

Lower dynamic array slices reduction. Extent is passed to
the init and combiner region as entry block arguments.

Depends on D155655 and D155656

[clementval]

Rebase

[clementval]

Rebase

[clementval]

Try to rebase again

flang/lib/Lower/OpenACC.cpp

for (auto b : bounds) {
           fir::getIntIfConstant(boundsOp.getUpperbound())) {
         int64_t ext = *fir::getIntIfConstant(boundsOp.getUpperbound()) -
                       *fir::getIntIfConstant(boundsOp.getLowerbound()) + 1;
         shape.push_back(ext);
       } else if (boundsOp.getExtent() &&
                  fir::getIntIfConstant(boundsOp.getExtent())) {
         shape.push_back(*fir::getIntIfConstant(boundsOp.getExtent()));
       } else {
-        return ty; // TODO: handle dynamic shaped array slice.
+        shape.push_back(fir::SequenceType::getUnknownExtent());
       }
     }
     if (shape.empty() || shape.size() != bounds.size())
       return ty;
     auto newSeqTy = fir::SequenceType::get(shape, seqTy.getEleTy());
     if (mlir::isa<fir::ReferenceType>(ty))
       return fir::ReferenceType::get(newSeqTy);
     return newSeqTy;
if (op == mlir::acc::ReductionOperator::AccMin) {
     }
 
     if constexpr (std::is_same_v<R, int64_t>)
       return value;
   }
   llvm_unreachable("OpenACC reduction unsupported type");
 }
 
+/// Check if the DataBoundsOp is a constant bound (lb and ub are constants or 
+/// extent is a constant).
+bool isConstantBound(mlir::acc::DataBoundsOp &op) {
+  if (op.getLowerbound() && fir::getIntIfConstant(op.getLowerbound()) &&
+      op.getUpperbound() && fir::getIntIfConstant(op.getUpperbound()))
+    return true;
+  if (op.getExtent() && fir::getIntIfConstant(op.getExtent()))
+    return true;
+  return false;
+}
+
 /// 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()))
+    if (!isConstantBound(op))
       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) {
+static mlir::Value genScalarReductionValue(fir::FirOpBuilder &builder,
+                                         mlir::Location loc, mlir::Type ty,
+                                         mlir::acc::ReductionOperator op) {
   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::AccMin ||
     mlir::Type floatTy =
         Fortran::lower::convertReal(builder.getContext(), cmplxTy.getFKind());
     mlir::Value realInit = builder.createRealConstant(
         loc, floatTy, getReductionInitValue<int64_t>(op, cmplxTy));
     mlir::Value imagInit = builder.createRealConstant(loc, floatTy, 0.0);
     return fir::factory::Complex{builder, loc}.createComplex(
         cmplxTy.getFKind(), realInit, imagInit);
   }
+  llvm::report_fatal_error("Unsupported OpenACC reduction type");
+}
+
+static mlir::Value genReductionInitValue(fir::FirOpBuilder &builder,
+                                         mlir::Location loc, mlir::Type ty,
+                                         mlir::acc::ReductionOperator op,
+                                         llvm::SmallVector<mlir::Value> &bounds) {
   if (auto refTy = mlir::dyn_cast<fir::ReferenceType>(ty)) {
     if (auto seqTy = mlir::dyn_cast<fir::SequenceType>(refTy.getEleTy())) {
+      if (seqTy.hasDynamicExtents()) {
+        mlir::Block* block = builder.getBlock();
+        assert(seqTy.getDimension() == (block->getNumArguments() - 1) &&
+            "expect dynamic extents passed as block arguments");
+        llvm::SmallVector<mlir::Value> extents;
+        for (unsigned i = 1; i < block->getNumArguments(); ++i) 
+          extents.push_back(block->getArgument(i));
+        mlir::Value temp = builder.createTemporary(
+            loc, seqTy, ".reduction_init", extents, {}, {});
+        mlir::Value shapeOp = builder.create<fir::ShapeOp>(loc, extents);
+        mlir::Value arrLoad = builder.create<fir::ArrayLoadOp>(
+            loc, seqTy, temp, shapeOp, /*slice=*/mlir::Value{}, std::nullopt);  
+
+        mlir::Type idxTy = builder.getIndexType();
+        llvm::SmallVector<fir::DoLoopOp> loops;
+        llvm::SmallVector<mlir::Value> ivs;
+        mlir::Value innerArg;
+        for (auto ext : llvm::reverse(extents)) {
+          mlir::Value lb = builder.createIntegerConstant(loc, idxTy, 0);
+          mlir::Value c1 = builder.createIntegerConstant(loc, idxTy, 1);
+          mlir::Value ub = builder.create<mlir::arith::SubIOp>(loc, ext, c1);
+          auto loop = builder.create<fir::DoLoopOp>(loc, lb, ub, c1,
+                                                    /*unordered=*/false,
+                                                    /*finalCount=*/false,
+                                                    mlir::ValueRange{arrLoad});
+          builder.setInsertionPointToStart(loop.getBody());
+          loops.push_back(loop);
+          ivs.push_back(loop.getInductionVar());
+          innerArg = loop.getRegionIterArgs().front();
+        }
+        mlir::Value initValue =
+            genScalarReductionValue(builder, loc, seqTy.getEleTy(), op);
+        mlir::Value arrUp = builder.create<fir::ArrayUpdateOp>(loc, seqTy,
+            innerArg, initValue, ivs, mlir::ValueRange());
+        builder.create<fir::ResultOp>(loc, arrUp);
+        builder.setInsertionPointAfter(loops[0]);
+        builder.create<fir::ArrayMergeStoreOp>(loc, arrLoad, loops[0].getResult(0), temp, mlir::Value(), mlir::ValueRange());
+        return temp;
+      } else {
         mlir::Type vecTy =
             mlir::VectorType::get(seqTy.getShape(), seqTy.getEleTy());
         auto shTy = vecTy.cast<mlir::ShapedType>();
         if (seqTy.getEleTy().isIntOrIndex())
           return builder.create<mlir::arith::ConstantOp>(
               loc, vecTy,
               mlir::DenseElementsAttr::get(
                   shTy,
                   getReductionInitValue<llvm::APInt>(op, seqTy.getEleTy())));
         if (mlir::isa<mlir::FloatType>(seqTy.getEleTy()))
           return builder.create<mlir::arith::ConstantOp>(
               loc, vecTy,
               mlir::DenseElementsAttr::get(
                   shTy,
                   getReductionInitValue<llvm::APFloat>(op, seqTy.getEleTy())));
       }
     }
+  }
 
-  llvm::report_fatal_error("Unsupported OpenACC reduction type");
+  return genScalarReductionValue(builder, loc, ty, op);
 }
 
 template <typename Op>
 static mlir::Value genLogicalCombiner(fir::FirOpBuilder &builder,
                                       mlir::Location loc, mlir::Value value1,
                                       mlir::Value value2) {
   mlir::Type i1 = builder.getI1Type();
   mlir::Value v1 = builder.create<fir::ConvertOp>(loc, i1, value1);

 
 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.hasDynamicExtents())
-        TODO(loc, "OpenACC reduction on array with dynamic extents");
+      if (seqTy.hasDynamicExtents()) {
+        mlir::Block* block = builder.getBlock();
+        assert(seqTy.getDimension() == (block->getNumArguments() - 2) &&
+            "expect dynamic extents passed as block arguments");
+        llvm::SmallVector<mlir::Value> extents;
+        for (unsigned i = 2; i < block->getNumArguments(); ++i) 
+          extents.push_back(block->getArgument(i));
+        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 : llvm::reverse(extents)) {
+          auto lb = builder.create<mlir::arith::ConstantOp>(
+              loc, idxTy, builder.getIntegerAttr(idxTy, 0));
+          auto c1 = builder.createIntegerConstant(loc, idxTy, 1);
+          auto ub = builder.create<mlir::arith::SubIOp>(loc, ext, c1);
+          auto loop = builder.create<fir::DoLoopOp>(loc, lb, ub, c1,
+                                                    /*unordered=*/false);
+          builder.setInsertionPointToStart(loop.getBody());
+          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(loops[0]);
+
+        return value1;
+      } else {
         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 : llvm::reverse(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, 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());
           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(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))
       return builder.create<mlir::arith::AddFOp>(loc, value1, value2);
     if (auto cmplxTy = mlir::dyn_cast_or_null<fir::ComplexType>(ty))
       return builder.create<fir::AddcOp>(loc, value1, value2);
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);
+  
+  llvm::SmallVector<mlir::Type> argTypes;
+  llvm::SmallVector<mlir::Location> argLocs;
+  argTypes.push_back(ty);
+  argLocs.push_back(loc);
+  if (hasDynamicShape(bounds)) {
+    for (unsigned i; i < bounds.size(); ++i) {
+      argTypes.push_back(builder.getIndexType());
+      argLocs.push_back(loc);
+    }
+  }
   builder.createBlock(&recipe.getInitRegion(), recipe.getInitRegion().end(),
-                      {ty}, {loc});
+                      argTypes, argLocs);
   builder.setInsertionPointToEnd(&recipe.getInitRegion().back());
   mlir::Value initValue = genReductionInitValue(builder, loc, ty, op, bounds);
   builder.create<mlir::acc::YieldOp>(loc, initValue);
 
+  argTypes.insert(argTypes.begin(), ty);
+  argLocs.insert(argLocs.begin(), loc);
   builder.createBlock(&recipe.getCombinerRegion(),
-                      recipe.getCombinerRegion().end(), {ty, ty}, {loc, loc});
+                      recipe.getCombinerRegion().end(), argTypes, argLocs);
   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;
 }
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 (hasDynamicShape(bounds))
-      TODO(operandLocation, "OpenACC reductions with dynamic shaped array");
-
     mlir::Type reductionTy = fir::unwrapRefType(baseAddr.getType());
     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()))

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_?xi32" : !fir.ref<!fir.array<?xi32>> reduction_operator <add> init {
+! CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.array<?xi32>>, %[[EXT:.*]]: index):
+! CHECK:   %[[ALLOCA:.*]] = fir.alloca !fir.array<?xi32>, %arg1 {bindc_name = ".reduction_init"}
+! CHECK:   %[[SHAPE:.*]] = fir.shape %[[EXT]] : (index) -> !fir.shape<1>
+! CHECK:   %[[ARR_LOAD:.*]] = fir.array_load %[[ALLOCA]](%[[SHAPE]]) : (!fir.ref<!fir.array<?xi32>>, !fir.shape<1>) -> !fir.array<?xi32>
+! CHECK:   %[[C0:.*]] = arith.constant 0 : index
+! CHECK:   %[[C1:.*]] = arith.constant 1 : index
+! CHECK:   %[[UB:.*]] = arith.subi %[[EXT]], %[[C1]] : index
+! CHECK:   %[[DO_RES:.*]] = fir.do_loop %[[IV:.*]] = %[[C0]] to %[[UB]] step %[[C1]] iter_args(%[[ARGS:.*]] = %[[ARR_LOAD]]) -> (!fir.array<?xi32>) {
+! CHECK:     %[[VALUE:.*]] = arith.constant 0 : i32
+! CHECK:     %[[ARR_UP:.*]] = fir.array_update %[[ARGS]], %[[VALUE]], %[[IV]] : (!fir.array<?xi32>, i32, index) -> !fir.array<?xi32>
+! CHECK:     fir.result %[[ARR_UP]] : !fir.array<?xi32>
+! CHECK:   }
+! CHECK:   fir.array_merge_store %[[ARR_LOAD]], %[[DO_RES]] to %[[ALLOCA]] : !fir.array<?xi32>, !fir.array<?xi32>, !fir.ref<!fir.array<?xi32>>
+! CHECK:   acc.yield %[[ALLOCA]] : !fir.ref<!fir.array<?xi32>>
+! CHECK: } combiner {
+! CHECK: ^bb0(%[[ARG0:.*]]: !fir.ref<!fir.array<?xi32>>, %[[ARG1:.*]]: !fir.ref<!fir.array<?xi32>>, %[[EXT:.*]]: index):
+! CHECK:   %[[C0:.*]] = arith.constant 0 : index
+! CHECK:   %[[C1:.*]] = arith.constant 1 : index
+! CHECK:   %[[UB:.*]] = arith.subi %[[EXT]], %[[C1]] : index
+! CHECK:   fir.do_loop %[[IV:.*]] = %[[C0]] to %[[UB]] step %[[C1]] {
+! CHECK:     %[[COORD0:.*]] = fir.coordinate_of %[[ARG0]], %[[IV]] : (!fir.ref<!fir.array<?xi32>>, index) -> !fir.ref<i32>
+! CHECK:     %[[COORD1:.*]] = fir.coordinate_of %[[ARG1]], %[[IV]] : (!fir.ref<!fir.array<?xi32>>, index) -> !fir.ref<i32>
+! CHECK:     %[[LOAD0:.*]] = fir.load %[[COORD0]] : !fir.ref<i32>
+! CHECK:     %[[LOAD1:.*]] = fir.load %[[COORD1]] : !fir.ref<i32>
+! CHECK:     %[[COMBINED:.*]] = arith.addi %[[LOAD0]], %[[LOAD1]] : i32
+! CHECK:     fir.store %[[COMBINED]] to %[[COORD0]] : !fir.ref<i32>
+! CHECK:   }
+! CHECK:   acc.yield %[[ARG0]] : !fir.ref<!fir.array<?xi32>>
+! CHECK: }
+
 ! 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-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>>)
+
+subroutine acc_reduction_add_dynamic_slice(a, n, m)
+  integer :: a(100)
+  integer :: n, m
+  !$acc parallel reduction(+:a(n:m))
+  !$acc end parallel
+end subroutine