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

[jeanPerier]

Summary

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.

[clementval]

LGTM

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

 
 //===----------------------------------------------------------------------===//
 // ForallIndexOp
 //===----------------------------------------------------------------------===//
 
 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();
-}
 
 #include "flang/Optimizer/HLFIR/HLFIROpInterfaces.cpp.inc"
 #define GET_OP_CLASSES
 #include "flang/Optimizer/HLFIR/HLFIROps.cpp.inc"

flang/test/HLFIR/forall-index.fir

hlfir.region_assign {
       %ival = fir.load %i : !fir.ref<i64>
       %yi = hlfir.designate %y(%ival) : (!fir.ref<!fir.array<10xf32>>, i64) -> !fir.ref<f32>
       hlfir.yield %yi : !fir.ref<f32>
     } to {
       %ival = fir.load %i : !fir.ref<i64>
       %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
     } to {
       %ival = fir.load %i : !fir.ref<i64>
       %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_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>

flang/test/Lower/HLFIR/forall.f90

 ! CHECK:  }  (%[[VAL_9:.*]]: i64) {
 ! CHECK:    %[[VAL_10:.*]] = hlfir.forall_index "i" %[[VAL_9]] : (i64) -> !fir.ref<i64>
 ! CHECK:    hlfir.forall_mask {
 ! CHECK:      %[[VAL_11:.*]] = fir.call @_QPpredicate(%[[VAL_10]]) fastmath<contract> : (!fir.ref<i64>) -> !fir.logical<4>
 ! CHECK:      %[[VAL_12:.*]] = fir.convert %[[VAL_11]] : (!fir.logical<4>) -> i1
 ! 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:    }
 ! CHECK:  }
 
 subroutine test_forall_several_indices()
   use forall_defs
   ! Test outer forall controls are lowered outside.