[flang] lower sum intrinsic to hlfir.sum operation

[tblah]

Summary

See https://reviews.llvm.org/D142897

[clementval]

flang/lib/Lower/ConvertCall.cpp

792

We probably want to have a more generic way to lower those intrinsic operations. Something like the table in flang/lib/Lower/IntrinsicCall.cpp.

[jeanPerier]

Nits about the builder and extendedValueToHlfirEntity usage. Looks good otherwise.

flang/lib/Lower/ConvertCall.cpp

552

fir::ExtendedValue should not contain hlfir::ExprType (all the helpers and functions only expect to work on fir::ExtendedValue that wraps fir/mlir types).

792

When lowering fully switch to the new expression representation, it will be possible to do so I think, but right now, it is not really possible because IntrinsicCall.cpp is based on the fir::ExtendedValue representation which requires array operands to be in memory, and the whole point of HLFIR is to try to avoid creating the temporary (to directly inline SUM(A+B) without ever creating a temp for A+B).

So I am OK with this not very optimal approach to start with, then we can improve this when we do not have to jungle between both lowering.

801

sumOp is not a fir::ExtendedValue. hlfir.epxr already are valid HLFIR entities, you can directly return hlfir::EntityWithAttributes{sumOp.getResult()}; (that will only assert that sumOp.getResult() is indeed a valid HLFIR entity type).

flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp

550–560

You could have used numType = hlfir::getFortranElementType(array.getType()) here, but I do not think you even need it if you always use the element type from stmtResultType (and since the assert should not be added here for properties that are verified in the verifier).

560

Better to not assert here since the verifier deals with this.

564

Is there a case where stmtResultType would be absent ? I think it would be simpler to require it to be present.

566–567

Same here.

570

Same here.

[jeanPerier]

flang/lib/Lower/ConvertCall.cpp

796

Here, I think it would be best to add dim = dim ? hlfir::loadTrivialScalar(loc, builder, dim) as well as array = hlfir::derefPointersAndAllocatables(loc, builder, array) so that the operands of the hlfir.sum are decorelated from the dereferences as much as possible.

The mask case is a bit trickier since it may be dynamically optional. hlfir.sum may need an extra optional i1 value argument to hold the dynamic presence of the mask in this case, but it is also ok to call derefPointersAndAllocatables on deallocated pointers/allocatable: the result is a fir.box that contains a null address.

[tblah]

Thanks for reviewing
Changes:

• Respond to review feedback
• Add tests

flang/lib/Lower/ConvertCall.cpp

796

See discussion on the first patch in the series. This omission is deliberate because I think it will make it simpler for passes to pattern match on this if we don't convert the operands until lowering to the runtime call.

For example, a pass might want to match MATMUL(TRANSPOSE(...), ....). This is simpler if the the generated HLFIR is something like

%transpose_expr = hlfir.transpose [...] : hlfir.expr<[...]>
%other_expr = [...] : hlfir.expr<[...]>
%mul_expr = hlfir.matmul %transpose_expr %other_expr

But I am open to your input on this

[jeanPerier]

flang/lib/Lower/ConvertCall.cpp

796

Your example makes absolute sense to me, but I think pattern matching will not be impacted by loads of scalars when obvious (like here with DIM), or dereference of allocatable and pointers.
I am inclined to think having a single place where this special cases are handled is simpler than in the operation transform that are likely to not be done with the allocatable/pointer case in mind.

For instance, would the current hlfir.sum handle:

subroutine foo(s, x, dim)
  integer, pointer :: dim
  real :: x(:, :), s(:)
  s = sum(x, dim)
end subroutine

It is a stupid edge case that is unlikely to be found in real apps, and, if an issue, would only be discovered quite late.

Using derefPointersAndAllocatables/loadTrivialScalar when possible does not affect the chaining of transformational since it does not touch values (hlfir.expr and simple scalars, which should be what is returned by sub-expression lowered nodes) and in my opinion will allow for simpler and more robust transformational operations.

I understand however that is is a bit inconvenient/not so nice to have to "split" the argument cooking between lowering and the pass that translate the operation to HLFIR.

[tblah]

Changes:

• Add load operation for pass-by-ref DIM argument
• derefPointersAndAllocatables on ARRAY argument
• Add a test for pass-by-pointer DIM argument

[jeanPerier]

Small nits. looks good otherwise!
Please merge at the same time of after the codegen patch (D143512) so that this patch does not break the end-to-end flow.

flang/lib/Lower/ConvertCall.cpp

779–790

It may be better to place this code in a helper lambda and call it when needed (only for SUM here), rather than doing this for all other intrinsics where "operands" will not be used.

792

Could you add a static llvm::cl::opt<bool> to control the intrinsic lowering to use the hlfir ops or not ?

I am OK with having it on by default, but I think it will be interesting to be able to run some test suits/benchmark with it on/off to catch potential issues.

796

If you want to assert, I think it would be most useful on operands[0] since the line above would crash if it is null (hlfir::Entity cannot be created from null values).

[tblah]

Fix nits in review

[jeanPerier]

LGTM

flang/include/flang/Optimizer/HLFIR/HLFIROps.td

def hlfir_SumOp : hlfir_Op<"sum", []> {
   );
 
   let results = (outs hlfir_ExprType);
 
   let assemblyFormat = [{
     $array $optionals attr-dict `:` functional-type(operands, results)
   }];
 
+  // dim and mask can be NULL, array must not be.
+  let builders = [OpBuilder<(ins "mlir::Value":$array,
+                                 "mlir::Value":$dim,
+                                 "mlir::Value":$mask,
+                                 "mlir::Type":$resultType)>];
+
   let hasVerifier = 1;
 }
 
 def hlfir_AssociateOp : hlfir_Op<"associate", [AttrSizedOperandSegments,
     DeclareOpInterfaceMethods<fir_FortranVariableOpInterface>]> {
   let summary = "Create a variable from an expression value";
   let description = [{
     Create a variable from an expression value.

flang/lib/Lower/ConvertCall.cpp

using PreparedActualArguments =
     llvm::SmallVector<std::optional<PreparedActualArgument>>;
 
 // Helper to transform a fir::ExtendedValue to an hlfir::EntityWithAttributes.
 static hlfir::EntityWithAttributes
 extendedValueToHlfirEntity(mlir::Location loc, fir::FirOpBuilder &builder,
                            const fir::ExtendedValue &exv,
                            llvm::StringRef name) {
   mlir::Value firBase = fir::getBase(exv);
-  if (fir::isa_trivial(firBase.getType()))
+  mlir::Type firBaseTy = firBase.getType();
+  if (fir::isa_trivial(firBaseTy) || firBaseTy.isa<hlfir::ExprType>())
     return hlfir::EntityWithAttributes{firBase};
   return hlfir::genDeclare(loc, builder, exv, name,
                            fir::FortranVariableFlagsAttr{});
 }
 
 /// Lower calls to user procedures with actual arguments that have been
 /// pre-lowered but not yet prepared according to the interface.
 /// This can be called for elemental procedures, but only with scalar
if (!fir::isa_trivial(resultEntity.getType())) {
     else
       asExpr = builder.create<hlfir::AsExprOp>(
           loc, resultEntity, builder.createBool(loc, mustBeFreed));
     resultEntity = hlfir::EntityWithAttributes{asExpr.getResult()};
   }
   return resultEntity;
 }
 
+/// Lower calls to intrinsic procedures with actual arguments that have been
+/// pre-lowered but have not yet been prepared according to the interface.
+static hlfir::EntityWithAttributes genHLFIRIntrinsicRefCore(
+    PreparedActualArguments &loweredActuals,
+    const Fortran::evaluate::SpecificIntrinsic &intrinsic,
+    const Fortran::lower::IntrinsicArgumentLoweringRules *argLowering,
+    CallContext &callContext) {
+  llvm::SmallVector<mlir::Value> operands;
+  operands.reserve(loweredActuals.size());
+  for (auto arg : llvm::enumerate(loweredActuals)) {
+    if (!arg.value()) {
+      operands.emplace_back();
+      continue;
+    }
+    hlfir::Entity actual = arg.value()->actual;
+    operands.emplace_back(actual.getBase());
+  }
+  fir::FirOpBuilder &builder = callContext.getBuilder();
+  mlir::Location loc = callContext.loc;
+
+  if (intrinsic.name == "sum") {
+    assert(operands.size() == 3);
+    mlir::Value array = operands[0];
+    assert(array);
+    mlir::Value dim = operands[1];
+    mlir::Value mask = operands[2];
+    // dim, mask can be NULL if these arguments were not given
+    hlfir::SumOp sumOp = builder.create<hlfir::SumOp>(loc, array, dim, mask,
+                                                      *callContext.resultType);
+    hlfir::EntityWithAttributes resultEntity = extendedValueToHlfirEntity(
+        loc, builder, sumOp.getResult(), ".tmp.intrinsic_result");
+    return resultEntity;
+  }
+
+  // TODO add hlfir operations for other transformational intrinsics here
+
+  // fallback to calling the intrinsic via fir.call
+  return genIntrinsicRefCore(loweredActuals, intrinsic, argLowering,
+                             callContext);
+}
+
 namespace {
 template <typename ElementalCallBuilderImpl>
 class ElementalCallBuilder {
 public:
   std::optional<hlfir::EntityWithAttributes>
   genElementalCall(PreparedActualArguments &loweredActuals, bool isImpure,
                    CallContext &callContext) {
     mlir::Location loc = callContext.loc;
ElementalIntrinsicCallBuilder(
       const Fortran::evaluate::SpecificIntrinsic &intrinsic,
       const Fortran::lower::IntrinsicArgumentLoweringRules *argLowering,
       bool isFunction)
       : intrinsic{intrinsic}, argLowering{argLowering}, isFunction{isFunction} {
   }
   std::optional<hlfir::Entity>
   genElementalKernel(PreparedActualArguments &loweredActuals,
                      CallContext &callContext) {
-    return genIntrinsicRefCore(loweredActuals, intrinsic, argLowering,
+    return genHLFIRIntrinsicRefCore(loweredActuals, intrinsic, argLowering,
                                     callContext);
   }
   // Elemental intrinsic functions cannot modify their arguments.
   bool argMayBeModifiedByCall(int) const { return !isFunction; }
   bool canLoadActualArgumentBeforeLoop(int) const {
     // Elemental intrinsic functions never need the actual addresses
     // of their arguments.
     return isFunction;
   }
genIntrinsicRef(const Fortran::evaluate::SpecificIntrinsic &intrinsic,
 
   if (callContext.isElementalProcWithArrayArgs()) {
     // All intrinsic elemental functions are pure.
     const bool isFunction = callContext.resultType.has_value();
     return ElementalIntrinsicCallBuilder{intrinsic, argLowering, isFunction}
         .genElementalCall(loweredActuals, /*isImpure=*/!isFunction, callContext)
         .value();
   }
-  hlfir::EntityWithAttributes result =
-      genIntrinsicRefCore(loweredActuals, intrinsic, argLowering, callContext);
+  hlfir::EntityWithAttributes result = genHLFIRIntrinsicRefCore(
+      loweredActuals, intrinsic, argLowering, callContext);
   if (result.getType().isa<hlfir::ExprType>()) {
     fir::FirOpBuilder *bldr = &callContext.getBuilder();
     callContext.stmtCtx.attachCleanup(
         [=]() { bldr->create<hlfir::DestroyOp>(loc, result); });
   }
   return result;
 }
 

flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp

if (resultTy.isArray()) {
     if (resultTy.getElementType() != numTy)
       return emitOpError(
           "result must have the same element type as ARRAY argument");
   }
 
   return mlir::success();
 }
 
+void hlfir::SumOp::build(mlir::OpBuilder &builder, mlir::OperationState &result,
+                         mlir::Value array, mlir::Value dim, mlir::Value mask,
+                         mlir::Type stmtResultType) {
+  assert(array && "array argument is not optional");
+  mlir::Type arrayTy = unwrapType(array.getType());
+  mlir::Type numTy;
+  if (auto arrayExprTy = arrayTy.dyn_cast<hlfir::ExprType>()) {
+    numTy = arrayExprTy.getEleTy();
+  } else if (auto arraySeqTy = arrayTy.dyn_cast<fir::SequenceType>()) {
+    numTy = arraySeqTy.getEleTy();
+  } else {
+    llvm_unreachable("bad array type");
+  }
+  assert(hlfir::isFortranScalarNumericalType(numTy));
+
+  // get the result shape from the statement context
+  hlfir::ExprType::Shape resultShape;
+  if (stmtResultType) {
+    if (auto array = stmtResultType.dyn_cast<fir::SequenceType>()) {
+      assert(array.getEleTy() == numTy && "Unexpected array result type");
+      assert(dim && "Non-scalar result type without DIM argument");
+      resultShape = hlfir::ExprType::Shape{array.getShape()};
+    } else {
+      assert(stmtResultType == numTy && "Scalar result should match ARRAY arg");
+      // we might get a scalar result even with a DIM argument if the array
+      // has rank 1, so don't assert !dim
+    }
+  }
+  mlir::Type resultType = hlfir::ExprType::get(
+      builder.getContext(), resultShape, numTy, /*polymorphic=*/false);
+
+  llvm::SmallVector<mlir::Value, 2> optionals;
+  for (mlir::Value arg : {dim, mask}) {
+    if (arg)
+      optionals.push_back(arg);
+  }
+
+  build(builder, result, resultType, array, optionals);
+}
+
 //===----------------------------------------------------------------------===//
 // AssociateOp
 //===----------------------------------------------------------------------===//
 
 void hlfir::AssociateOp::build(mlir::OpBuilder &builder,
                                mlir::OperationState &result, mlir::Value source,
                                llvm::StringRef uniq_name, mlir::Value shape,
                                mlir::ValueRange typeparams,