[flang] Fix replaceAllUsesWith API violations (1/N) (#154698)

`replaceAllUsesWith` is not safe to use in a dialect conversion and will
be deactivated soon (#154112). Fix commit fixes some API violations.
Also some general improvements.

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

@@ -1900,8 +1900,7 @@ hlfir::ShapeOfOp::canonicalize(ShapeOfOp shapeOf,
     // shape information is not available at compile time
     return llvm::LogicalResult::failure();
 
-  rewriter.replaceAllUsesWith(shapeOf.getResult(), shape);
-  rewriter.eraseOp(shapeOf);
+  rewriter.replaceOp(shapeOf, shape);
   return llvm::LogicalResult::success();
 }
 

flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp

@@ -455,12 +455,8 @@ struct AssociateOpConversion
 
       mlir::Type associateHlfirVarType = associate.getResultTypes()[0];
       hlfirVar = adjustVar(hlfirVar, associateHlfirVarType);
-      associate.getResult(0).replaceAllUsesWith(hlfirVar);
-
       mlir::Type associateFirVarType = associate.getResultTypes()[1];
       firVar = adjustVar(firVar, associateFirVarType);
-      associate.getResult(1).replaceAllUsesWith(firVar);
-      associate.getResult(2).replaceAllUsesWith(flag);
       // FIXME: note that the AssociateOp that is being erased
       // here will continue to be a user of the original Source
       // operand (e.g. a result of hlfir.elemental), because
@@ -472,7 +468,7 @@ struct AssociateOpConversion
       // the conversions, so that we can analyze HLFIR in its
       // original form and decide which of the AssociateOp
       // users of hlfir.expr can reuse the buffer (if it can).
-      rewriter.eraseOp(associate);
+      rewriter.replaceOp(associate, {hlfirVar, firVar, flag});
     };
 
     // If this is the last use of the expression value and this is an hlfir.expr

flang/lib/Optimizer/HLFIR/Transforms/InlineElementals.cpp

@@ -101,9 +101,8 @@ public:
         elemental.getLoc(), builder, elemental, apply.getIndices());
 
     // remove the old elemental and all of the bookkeeping
-    rewriter.replaceAllUsesWith(apply.getResult(), yield.getElementValue());
+    rewriter.replaceOp(apply, {yield.getElementValue()});
     rewriter.eraseOp(yield);
-    rewriter.eraseOp(apply);
     rewriter.eraseOp(destroy);
     rewriter.eraseOp(elemental);
 

flang/lib/Optimizer/OpenMP/SimdOnly.cpp

@@ -95,16 +95,13 @@ public:
       return mlir::success();
     }
     if (auto mapInfoOp = mlir::dyn_cast<mlir::omp::MapInfoOp>(op)) {
-      mapInfoOp.getResult().replaceAllUsesWith(mapInfoOp.getVarPtr());
-      rewriter.eraseOp(mapInfoOp);
+      rewriter.replaceOp(mapInfoOp, {mapInfoOp.getVarPtr()});
       return mlir::success();
     }
 
     // Might be leftover after parse tree rewriting
     if (auto threadPrivateOp = mlir::dyn_cast<mlir::omp::ThreadprivateOp>(op)) {
-      threadPrivateOp.getTlsAddr().replaceAllUsesWith(
-          threadPrivateOp.getSymAddr());
-      rewriter.eraseOp(threadPrivateOp);
+      rewriter.replaceOp(threadPrivateOp, {threadPrivateOp.getSymAddr()});
       return mlir::success();
     }
 

flang/lib/Optimizer/Transforms/AffineDemotion.cpp

@@ -117,10 +117,7 @@ public:
                                                       op.getValue());
           return success();
         }
-      rewriter.startOpModification(op->getParentOp());
-      op.getResult().replaceAllUsesWith(op.getValue());
-      rewriter.finalizeOpModification(op->getParentOp());
-      rewriter.eraseOp(op);
+      rewriter.replaceOp(op, op.getValue());
     }
     return success();
   }

flang/lib/Optimizer/Transforms/ArrayValueCopy.cpp

@@ -1264,7 +1264,6 @@ public:
     auto lhsEltRefType = toRefType(update.getMerge().getType());
     auto [_, lhsLoadResult] = materializeAssignment(
         loc, rewriter, update, assignElement, lhsEltRefType);
-    update.replaceAllUsesWith(lhsLoadResult);
     rewriter.replaceOp(update, lhsLoadResult);
     return mlir::success();
   }
@@ -1287,7 +1286,6 @@ public:
     auto lhsEltRefType = modify.getResult(0).getType();
     auto [lhsEltCoor, lhsLoadResult] = materializeAssignment(
         loc, rewriter, modify, assignElement, lhsEltRefType);
-    modify.replaceAllUsesWith(mlir::ValueRange{lhsEltCoor, lhsLoadResult});
     rewriter.replaceOp(modify, mlir::ValueRange{lhsEltCoor, lhsLoadResult});
     return mlir::success();
   }
@@ -1339,7 +1337,6 @@ public:
       // This array_access is associated with an array_amend and there is a
       // conflict. Make a copy to store into.
       auto result = referenceToClone(loc, rewriter, access);
-      access.replaceAllUsesWith(result);
       rewriter.replaceOp(access, result);
       return mlir::success();
     }

flang/lib/Optimizer/Transforms/FIRToSCF.cpp

@@ -9,7 +9,7 @@
 #include "flang/Optimizer/Dialect/FIRDialect.h"
 #include "flang/Optimizer/Transforms/Passes.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
-#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/WalkPatternRewriteDriver.h"
 
 namespace fir {
 #define GEN_PASS_DEF_FIRTOSCFPASS
@@ -201,12 +201,7 @@ void FIRToSCFPass::runOnOperation() {
   mlir::RewritePatternSet patterns(&getContext());
   patterns.add<DoLoopConversion, IterWhileConversion, IfConversion>(
       patterns.getContext());
-  mlir::ConversionTarget target(getContext());
-  target.addIllegalOp<fir::DoLoopOp, fir::IterWhileOp, fir::IfOp>();
-  target.markUnknownOpDynamicallyLegal([](mlir::Operation *) { return true; });
-  if (failed(
-          applyPartialConversion(getOperation(), target, std::move(patterns))))
-    signalPassFailure();
+  walkAndApplyPatterns(getOperation(), std::move(patterns));
 }
 
 std::unique_ptr<mlir::Pass> fir::createFIRToSCFPass() {

flang/lib/Optimizer/Transforms/OptimizeArrayRepacking.cpp

@@ -58,8 +58,7 @@ PackingOfContiguous::matchAndRewrite(fir::PackArrayOp op,
                                      mlir::PatternRewriter &rewriter) const {
   mlir::Value box = op.getArray();
   if (hlfir::isSimplyContiguous(box, !op.getInnermost())) {
-    rewriter.replaceAllUsesWith(op, box);
-    rewriter.eraseOp(op);
+    rewriter.replaceOp(op, box);
     return mlir::success();
   }
   return mlir::failure();

flang/lib/Optimizer/Transforms/StackArrays.cpp

@@ -600,10 +600,7 @@ AllocMemConversion::matchAndRewrite(fir::AllocMemOp allocmem,
   // replace references to heap allocation with references to stack allocation
   mlir::Value newValue = convertAllocationType(
       rewriter, allocmem.getLoc(), allocmem.getResult(), alloca->getResult());
-  rewriter.replaceAllUsesWith(allocmem.getResult(), newValue);
-
-  // remove allocmem operation
-  rewriter.eraseOp(allocmem.getOperation());
+  rewriter.replaceOp(allocmem, newValue);
 
   return mlir::success();
 }