[flang] fix cg-rewrite DCE (#99653)

cg-rewrite runs regionDCE to get rid of the unused fir.shape/shift/slice
before codegen since those operations have no codegen.
I came across an issue where unreachable code would cause the pass to
fail with `error: loc(...): null operand found`.

It turns out `mlir::RegionDCE` does not work properly in presence of
unreachable code because it delete operations in reachable code that are
unused in reachable code, but still used in unreachable code (like the
constant in the added test case). It seems `mlir::RegionDCE` is always
run after `mlir::eraseUnreachableBlock` outside of this pass.

A solution could be to run `mlir::eraseUnreachableBlock` here or to try
modifying `mlir::RegionDCE`. But the current behavior may be
intentional, and both of these calls are actually quite expensive. For
instance, RegionDCE will does liveness analysis, and removes unused
block arguments, which is way more than what is needed here. I am not
very found of having this rather heavy transformation inside this pass
(they should be run after or before if they matter in the overall
pipeline).

Do a naïve backward deletion of the trivially dead operations instead.
It is cheaper, and works with unreachable code.

2 files changed, 35 insertions, 2 deletions

diff --git a/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp b/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
index c57f514..fdf2402 100644
--- a/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
+++ b/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
@@ -18,8 +18,8 @@
 #include "flang/Optimizer/Dialect/FIROps.h"
 #include "flang/Optimizer/Dialect/FIRType.h"
 #include "flang/Optimizer/Dialect/Support/FIRContext.h"
+#include "mlir/IR/Iterators.h"
 #include "mlir/Transforms/DialectConversion.h"
-#include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Debug.h"
 
@@ -325,6 +325,25 @@ public:
   }
 };
 
+/// Simple DCE to erase fir.shape/shift/slice/unused shape operands after this
+/// pass (fir.shape and like have no codegen).
+/// mlir::RegionDCE is expensive and requires running
+/// mlir::eraseUnreachableBlocks. It does things that are not needed here, like
+/// removing unused block arguments. fir.shape/shift/slice cannot be block
+/// arguments.
+/// This helper does a naive backward walk of the IR. It is not even guaranteed
+/// to walk blocks according to backward dominance, but that is good enough for
+/// what is done here, fir.shape/shift/slice have no usages anymore. The
+/// backward walk allows getting rid of most of the unused operands, it is not a
+/// problem to leave some in the weird cases.
+static void simpleDCE(mlir::RewriterBase &rewriter, mlir::Operation *op) {
+  op->walk<mlir::WalkOrder::PostOrder, mlir::ReverseIterator>(
+      [&](mlir::Operation *subOp) {
+        if (mlir::isOpTriviallyDead(subOp))
+          rewriter.eraseOp(subOp);
+      });
+}
+
 class CodeGenRewrite : public fir::impl::CodeGenRewriteBase<CodeGenRewrite> {
 public:
   using CodeGenRewriteBase<CodeGenRewrite>::CodeGenRewriteBase;
@@ -356,7 +375,7 @@ public:
     }
     // Erase any residual (fir.shape, fir.slice...).
     mlir::IRRewriter rewriter(&context);
-    (void)mlir::runRegionDCE(rewriter, mod->getRegions());
+    simpleDCE(rewriter, mod.getOperation());
   }
 };
 
diff --git a/flang/test/Fir/declare-codegen.fir b/flang/test/Fir/declare-codegen.fir
index c5879fa..fe8d84e 100644
--- a/flang/test/Fir/declare-codegen.fir
+++ b/flang/test/Fir/declare-codegen.fir
@@ -38,3 +38,17 @@ func.func @useless_shape_with_duplicate_extent_operand(%arg0: !fir.ref<!fir.arra
 
 // DECL-LABEL: func.func @useless_shape_with_duplicate_extent_operand(
 // DECL: fircg.ext_declare
+
+// Test DCE does not crash because of unreachable code.
+func.func @unreachable_code(%arg0: !fir.ref<!fir.char<1,10>>) {
+  %c10 = arith.constant 10 : index
+  %2 = fir.declare %arg0 typeparams %c10 {uniq_name = "live_code"} : (!fir.ref<!fir.char<1,10>>, index) -> (!fir.ref<!fir.char<1,10>>)
+  return
+^bb2:  // no predecessors
+  %3 = fir.declare %arg0 typeparams %c10 {uniq_name = "dead_code"} : (!fir.ref<!fir.char<1,10>>, index) -> (!fir.ref<!fir.char<1,10>>)
+  fir.unreachable
+}
+// NODECL-LABEL:   func.func @unreachable_code(
+// NODECL-NOT:       uniq_name = "live_code"
+// DECL-LABEL:    func.func @unreachable_code(
+// DECL:             uniq_name = "live_code"