[flang] fix cg-rewrite DCE

[jeanPerier]

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.

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

@llvm/pr-subscribers-flang-codegen

Author: None (jeanPerier)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/99653.diff

2 Files Affected:

• (modified) flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp (+21-2)
• (modified) flang/test/Fir/declare-codegen.fir (+14)

diff --git a/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp b/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
index c57f5140c29eb..fdf24028c5f9b 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 @@ class DummyScopeOpConversion
   }
 };
 
+/// 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 @@ class CodeGenRewrite : public fir::impl::CodeGenRewriteBase<CodeGenRewrite> {
     }
     // 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 c5879facb1572..fe8d84ef2d19f 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"

[vzakhari]

LGTM. Thank you, Jean!