[flang] lower matmul intrinsic to hlfir.matmul operation

[jeanPerier]

flang/lib/Lower/ConvertCall.cpp

1313–1314

Same as for SUM, I think it is probably best to dereference allocatable/pointers here to save complexity in MATAMUL.

It would be great if we could somehow deduce and do that automatically based on some knowledge about the arguments. I think it is doable for all arguments that are not LowerIntrinsicArgAs::AsInquired and for which handleDynamicOptional is false.
loadTrivialScalar should be callable for all arguments that are LowerIntrinsicArgAs::AsValue and for which handleDynamicOptional is false.

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

574

I have a feeling that this builder and the SUM builder are not very specific to the intrinsics and that the transformational operation hlfir.expr result type could be computed in lowering regardless of the intrinsic given stmtResultType, and that the default builder could be used.

[tblah]

Thanks for reviewing this. Changes are:

• Share argument dereferencing with hlfir.sum
• Share result type calculation with hlfir.sum

[jeanPerier]

Looks great, thanks for dealing with all the comments.

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

def hlfir_SumOp : hlfir_Op<"sum", [AttrSizedOperandSegments,
   );
 
   let results = (outs hlfir_ExprType);
 
   let assemblyFormat = [{
     $array (`dim` $dim^)? (`mask` $mask^)? 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_MatmulOp : hlfir_Op<"matmul",
     [DeclareOpInterfaceMethods<ArithFastMathInterface>]> {
   let summary = "MATMUL transformational intrinsic";
   let description = [{
     Matrix multiplication
   }];
 
   let arguments = (ins
     AnyFortranNumericalOrLogicalArrayObject:$lhs,
     AnyFortranNumericalOrLogicalArrayObject:$rhs,
     DefaultValuedAttr<Arith_FastMathAttr,
                       "::mlir::arith::FastMathFlags::none">:$fastmath
   );
 
   let results = (outs hlfir_ExprType);
 
   let assemblyFormat = [{
     $lhs $rhs attr-dict `:` functional-type(operands, results)
   }];
 
-  let builders = [OpBuilder<(ins "mlir::Value":$lhs,
-                                 "mlir::Value":$rhs,
-                                 "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

 static std::optional<hlfir::EntityWithAttributes>
 genHLFIRIntrinsicRefCore(PreparedActualArguments &loweredActuals,
                          const Fortran::evaluate::SpecificIntrinsic &intrinsic,
                          const fir::IntrinsicArgumentLoweringRules *argLowering,
                          CallContext &callContext) {
   if (!useHlfirIntrinsicOps)
     return genIntrinsicRefCore(loweredActuals, intrinsic, argLowering, callContext);
 
-  auto getOperandVector =
-    [](PreparedActualArguments &loweredActuals) {
-      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()->getOriginalActual();
-        operands.emplace_back(actual.getBase());
-      }
-      return operands;
-    };
+  fir::FirOpBuilder &builder = callContext.getBuilder();
+  mlir::Location loc = callContext.loc;
+
+  auto getOperandVector = [&](PreparedActualArguments &loweredActuals) {
+
+    for (size_t i = 0; i < loweredActuals.size(); ++i) {
+      std::optional<PreparedActualArgument> arg = loweredActuals[i];
+      if (!arg) {
+      hlfir::Entity actual = arg->getOriginalActual();
+      mlir::Value valArg;
+
+      fir::ArgLoweringRule argRules =
+          fir::lowerIntrinsicArgumentAs(*argLowering, i);
+      if (!argRules.handleDynamicOptional &&
+          argRules.lowerAs != fir::LowerIntrinsicArgAs::Inquired)
+        valArg = hlfir::derefPointersAndAllocatables(loc, builder, actual);
+      else
+        valArg = actual.getBase();
+
+      operands.emplace_back(valArg);
 
-  fir::FirOpBuilder &builder = callContext.getBuilder();
-  mlir::Location loc = callContext.loc;
+  auto computeResultType = [&](mlir::Value argArray,
+                               mlir::Type stmtResultType) -> mlir::Type {
+    hlfir::ExprType::Shape resultShape;
+    mlir::Type normalisedResult =
+        hlfir::getFortranElementOrSequenceType(stmtResultType);
+    mlir::Type elementType;
+    if (auto array = normalisedResult.dyn_cast<fir::SequenceType>()) {
+      resultShape = hlfir::ExprType::Shape{array.getShape()};
+      elementType = array.getEleTy();
+    } else {
+      elementType = normalisedResult;
+    }
+    return hlfir::ExprType::get(builder.getContext(), resultShape, elementType,
+                                /*polymorphic=*/false);
+  };
 
   if (intrinsic.name == "sum") {
     llvm::SmallVector<mlir::Value> operands = getOperandVector(loweredActuals);
     assert(operands.size() == 3);
-    mlir::Value array = hlfir::derefPointersAndAllocatables(
-        loc, builder, hlfir::Entity{operands[0]});
+    mlir::Value array = operands[0];
     mlir::Value dim = operands[1];
     if (dim)
       dim = hlfir::loadTrivialScalar(loc, builder, hlfir::Entity{dim});
     mlir::Value mask = operands[2];
+    mlir::Type resultTy = computeResultType(array, *callContext.resultType);
     // 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::SumOp sumOp =
+        builder.create<hlfir::SumOp>(loc, resultTy, array, dim, mask);
     return {hlfir::EntityWithAttributes{sumOp.getResult()}};
   }
+  if (intrinsic.name == "matmul") {
+    llvm::SmallVector<mlir::Value> operands = getOperandVector(loweredActuals);
+    mlir::Type resultTy =
+        computeResultType(operands[0], *callContext.resultType);
+    hlfir::MatmulOp matmulOp = builder.create<hlfir::MatmulOp>(
+        loc, resultTy, operands[0], operands[1]);
+
+    return {hlfir::EntityWithAttributes{matmulOp.getResult()}};
+  }
 
   // TODO add hlfir operations for other transformational intrinsics here
 
   // fallback to calling the intrinsic via fir.call
   return genIntrinsicRefCore(loweredActuals, intrinsic, argLowering,
                              callContext);
 }
 

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");
-
-  fir::SequenceType arrayTy =
-      hlfir::getFortranElementOrSequenceType(array.getType())
-          .dyn_cast<fir::SequenceType>();
-  assert(arrayTy && "array must be of array type");
-  mlir::Type numTy = arrayTy.getEleTy();
-
-  // get the result shape from the statement context
-  hlfir::ExprType::Shape resultShape;
-  if (auto array = stmtResultType.dyn_cast<fir::SequenceType>()) {
-    resultShape = hlfir::ExprType::Shape{array.getShape()};
-  }
-  mlir::Type resultType = hlfir::ExprType::get(
-      builder.getContext(), resultShape, numTy, /*polymorphic=*/false);
-
-  build(builder, result, resultType, array, dim, mask);
-}
-
 //===----------------------------------------------------------------------===//
 // MatmulOp
 //===----------------------------------------------------------------------===//
 
 mlir::LogicalResult hlfir::MatmulOp::verify() {
   mlir::Value lhs = getLhs();
   mlir::Value rhs = getRhs();
   fir::SequenceType lhsTy =
void hlfir::AssociateOp::build(mlir::OpBuilder &builder,
                                mlir::OperationState &result, mlir::Value source,
                                llvm::StringRef uniq_name, mlir::Value shape,
                                mlir::ValueRange typeparams,
                                fir::FortranVariableFlagsAttr fortran_attrs) {
   auto nameAttr = builder.getStringAttr(uniq_name);
   // TODO: preserve polymorphism of polymorphic expr.
   mlir::Type firVarType = fir::ReferenceType::get(
       getFortranElementOrSequenceType(source.getType()));
   mlir::Type hlfirVariableType =
       DeclareOp::getHLFIRVariableType(firVarType, /*hasExplicitLbs=*/false);
   mlir::Type i1Type = builder.getI1Type();
   build(builder, result, {hlfirVariableType, firVarType, i1Type}, source, shape,
         typeparams, nameAttr, fortran_attrs);
 }
 
 //===----------------------------------------------------------------------===//
 // EndAssociateOp

flang/test/Lower/HLFIR/matmul.f90

+! Test lowering of MATMUL intrinsic to HLFIR
+! RUN: bbc -emit-fir -hlfir -o - %s 2>&1 | FileCheck %s
+
+subroutine matmul1(lhs, rhs, res)
+  integer :: lhs(:,:), rhs(:,:), res(:,:)
+  res = MATMUL(lhs, rhs)
+endsubroutine
+! CHECK-LABEL: func.func @_QPmatmul1
+! CHECK:           %[[LHS:.*]]: !fir.box<!fir.array<?x?xi32>> {fir.bindc_name = "lhs"}
+! CHECK:           %[[RHS:.*]]: !fir.box<!fir.array<?x?xi32>> {fir.bindc_name = "rhs"}
+! CHECK:           %[[RES:.*]]: !fir.box<!fir.array<?x?xi32>> {fir.bindc_name = "res"}
+! CHECK-DAG:     %[[LHS_VAR:.*]]:2 = hlfir.declare %[[LHS]]
+! CHECK-DAG:     %[[RHS_VAR:.*]]:2 = hlfir.declare %[[RHS]]
+! CHECK-DAG:     %[[RES_VAR:.*]]:2 = hlfir.declare %[[RES]]
+! CHECK-NEXT:    %[[EXPR:.*]] = hlfir.matmul %[[LHS_VAR]]#0 %[[RHS_VAR]]#0 {fastmath = #arith.fastmath<contract>} : (!fir.box<!fir.array<?x?xi32>>, !fir.box<!fir.array<?x?xi32>>) -> !hlfir.expr<?x?xi32>
+! CHECK-NEXT:    hlfir.assign %[[EXPR]] to %[[RES_VAR]]#0 : !hlfir.expr<?x?xi32>, !fir.box<!fir.array<?x?xi32>>
+! CHECK-NEXT:    hlfir.destroy %[[EXPR]]
+! CHECK-NEXT:    return
+! CHECK-NEXT:   }