[flang] Fix simplify intrinsic for count not checking for rank = 0 properly

[SBallantyne]

Summary

Simple fix to check for rank in the same way as other intrinsics to allow
runtime count to take over when dealing with unknown dimension arrays.

Fixes #60356

[kiranchandramohan]

flang/lib/Optimizer/Transforms/SimplifyIntrinsics.cpp

626–628

Could you add a comment, detailing what this check means? (possibly just copying the summary here),
Could you add a test to ensure that simplify logical reduction does not happen for this case?

[SBallantyne]

Add test and comment

[Leporacanthicus]

Looks OK except for the test checking for the wrong "Not".

flang/test/Transforms/simplifyintrinsics.fir

1239

Should be FortranACount without Dim, right? (And two lines down too)

[SBallantyne]

Fix minor issue with tests.

[Leporacanthicus]

LGTM, thanks for fixing!

[kiranchandramohan]

flang/lib/Optimizer/Transforms/SimplifyIntrinsics.cpp

626

Could you check why the rank is 0 for assumed shape arrays?

flang/lib/Optimizer/Transforms/SimplifyIntrinsics.cpp

void SimplifyIntrinsicsPass::simplifyIntOrFloatReduction(
   const mlir::Value &dim = args[3];
   const mlir::Value &mask = args[4];
   // dim is zero when it is absent, which is an implementation
   // detail in the runtime library.
 
   bool dimAndMaskAbsent = isZero(dim) && isOperandAbsent(mask);
   unsigned rank = getDimCount(args[0]);
 
+  // Rank is set to 0 for assumed shape arrays, don't simplify
+  // in these cases
   if (!(dimAndMaskAbsent && rank > 0))
     return;
 
   mlir::Type resultType = call.getResult(0).getType();
 
   if (!resultType.isa<mlir::FloatType>() &&
       !resultType.isa<mlir::IntegerType>())
     return;

 }
 
 void SimplifyIntrinsicsPass::simplifyLogicalReduction(
     fir::CallOp call, const fir::KindMapping &kindMap,
     GenReductionBodyTy genBodyFunc) {
 
   mlir::Operation::operand_range args = call.getArgs();
   const mlir::Value &dim = args[3];
+  unsigned rank = getDimCount(args[0]);
 
-  if (!isZero(dim))
+  // Rank is set to 0 for assumed shape arrays, don't simplify
+  // in these cases
+  if (!(isZero(dim) && rank > 0))
     return;
 
-  unsigned rank = getDimCount(args[0]);
   mlir::SymbolRefAttr callee = call.getCalleeAttr();
 
   fir::FirOpBuilder builder{getSimplificationBuilder(call, kindMap)};
   std::string funcName =
       (mlir::Twine{callee.getLeafReference().getValue(), "x"} +
        mlir::Twine{rank})
           .str();
 

flang/test/Transforms/simplifyintrinsics.fir

 }
 func.func private @_FortranACountDim(!fir.ref<!fir.box<none>>, !fir.box<none>, i32, i32, !fir.ref<i8>, i32) -> none attributes {fir.runtime}
 
 // CHECK-LABEL:   func.func @_QMtestPcount_generate_mask(
 // CHECK-SAME:                                           %[[A:.*]]: !fir.ref<!fir.array<10x10x!fir.logical<4>>> {fir.bindc_name = "mask"}) -> !fir.array<10xi32> {
 // CHECK-NOT        fir.call @_FortranACountDim_simplified({{.*}})
 // CHECK:           %[[RES:.*]] = fir.call @_FortranACountDim({{.*}}) fastmath<contract> : (!fir.ref<!fir.box<none>>, !fir.box<none>, i32, i32, !fir.ref<i8>, i32) -> none
 // CHECK-NOT        fir.call @_FortranACountDim_simplified({{.*}})
+
+// -----
+// Ensure count isn't simplified for unknown dimension arrays
+
+func.func @_QPmc(%arg0: !fir.box<!fir.array<?x?x?x!fir.logical<4>>> {fir.bindc_name = "m"}) -> i32 {
+  %0 = fir.alloca i32 {bindc_name = "mc", uniq_name = "_QFmcEmc"}
+  %c0 = arith.constant 0 : index
+  %1 = fir.address_of(@_QQcl.95529933117f47914fc21e220c1a0896) : !fir.ref<!fir.char<1,58>>
+  %c3_i32 = arith.constant 3 : i32
+  %2 = fir.convert %arg0 : (!fir.box<!fir.array<?x?x?x!fir.logical<4>>>) -> !fir.box<none>
+  %3 = fir.convert %1 : (!fir.ref<!fir.char<1,58>>) -> !fir.ref<i8>
+  %4 = fir.convert %c0 : (index) -> i32
+  %5 = fir.call @_FortranACount(%2, %3, %c3_i32, %4) fastmath<contract> : (!fir.box<none>, !fir.ref<i8>, i32, i32) -> i64
+  %6 = fir.convert %5 : (i64) -> i32
+  fir.store %6 to %0 : !fir.ref<i32>
+  %7 = fir.load %0 : !fir.ref<i32>
+  return %7 : i32
+}
+func.func private @_FortranACount(!fir.box<none>, !fir.ref<i8>, i32, i32) -> i64 attributes {fir.runtime}
+
+// CHECK-LABEL:   func.func @_QPmc(
+// CHECK-SAME:                     %[[VAL_0:.*]]: !fir.box<!fir.array<?x?x?x!fir.logical<4>>> {fir.bindc_name = "m"}) -> i32 {
+// CHECK-NOT        fir.call @_FortranACount_simplified({{.*}})
+// CHECK:           %[[RES:.*]] = fir.call @_FortranACount({{.*}}) fastmath<contract> : (!fir.box<none>, !fir.ref<i8>, i32, i32) -> i64
+// CHECK-NOT        fir.call @_FortranACount_simplified({{.*}})