[flang][hlfir] Only canonicalize forall_index if it can be erased

It seems the canonicalization was not correct: it cannot return that
it failed if it did modify the IR.
This was exposed by a new MLIR sanity check added in
https://reviews.llvm.org/D144552.
I am not sure it is legit to return success if the operation being
canonicalized is not modified either. So only remove the loads if
they are the only uses of the forall_index.

Should fix (intermittent?) bot failures like
https://lab.llvm.org/buildbot/#/builders/179/builds/6251
since the new MLIR check was added.

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

diff --git a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
index 9619eb0..4547c42 100644 (file)
--- a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
+++ b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
@@ -1309,19 +1309,20 @@ mlir::LogicalResult hlfir::ElseWhereOp::verify() {
 mlir::LogicalResult
 hlfir::ForallIndexOp::canonicalize(hlfir::ForallIndexOp indexOp,
                                    mlir::PatternRewriter &rewriter) {
+  for (mlir::Operation *user : indexOp->getResult(0).getUsers())
+    if (!mlir::isa<fir::LoadOp>(user))
+      return mlir::failure();
+
   auto insertPt = rewriter.saveInsertionPoint();
   for (mlir::Operation *user : indexOp->getResult(0).getUsers())
-    if (auto loadOp = mlir::dyn_cast_or_null<fir::LoadOp>(user)) {
+    if (auto loadOp = mlir::dyn_cast<fir::LoadOp>(user)) {
       rewriter.setInsertionPoint(loadOp);
       rewriter.replaceOpWithNewOp<fir::ConvertOp>(
           user, loadOp.getResult().getType(), indexOp.getIndex());
     }
   rewriter.restoreInsertionPoint(insertPt);
-  if (indexOp.use_empty()) {
-    rewriter.eraseOp(indexOp);
-    return mlir::success();
-  }
-  return mlir::failure();
+  rewriter.eraseOp(indexOp);
+  return mlir::success();
 }
 
 #include "flang/Optimizer/HLFIR/HLFIROpInterfaces.cpp.inc"

diff --git a/flang/test/HLFIR/forall-index.fir b/flang/test/HLFIR/forall-index.fir
index 02c5f41..814d845 100644 (file)
--- a/flang/test/HLFIR/forall-index.fir
+++ b/flang/test/HLFIR/forall-index.fir
@@ -21,6 +21,32 @@ func.func @forall_index(%x: !fir.ref<!fir.array<10xf32>>, %y: !fir.ref<!fir.arra
       %xi = hlfir.designate %x(%ival) : (!fir.ref<!fir.array<10xf32>>, i64) -> !fir.ref<f32>
       hlfir.yield %xi : !fir.ref<f32>
     }
+  }
+  return
+}
+// CHECK-LABEL:  func.func @forall_index(
+// CHECK:           hlfir.forall lb {
+// CHECK:           } ub {
+// CHECK:           }  (%[[VAL_4:.*]]: i64) {
+// CHECK:             hlfir.forall_index "i" %[[VAL_4]] : (i64) -> !fir.ref<i64>
+
+// CANONICALIZATION-LABEL:  func.func @forall_index(
+// CANONICALIZATION:      hlfir.forall lb {
+// CANONICALIZATION:      } ub {
+// CANONICALIZATION:      }  (%[[VAL_4:.*]]: i64) {
+// CANONICALIZATION-NOT:    hlfir.forall_index
+// CANONICALIZATION:        hlfir.designate %{{.*}} (%[[VAL_4]])  : (!fir.ref<!fir.array<10xf32>>, i64) -> !fir.ref<f32>
+// CANONICALIZATION:        hlfir.designate %{{.*}} (%[[VAL_4]])  : (!fir.ref<!fir.array<10xf32>>, i64) -> !fir.ref<f32>
+
+func.func @forall_index_do_not_canonicalize(%x: !fir.ref<!fir.array<10xf32>>, %y: !fir.ref<!fir.array<10xf32>>) {
+  %c1 = arith.constant 1 : index
+  %c10 = arith.constant 10 : index
+  hlfir.forall lb {
+    hlfir.yield %c1 : index
+  } ub {
+    hlfir.yield %c10 : index
+  } (%arg0: i64) {
+    %i = hlfir.forall_index "i" %arg0 : (i64) -> !fir.ref<i64>
     hlfir.region_assign {
       %res = fir.call @taking_address(%i) : (!fir.ref<i64>) -> f32
       hlfir.yield %res : f32
@@ -32,13 +58,14 @@ func.func @forall_index(%x: !fir.ref<!fir.array<10xf32>>, %y: !fir.ref<!fir.arra
   }
   return
 }
+// CHECK-LABEL:  func.func @forall_index_do_not_canonicalize(
 // CHECK:           hlfir.forall lb {
 // CHECK:           } ub {
 // CHECK:           }  (%[[VAL_4:.*]]: i64) {
 // CHECK:             hlfir.forall_index "i" %[[VAL_4]] : (i64) -> !fir.ref<i64>
 
+// CANONICALIZATION-LABEL:  func.func @forall_index_do_not_canonicalize(
 // CANONICALIZATION:  %[[VAL_5:.*]] = hlfir.forall_index "i" %[[VAL_4:.*]] : (i64) -> !fir.ref<i64>
-// CANONICALIZATION:  hlfir.designate %{{.*}} (%[[VAL_4]])  : (!fir.ref<!fir.array<10xf32>>, i64) -> !fir.ref<f32>
-// CANONICALIZATION:  hlfir.designate %{{.*}} (%[[VAL_4]])  : (!fir.ref<!fir.array<10xf32>>, i64) -> !fir.ref<f32>
 // CANONICALIZATION:  fir.call @taking_address(%[[VAL_5]]) : (!fir.ref<i64>) -> f32
-// CANONICALIZATION:  hlfir.designate %{{.*}} (%[[VAL_4]])  : (!fir.ref<!fir.array<10xf32>>, i64) -> !fir.ref<f32>
+// CANONICALIZATION:  %[[VAL_6:.*]] = fir.load %[[VAL_5]] : !fir.ref<i64>
+// CANONICALIZATION:  hlfir.designate %{{.*}} (%[[VAL_6]])  : (!fir.ref<!fir.array<10xf32>>, i64) -> !fir.ref<f32>

diff --git a/flang/test/Lower/HLFIR/forall.f90 b/flang/test/Lower/HLFIR/forall.f90
index 5091dde..1485ceb 100644 (file)
--- a/flang/test/Lower/HLFIR/forall.f90
+++ b/flang/test/Lower/HLFIR/forall.f90
@@ -91,11 +91,13 @@ end subroutine
 ! CHECK:      hlfir.yield %[[VAL_12]] : i1
 ! CHECK:    } do {
 ! CHECK:      hlfir.region_assign {
-! CHECK:        %[[VAL_13:.*]] = hlfir.designate %[[VAL_8]]#0 (%[[VAL_9]])  : (!fir.ref<!fir.array<10xi32>>, i64) -> !fir.ref<i32>
+! CHECK:        %[[I_LOAD:.*]] = fir.load %[[VAL_10]] : !fir.ref<i64>
+! CHECK:        %[[VAL_13:.*]] = hlfir.designate %[[VAL_8]]#0 (%[[I_LOAD]])  : (!fir.ref<!fir.array<10xi32>>, i64) -> !fir.ref<i32>
 ! CHECK:        %[[VAL_14:.*]] = fir.load %[[VAL_13]] : !fir.ref<i32>
 ! CHECK:        hlfir.yield %[[VAL_14]] : i32
 ! CHECK:      } to {
-! CHECK:        %[[VAL_15:.*]] = hlfir.designate %[[VAL_5]]#0 (%[[VAL_9]], %[[VAL_9]])  : (!fir.ref<!fir.array<10x10xi32>>, i64, i64) -> !fir.ref<i32>
+! CHECK:        %[[I_LOAD:.*]] = fir.load %[[VAL_10]] : !fir.ref<i64>
+! CHECK:        %[[VAL_15:.*]] = hlfir.designate %[[VAL_5]]#0 (%[[I_LOAD]], %[[I_LOAD]])  : (!fir.ref<!fir.array<10x10xi32>>, i64, i64) -> !fir.ref<i32>
 ! CHECK:        hlfir.yield %[[VAL_15]] : !fir.ref<i32>
 ! CHECK:      }
 ! CHECK:    }