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

[Mats Petersson via flang-commits]

Author: Sacha Ballantyne
Date: 2023-01-30T12:23:37Z
New Revision: bb94d33aac2a617a762d0662648f8ec496c4c364

URL: https://github.com/llvm/llvm-project/commit/bb94d33aac2a617a762d0662648f8ec496c4c364
DIFF: https://github.com/llvm/llvm-project/commit/bb94d33aac2a617a762d0662648f8ec496c4c364.diff

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

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

Reviewed By: Leporacanthicus

Differential Revision: https://reviews.llvm.org/D142877

Added: 
    

Modified: 
    flang/lib/Optimizer/Transforms/SimplifyIntrinsics.cpp
    flang/test/Transforms/simplifyintrinsics.fir

Removed: 
    


################################################################################
diff  --git a/flang/lib/Optimizer/Transforms/SimplifyIntrinsics.cpp b/flang/lib/Optimizer/Transforms/SimplifyIntrinsics.cpp
index c1f7f39ad356..de65c487a034 100644
--- a/flang/lib/Optimizer/Transforms/SimplifyIntrinsics.cpp
+++ b/flang/lib/Optimizer/Transforms/SimplifyIntrinsics.cpp
@@ -583,6 +583,8 @@ void SimplifyIntrinsicsPass::simplifyIntOrFloatReduction(
   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;
 
@@ -619,11 +621,13 @@ void SimplifyIntrinsicsPass::simplifyLogicalReduction(
 
   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)};

diff  --git a/flang/test/Transforms/simplifyintrinsics.fir b/flang/test/Transforms/simplifyintrinsics.fir
index e4a2eaeb486d..282fafcd4918 100644
--- a/flang/test/Transforms/simplifyintrinsics.fir
+++ b/flang/test/Transforms/simplifyintrinsics.fir
@@ -1214,3 +1214,28 @@ func.func private @_FortranACountDim(!fir.ref<!fir.box<none>>, !fir.box<none>, i
 // 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({{.*}})