diff --git a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td index 2dd85a2c5c181..beb50a48d86df 100644 --- a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td +++ b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td @@ -359,6 +359,27 @@ def hlfir_AnyOp : hlfir_Op<"any", []> { let hasVerifier = 1; } +def hlfir_CountOp : hlfir_Op<"count", [AttrSizedOperandSegments]> { + let summary = "COUNT transformational intrinsic"; + let description = [{ + Takes a logical and counts the number of true values. + }]; + + let arguments = (ins + AnyFortranLogicalArrayObject:$mask, + Optional:$dim, + Optional:$kind + ); + + let results = (outs AnyFortranValue); + + let assemblyFormat = [{ + $mask (`dim` $dim^)? (`kind` $kind^)? attr-dict `:` functional-type(operands, results) + }]; + + let hasVerifier = 1; +} + def hlfir_ProductOp : hlfir_Op<"product", [AttrSizedOperandSegments, DeclareOpInterfaceMethods]> { @@ -430,7 +451,7 @@ def hlfir_SumOp : hlfir_Op<"sum", [AttrSizedOperandSegments, let hasVerifier = 1; } -def hlifr_DotProductOp : hlfir_Op<"dot_product", +def hlfir_DotProductOp : hlfir_Op<"dot_product", [DeclareOpInterfaceMethods]> { let summary = "DOT_PRODUCT transformational intrinsic"; let description = [{ diff --git a/flang/lib/Lower/ConvertCall.cpp b/flang/lib/Lower/ConvertCall.cpp index 604291cdbec6d..4f4505c8b0664 100644 --- a/flang/lib/Lower/ConvertCall.cpp +++ b/flang/lib/Lower/ConvertCall.cpp @@ -1497,6 +1497,18 @@ genHLFIRIntrinsicRefCore(PreparedActualArguments &loweredActuals, return {hlfir::EntityWithAttributes{dotProductOp.getResult()}}; } + if (intrinsicName == "count") { + llvm::SmallVector operands = getOperandVector(loweredActuals); + mlir::Value array = operands[0]; + mlir::Value dim = operands[1]; + if (dim) + dim = hlfir::loadTrivialScalar(loc, builder, hlfir::Entity{dim}); + mlir::Value kind = operands[2]; + mlir::Type resultTy = computeResultType(array, *callContext.resultType); + hlfir::CountOp countOp = + builder.create(loc, resultTy, array, dim, kind); + return {hlfir::EntityWithAttributes{countOp.getResult()}}; + } // TODO add hlfir operations for other transformational intrinsics here diff --git a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp index 7a00bfae75ed1..21a44c07953b6 100644 --- a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp +++ b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp @@ -509,6 +509,43 @@ mlir::LogicalResult hlfir::AnyOp::verify() { return verifyLogicalReductionOp(this); } +//===----------------------------------------------------------------------===// +// CountOp +//===----------------------------------------------------------------------===// + +mlir::LogicalResult hlfir::CountOp::verify() { + mlir::Operation *op = getOperation(); + + auto results = op->getResultTypes(); + assert(results.size() == 1); + mlir::Value mask = getMask(); + mlir::Value dim = getDim(); + + fir::SequenceType maskTy = + hlfir::getFortranElementOrSequenceType(mask.getType()) + .cast(); + llvm::ArrayRef maskShape = maskTy.getShape(); + + mlir::Type resultType = results[0]; + if (auto resultExpr = mlir::dyn_cast_or_null(resultType)) { + if (maskShape.size() > 1 && dim != nullptr) { + if (!resultExpr.isArray()) + return emitOpError("result must be an array"); + + llvm::ArrayRef resultShape = resultExpr.getShape(); + // Result has rank n-1 + if (resultShape.size() != (maskShape.size() - 1)) + return emitOpError("result rank must be one less than MASK"); + } else { + return emitOpError("result must be of numerical scalar type"); + } + } else if (!hlfir::isFortranScalarNumericalType(resultType)) { + return emitOpError("result must be of numerical scalar type"); + } + + return mlir::success(); +} + //===----------------------------------------------------------------------===// // ConcatOp //===----------------------------------------------------------------------===// diff --git a/flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIRIntrinsics.cpp b/flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIRIntrinsics.cpp index b88b30e235a15..d3c604e249a9c 100644 --- a/flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIRIntrinsics.cpp +++ b/flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIRIntrinsics.cpp @@ -262,6 +262,39 @@ using AnyOpConversion = HlfirReductionIntrinsicConversion; using AllOpConversion = HlfirReductionIntrinsicConversion; +struct CountOpConversion : public HlfirIntrinsicConversion { + using HlfirIntrinsicConversion::HlfirIntrinsicConversion; + + mlir::LogicalResult + matchAndRewrite(hlfir::CountOp count, + mlir::PatternRewriter &rewriter) const override { + fir::KindMapping kindMapping{rewriter.getContext()}; + fir::FirOpBuilder builder{rewriter, kindMapping}; + const mlir::Location &loc = count->getLoc(); + + mlir::Type i32 = builder.getI32Type(); + mlir::Type logicalType = fir::LogicalType::get( + builder.getContext(), builder.getKindMap().defaultLogicalKind()); + + llvm::SmallVector inArgs; + inArgs.push_back({count.getMask(), logicalType}); + inArgs.push_back({count.getDim(), i32}); + inArgs.push_back({count.getKind(), i32}); + + auto *argLowering = fir::getIntrinsicArgumentLowering("count"); + llvm::SmallVector args = + lowerArguments(count, inArgs, rewriter, argLowering); + + mlir::Type scalarResultType = hlfir::getFortranElementType(count.getType()); + + auto [resultExv, mustBeFreed] = + fir::genIntrinsicCall(builder, loc, "count", scalarResultType, args); + + processReturnValue(count, resultExv, mustBeFreed, builder, rewriter); + return mlir::success(); + } +}; + struct MatmulOpConversion : public HlfirIntrinsicConversion { using HlfirIntrinsicConversion::HlfirIntrinsicConversion; @@ -405,14 +438,14 @@ class LowerHLFIRIntrinsics patterns.insert(context); + CountOpConversion, DotProductOpConversion>(context); mlir::ConversionTarget target(*context); target.addLegalDialect(); target.addIllegalOp(); + hlfir::AllOp, hlfir::DotProductOp, hlfir::CountOp>(); target.markUnknownOpDynamicallyLegal( [](mlir::Operation *) { return true; }); if (mlir::failed( diff --git a/flang/test/HLFIR/count-lowering.fir b/flang/test/HLFIR/count-lowering.fir new file mode 100644 index 0000000000000..0d9cc34a316eb --- /dev/null +++ b/flang/test/HLFIR/count-lowering.fir @@ -0,0 +1,164 @@ +// Test hlfir.count operation lowering to fir runtime call +// RUN: fir-opt %s -lower-hlfir-intrinsics | FileCheck %s + +func.func @_QPcount1(%arg0: !fir.box>> {fir.bindc_name = "a"}, %arg1: !fir.ref {fir.bindc_name = "s"}) { + %0:2 = hlfir.declare %arg0 {uniq_name = "_QFcount1Ea"} : (!fir.box>>) -> (!fir.box>>, !fir.box>>) + %1:2 = hlfir.declare %arg1 {uniq_name = "_QFcount1Es"} : (!fir.ref) -> (!fir.ref, !fir.ref) + %2 = hlfir.count %0#0 {fastmath = #arith.fastmath} : (!fir.box>>) -> i32 + hlfir.assign %2 to %1#0 : i32, !fir.ref + return +} +// CHECK-LABEL: func.func @_QPcount1( +// CHECK: %[[ARG0:.*]]: !fir.box>> +// CHECK: %[[ARG1:.*]]: !fir.ref +// CHECK-DAG: %[[MASK:.*]]:2 = hlfir.declare %[[ARG0]] +// CHECK-DAG: %[[RES:.*]]:2 = hlfir.declare %[[ARG1]] +// CHECK-DAG: %[[MASK_ARG:.*]] = fir.convert %[[MASK]]#1 : (!fir.box>>) -> !fir.box +// CHECK: %[[RET_ARG:.*]] = fir.call @_FortranACount(%[[MASK_ARG]], %[[LOC_STR:.*]], %[[LOC_N:.*]], %[[C1:.*]]) : (!fir.box, !fir.ref, i32, i32) -> i64 +// CHECK-NEXT: %[[RET:.*]] = fir.convert %[[RET_ARG]] : (i64) -> i32 +// CHECK-NEXT: hlfir.assign %[[RET]] to %[[RES]]#0 : i32, !fir.ref +// CHECK-NEXT: return +// CHECK-NEXT: } + +func.func @_QPcount2(%arg0: !fir.box>> {fir.bindc_name = "a"}, %arg1: !fir.box> {fir.bindc_name = "s"}, %arg2: !fir.ref {fir.bindc_name = "d"}) { + %0:2 = hlfir.declare %arg0 {uniq_name = "_QFcount2Ea"} : (!fir.box>>) -> (!fir.box>>, !fir.box>>) + %1:2 = hlfir.declare %arg2 {uniq_name = "_QFcount2Ed"} : (!fir.ref) -> (!fir.ref, !fir.ref) + %2:2 = hlfir.declare %arg1 {uniq_name = "_QFcount2Es"} : (!fir.box>) -> (!fir.box>, !fir.box>) + %3 = fir.load %1#0 : !fir.ref + %4 = hlfir.count %0#0 dim %3 {fastmath = #arith.fastmath} : (!fir.box>>, i32) -> !hlfir.expr + hlfir.assign %4 to %2#0 : !hlfir.expr, !fir.box> + hlfir.destroy %4 : !hlfir.expr + return +} +// CHECK-LABEL: func.func @_QPcount2( +// CHECK: %[[ARG0:.*]]: !fir.box>> +// CHECK: %[[ARG1:.*]]: !fir.box +// CHECK: %[[ARG2:.*]]: !fir.ref +// CHECK-DAG: %[[MASK:.*]]:2 = hlfir.declare %[[ARG0]] +// CHECK-DAG: %[[DIM_VAR:.*]]:2 = hlfir.declare %[[ARG2]] +// CHECK-DAG: %[[RES:.*]]:2 = hlfir.declare %[[ARG1]] + +// CHECK-DAG: %[[RET_BOX:.*]] = fir.alloca !fir.box>> +// CHECK-DAG: %[[RET_ADDR:.*]] = fir.zero_bits !fir.heap> +// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index +// CHECK-DAG: %[[RET_SHAPE:.*]] = fir.shape %[[C0]] : (index) -> !fir.shape<1> +// CHECK-DAG: %[[RET_EMBOX:.*]] = fir.embox %[[RET_ADDR]](%[[RET_SHAPE]]) +// CHECK-DAG: fir.store %[[RET_EMBOX]] to %[[RET_BOX]] + +// CHECK-DAG: %[[DIM:.*]] = fir.load %[[DIM_VAR]]#0 : !fir.ref +// CHECK-DAG: %[[RET_ARG:.*]] = fir.convert %[[RET_BOX]] +// CHECK-DAG: %[[MASK_ARG:.*]] = fir.convert %[[MASK]]#1 + +// CHECK: %[[NONE:.*]] = fir.call @_FortranACountDim(%[[RET_ARG]], %[[MASK_ARG]], %[[DIM]], %[[LOC_STR:.*]], %[[LOC_N:.*]]) : (!fir.ref>, !fir.box, i32, i32, !fir.ref, i32) -> none +// CHECK: %[[RET:.*]] = fir.load %[[RET_BOX]] +// CHECK: %[[BOX_DIMS:.*]]:3 = fir.box_dims %[[RET]] +// CHECK-NEXT: %[[ADDR:.*]] = fir.box_addr %[[RET]] +// CHECK-NEXT: %[[SHIFT:.*]] = fir.shape_shift %[[BOX_DIMS]]#0, %[[BOX_DIMS]]#1 +// CHECK-NEXT: %[[TMP:.*]]:2 = hlfir.declare %[[ADDR]](%[[SHIFT]]) {uniq_name = ".tmp.intrinsic_result"} +// CHECK: %[[TRUE:.*]] = arith.constant true +// CHECK: %[[EXPR:.*]] = hlfir.as_expr %[[TMP]]#0 move %[[TRUE]] : (!fir.box>, i1) -> !hlfir.expr +// CHECK: hlfir.assign %[[EXPR]] to %[[RES]]#0 +// CHECK: hlfir.destroy %[[EXPR]] +// CHECK-NEXT: return +// CHECK-NEXT: } + +func.func @_QPcount3(%arg0: !fir.ref> {fir.bindc_name = "s"}) { + %0 = fir.address_of(@_QFcount3Ea) : !fir.ref>> + %c2 = arith.constant 2 : index + %c2_0 = arith.constant 2 : index + %1 = fir.shape %c2, %c2_0 : (index, index) -> !fir.shape<2> + %2:2 = hlfir.declare %0(%1) {uniq_name = "_QFcount3Ea"} : (!fir.ref>>, !fir.shape<2>) -> (!fir.ref>>, !fir.ref>>) + %c2_1 = arith.constant 2 : index + %3 = fir.shape %c2_1 : (index) -> !fir.shape<1> + %4:2 = hlfir.declare %arg0(%3) {uniq_name = "_QFcount3Es"} : (!fir.ref>, !fir.shape<1>) -> (!fir.ref>, !fir.ref>) + %c1_i32 = arith.constant 1 : i32 + %5 = hlfir.count %2#0 dim %c1_i32 {fastmath = #arith.fastmath} : (!fir.ref>>, i32) -> !hlfir.expr<2xi32> + hlfir.assign %5 to %4#0 : !hlfir.expr<2xi32>, !fir.ref> + hlfir.destroy %5 : !hlfir.expr<2xi32> + return +} +// CHECK-LABEL: func.func @_QPcount3( +// CHECK: %[[ARG0:.*]]: !fir.ref> +// CHECK-DAG: %[[RET_BOX:.*]] = fir.alloca !fir.box>> +// CHECK-DAG: %[[RET_ADDR:.*]] = fir.zero_bits !fir.heap> +// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index +// CHECK-DAG: %[[RET_SHAPE:.*]] = fir.shape %[[C0]] : (index) -> !fir.shape<1> +// CHECK-DAG: %[[RET_EMBOX:.*]] = fir.embox %[[RET_ADDR]](%[[RET_SHAPE]]) +// CHECK-DAG: fir.store %[[RET_EMBOX]] to %[[RET_BOX]] +// CHECK-DAG: %[[RES:.*]]:2 = hlfir.declare %[[ARG0]](%[[RES_SHAPE:.*]]) + +// CHECK-DAG: %[[MASK_ADDR:.*]] = fir.address_of +// CHECK-DAG: %[[MASK_VAR:.*]]:2 = hlfir.declare %[[MASK_ADDR]](%[[MASK_SHAPE:.*]]) +// CHECK-DAG: %[[MASK_BOX:.*]] = fir.embox %[[MASK_VAR]]#1(%[[MASK_SHAPE:.*]]) + +// CHECK-DAG: %[[DIM:.*]] = arith.constant 1 : i32 + +// CHECK-DAG: %[[RET_ARG:.*]] = fir.convert %[[RET_BOX]] +// CHECK-DAG: %[[MASK_ARG:.*]] = fir.convert %[[MASK_BOX]] : (!fir.box>>) -> !fir.box +// CHECK: %[[NONE:.*]] = fir.call @_FortranACountDim(%[[RET_ARG]], %[[MASK_ARG]], %[[DIM]], %[[LOC_STR:.*]], %[[LOC_N:.*]]) +// CHECK: %[[RET:.*]] = fir.load %[[RET_BOX]] +// CHECK: %[[BOX_DIMS:.*]]:3 = fir.box_dims %[[RET]] +// CHECK-NEXT: %[[ADDR:.*]] = fir.box_addr %[[RET]] +// CHECK-NEXT: %[[SHIFT:.*]] = fir.shape_shift %[[BOX_DIMS]]#0, %[[BOX_DIMS]]#1 +// CHECK-NEXT: %[[TMP:.*]]:2 = hlfir.declare %[[ADDR]](%[[SHIFT]]) {uniq_name = ".tmp.intrinsic_result"} +// CHECK: %[[TRUE:.*]] = arith.constant true +// CHECK: %[[EXPR:.*]] = hlfir.as_expr %[[TMP]]#0 move %[[TRUE]] : (!fir.box>, i1) -> !hlfir.expr +// CHECK: hlfir.assign %[[EXPR]] to %[[RES]] +// CHECK: hlfir.destroy %[[EXPR]] +// CHECK-NEXT: return +// CHECK-NEXT: } + +func.func @_QPcount4(%arg0: !fir.box>> {fir.bindc_name = "a"}, %arg1: !fir.box> {fir.bindc_name = "s"}, %arg2: !fir.ref {fir.bindc_name = "d"}) { + %0:2 = hlfir.declare %arg0 {uniq_name = "_QFcount4Ea"} : (!fir.box>>) -> (!fir.box>>, !fir.box>>) + %1:2 = hlfir.declare %arg2 {uniq_name = "_QFcount4Ed"} : (!fir.ref) -> (!fir.ref, !fir.ref) + %2:2 = hlfir.declare %arg1 {uniq_name = "_QFcount4Es"} : (!fir.box>) -> (!fir.box>, !fir.box>) + %c8_i32 = arith.constant 8 : i32 + %3 = fir.load %1#0 : !fir.ref + %4 = hlfir.count %0#0 dim %3 kind %c8_i32 {fastmath = #arith.fastmath} : (!fir.box>>, i32, i32) -> !hlfir.expr + %5 = hlfir.shape_of %4 : (!hlfir.expr) -> !fir.shape<1> + %6 = hlfir.elemental %5 : (!fir.shape<1>) -> !hlfir.expr { + ^bb0(%arg3: index): + %7 = hlfir.apply %4, %arg3 : (!hlfir.expr, index) -> i64 + %8 = fir.convert %7 : (i64) -> i32 + hlfir.yield_element %8 : i32 + } + hlfir.assign %6 to %2#0 : !hlfir.expr, !fir.box> + hlfir.destroy %6 : !hlfir.expr + hlfir.destroy %4 : !hlfir.expr + return +} +// CHECK-LABEL: func.func @_QPcount4( +// CHECK: %[[ARG0:.*]]: !fir.box>> +// CHECK: %[[ARG1:.*]]: !fir.box +// CHECK: %[[ARG2:.*]]: !fir.ref +// CHECK-DAG: %[[MASK:.*]]:2 = hlfir.declare %[[ARG0]] +// CHECK-DAG: %[[DIM_VAR:.*]]:2 = hlfir.declare %[[ARG2]] +// CHECK-DAG: %[[RES:.*]]:2 = hlfir.declare %[[ARG1]] + +// CHECK-DAG: %[[RET_BOX:.*]] = fir.alloca !fir.box>> +// CHECK-DAG: %[[RET_ADDR:.*]] = fir.zero_bits !fir.heap> +// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index +// CHECK-DAG: %[[RET_SHAPE:.*]] = fir.shape %[[C0]] : (index) -> !fir.shape<1> +// CHECK-DAG: %[[RET_EMBOX:.*]] = fir.embox %[[RET_ADDR]](%[[RET_SHAPE]]) +// CHECK-DAG: fir.store %[[RET_EMBOX]] to %[[RET_BOX]] + +// CHECK-DAG: %[[DIM:.*]] = fir.load %[[DIM_VAR]]#0 : !fir.ref +// CHECK-DAG: %[[KIND:.*]] = arith.constant 8 : i32 +// CHECK-DAG: %[[RET_ARG:.*]] = fir.convert %[[RET_BOX]] +// CHECK-DAG: %[[MASK_ARG:.*]] = fir.convert %[[MASK]]#1 + +// CHECK: %[[NONE:.*]] = fir.call @_FortranACountDim(%[[RET_ARG]], %[[MASK_ARG]], %[[DIM]], %[[KIND]], %[[LOC_STR:.*]], %[[LOC_N:.*]]) : (!fir.ref>, !fir.box, i32, i32, !fir.ref, i32) -> none +// CHECK: %[[RET:.*]] = fir.load %[[RET_BOX]] +// CHECK: %[[BOX_DIMS:.*]]:3 = fir.box_dims %[[RET]] +// CHECK-NEXT: %[[ADDR:.*]] = fir.box_addr %[[RET]] +// CHECK-NEXT: %[[SHIFT:.*]] = fir.shape_shift %[[BOX_DIMS]]#0, %[[BOX_DIMS]]#1 +// CHECK-NEXT: %[[TMP:.*]]:2 = hlfir.declare %[[ADDR]](%[[SHIFT]]) {uniq_name = ".tmp.intrinsic_result"} +// CHECK: %[[TRUE:.*]] = arith.constant true +// CHECK: %[[EXPR:.*]] = hlfir.as_expr %[[TMP]]#0 move %[[TRUE]] : (!fir.box>, i1) -> !hlfir.expr +// CHECK-NEXT: %[[OUT_SHAPE:.*]] = hlfir.shape_of %[[EXPR]] +// CHECK-NEXT: %[[OUT:.*]] = hlfir.elemental %[[OUT_SHAPE]] : (!fir.shape<1>) -> !hlfir.expr +// CHECK-DAG: hlfir.assign %[[OUT]] to %[[RES]]#0 +// CHECK-NEXT: hlfir.destroy %[[OUT]] : !hlfir.expr +// CHECK-NEXT: hlfir.destroy %[[EXPR]] : !hlfir.expr +// CHECK-NEXT: return +// CHECK-NEXT: } diff --git a/flang/test/HLFIR/count.fir b/flang/test/HLFIR/count.fir new file mode 100644 index 0000000000000..c25f9b94124af --- /dev/null +++ b/flang/test/HLFIR/count.fir @@ -0,0 +1,83 @@ +// Test hlfir.count operation parse, verify (no errors), and unparse + +// RUN: fir-opt %s | fir-opt | FileCheck %s + +// mask is an expression of known shape +func.func @count0(%arg0: !hlfir.expr<2x!fir.logical<4>>) { + %count = hlfir.count %arg0 : (!hlfir.expr<2x!fir.logical<4>>) -> i32 + return +} +// CHECK: func.func @count0(%[[ARRAY:.*]]: !hlfir.expr<2x!fir.logical<4>>) { +// CHECK-NEXT: %[[COUNT:.*]] = hlfir.count %[[ARRAY]] : (!hlfir.expr<2x!fir.logical<4>>) -> i32 +// CHECK-NEXT: return +// CHECK-NEXT: } + +// mask is an expression of assumed shape +func.func @count1(%arg0: !hlfir.expr>) { + %count = hlfir.count %arg0 : (!hlfir.expr>) -> i32 + return +} +// CHECK: func.func @count1(%[[ARRAY:.*]]: !hlfir.expr>) { +// CHECK-NEXT: %[[COUNT:.*]] = hlfir.count %[[ARRAY]] : (!hlfir.expr>) -> i32 +// CHECK-NEXT: return +// CHECK-NEXT: } + +// mask is a boxed array +func.func @count2(%arg0: !fir.box>>) { + %count = hlfir.count %arg0 : (!fir.box>>) -> i32 + return +} +// CHECK: func.func @count2(%[[ARRAY:.*]]: !fir.box>>) { +// CHECK-NEXT: %[[COUNT:.*]] = hlfir.count %[[ARRAY]] : (!fir.box>>) -> i32 +// CHECK-NEXT: return +// CHECK-NEXT: } + +// mask is an assumed shape boxed array +func.func @count3(%arg0: !fir.box>>) { + %count = hlfir.count %arg0 : (!fir.box>>) -> i32 + return +} +// CHECK: func.func @count3(%[[ARRAY:.*]]: !fir.box>>) { +// CHECK-NEXT: %[[COUNT:.*]] = hlfir.count %[[ARRAY]] : (!fir.box>>) -> i32 +// CHECK-NEXT: return +// CHECK-NEXT: } + +// mask is a 2-dimensional array +func.func @count4(%arg0: !fir.box>>){ + %count = hlfir.count %arg0 : (!fir.box>>) -> i32 + return +} +// CHECK: func.func @count4(%[[ARRAY:.*]]: !fir.box>>) { +// CHECK-NEXT: %[[COUNT:.*]] = hlfir.count %[[ARRAY]] : (!fir.box>>) -> i32 +// CHECK-NEXT: return +// CHECK-NEXT: } + +// mask and dim argument +func.func @count5(%arg0: !fir.box>>, %arg1: i32) { + %count = hlfir.count %arg0 dim %arg1 : (!fir.box>>, i32) -> i32 + return +} +// CHECK: func.func @count5(%[[ARRAY:.*]]: !fir.box>>, %[[DIM:.*]]: i32) { +// CHECK-NEXT: %[[COUNT:.*]] = hlfir.count %[[ARRAY]] dim %[[DIM]] : (!fir.box>>, i32) -> i32 +// CHECK-NEXT: return +// CHECK-NEXT: } + +// mask is a 2 dimensional array with dim +func.func @count6(%arg0: !fir.box>>, %arg1: i32) { + %count = hlfir.count %arg0 dim %arg1 : (!fir.box>>, i32) -> i32 + return +} +// CHECK: func.func @count6(%[[ARRAY:.*]]: !fir.box>>, %[[DIM:.*]]: i32) { +// CHECK-NEXT: %[[ANY:.*]] = hlfir.count %[[ARRAY]] dim %[[DIM]] : (!fir.box>>, i32) -> i32 +// CHECK-NEXT: return +// CHECK-NEXT: } + +// mask is of a different logical type +func.func @count7(%arg0: !fir.box>>, %arg1: i32) { + %count = hlfir.count %arg0 dim %arg1 : (!fir.box>>, i32) -> i32 + return +} +// CHECK: func.func @count7(%[[ARRAY:.*]]: !fir.box>>, %[[DIM:.*]]: i32) { +// CHECK-NEXT: %[[ANY:.*]] = hlfir.count %[[ARRAY]] dim %[[DIM]] : (!fir.box>>, i32) -> i32 +// CHECK-NEXT: return +// CHECK-NEXT: } \ No newline at end of file diff --git a/flang/test/HLFIR/invalid.fir b/flang/test/HLFIR/invalid.fir index 01bccdf80428b..6db1b79e0818f 100644 --- a/flang/test/HLFIR/invalid.fir +++ b/flang/test/HLFIR/invalid.fir @@ -368,6 +368,30 @@ func.func @bad_all6(%arg0: !hlfir.expr>) { %0 = hlfir.all %arg0 : (!hlfir.expr>) -> !hlfir.expr> } +// ----- +func.func @bad_count1(%arg0: !hlfir.expr>, %arg1: i32) { + // expected-error@+1 {{'hlfir.count' op result must be an array}} + %0 = hlfir.count %arg0 dim %arg1 : (!hlfir.expr>, i32) -> !hlfir.expr +} + +// ----- +func.func @bad_count2(%arg0: !hlfir.expr>, %arg1: i32){ + // expected-error@+1 {{'hlfir.count' op result rank must be one less than MASK}} + %0 = hlfir.count %arg0 dim %arg1 : (!hlfir.expr>, i32) -> !hlfir.expr> +} + +// ----- +func.func @bad_count3(%arg0: !hlfir.expr>, %arg1: i32) { + // expected-error@+1 {{'hlfir.count' op result must be of numerical scalar type}} + %0 = hlfir.count %arg0 dim %arg1 : (!hlfir.expr>, i32) -> !hlfir.expr +} + +// ----- +func.func @bad_count4(%arg0: !hlfir.expr>, %arg1: i32) { + // expected-error@+1 {{'hlfir.count' op result must be of numerical scalar type}} + %0 = hlfir.count %arg0 dim %arg1 : (!hlfir.expr>, i32) -> !fir.logical<4> +} + // ----- func.func @bad_product1(%arg0: !hlfir.expr, %arg1: i32, %arg2: !fir.box>) { // expected-error@+1 {{'hlfir.product' op result must have the same element type as ARRAY argument}} diff --git a/flang/test/Lower/HLFIR/count.f90 b/flang/test/Lower/HLFIR/count.f90 new file mode 100644 index 0000000000000..25c74841514e9 --- /dev/null +++ b/flang/test/Lower/HLFIR/count.f90 @@ -0,0 +1,82 @@ +! Test lowering of COUNT intrinsic to HLFIR +! RUN: bbc -emit-hlfir -o - %s 2>&1 | FileCheck %s + +! simple 1 argument COUNT +subroutine count1(a, s) + logical :: a(:) + integer :: s + s = COUNT(a) +end subroutine +! CHECK-LABEL: func.func @_QPcount1( +! CHECK: %[[ARG0:.*]]: !fir.box>> +! CHECK: %[[ARG1:.*]]: !fir.ref +! CHECK-DAG: %[[MASK:.*]]:2 = hlfir.declare %[[ARG0]] +! CHECK-DAG: %[[OUT:.*]]:2 = hlfir.declare %[[ARG1]] +! CHECK-NEXT: %[[EXPR:.*]] = hlfir.count %[[MASK]]#0 : (!fir.box>>) -> i32 +! CHECK-NEXT: hlfir.assign %[[EXPR]] to %[[OUT]]#0 : i32, !fir.ref +! CHECK-NEXT: return +! CHECK-NEXT: } + +! count with by-ref DIM argument +subroutine count2(a, s, d) + logical :: a(:,:) + integer :: s(:), d + s = COUNT(a, d) +end subroutine +! CHECK-LABEL: func.func @_QPcount2( +! CHECK: %[[ARG0:.*]]: !fir.box>> +! CHECK: %[[ARG1:.*]]: !fir.box> +! CHECK: %[[ARG2:.*]]: !fir.ref +! CHECK-DAG: %[[MASK:.*]]:2 = hlfir.declare %[[ARG0]] +! CHECK-DAG: %[[DIM_REF:.*]]:2 = hlfir.declare %[[ARG2]] +! CHECK-DAG: %[[OUT:.*]]:2 = hlfir.declare %[[ARG1]] +! CHECK-DAG: %[[DIM:.*]] = fir.load %[[DIM_REF]]#0 : !fir.ref +! CHECK-DAG: %[[EXPR:.*]] = hlfir.count %[[MASK]]#0 dim %[[DIM]] : (!fir.box>>, i32) -> !hlfir.expr +! CHECK-NEXT: hlfir.assign %[[EXPR]] to %[[OUT]]#0 : !hlfir.expr, !fir.box> +! CHECK-NEXT: hlfir.destroy %[[EXPR]] : !hlfir.expr +! CHECK-NEXT: return +! CHECK-NEXT: } + +! count with DIM argument by-val, mask isn't boxed +subroutine count3(s) + integer :: s(2) + logical :: a(2,2) = reshape((/.true.,.false.,.true.,.false./), shape(a)) + s = COUNT(a, 1) +end subroutine +! CHECK-LABEL: func.func @_QPcount3( +! CHECK: %[[ARG0:.*]]: !fir.ref> +! CHECK-DAG: %[[ADDR:.*]] = fir.address_of{{.*}} : !fir.ref>> +! CHECK-DAG: %[[MASK_SHAPE:.*]] = fir.shape {{.*}} -> !fir.shape<2> +! CHECK-DAG: %[[MASK:.*]]:2 = hlfir.declare %[[ADDR]](%[[MASK_SHAPE]]) +! CHECK-DAG: %[[OUT_SHAPE:.*]] = fir.shape {{.*}} -> !fir.shape<1> +! CHECK-DAG: %[[OUT:.*]]:2 = hlfir.declare %[[ARG0]](%[[OUT_SHAPE]]) +! CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i32 +! CHECK-DAG: %[[EXPR:.*]] = hlfir.count %[[MASK]]#0 dim %[[C1]] : (!fir.ref>>, i32) -> !hlfir.expr<2xi32> +! CHECK-DAG: hlfir.assign %[[EXPR]] to %[[OUT]]#0 : !hlfir.expr<2xi32>, !fir.ref> +! CHECK-NEXT: hlfir.destroy %[[EXPR]] : !hlfir.expr<2xi32> +! CHECK-NEXT: return +! CHECK-NEXT: } + +! count with dim and kind arguments +subroutine count4(a, s, d) + logical :: a(:,:) + integer :: s(:), d + s = COUNT(a, d, 8) +end subroutine +! CHECK-LABEL: func.func @_QPcount4( +! CHECK: %[[ARG0:.*]]: !fir.box>> +! CHECK: %[[ARG1:.*]]: !fir.box> +! CHECK: %[[ARG2:.*]]: !fir.ref +! CHECK-DAG: %[[MASK:.*]]:2 = hlfir.declare %[[ARG0]] +! CHECK-DAG: %[[DIM_REF:.*]]:2 = hlfir.declare %[[ARG2]] +! CHECK-DAG: %[[OUT:.*]]:2 = hlfir.declare %[[ARG1]] +! CHECK-DAG: %[[C8:.*]] = arith.constant 8 : i32 +! CHECK-DAG: %[[DIM:.*]] = fir.load %[[DIM_REF]]#0 : !fir.ref +! CHECK-DAG: %[[EXPR:.*]] = hlfir.count %[[MASK]]#0 dim %[[DIM]] kind %[[C8]] : (!fir.box>>, i32, i32) -> !hlfir.expr +! CHECK-DAG: %[[RES_SHAPE:.*]] = hlfir.shape_of %[[EXPR]] +! CHECK-DAG: %[[RES:.*]] = hlfir.elemental %[[RES_SHAPE]] : (!fir.shape<1>) -> !hlfir.expr +! CHECK-DAG: hlfir.assign %[[RES]] to %[[OUT]]#0 +! CHECK-NEXT: hlfir.destroy %[[RES]] : !hlfir.expr +! CHECK-NEXT: hlfir.destroy %[[EXPR]] : !hlfir.expr +! CHECK-NEXT: return +! CHECK-NEXT: }