[flang-commits] [flang] [flang][codegen] Rebuild compile-time-only block arguments in cg-rewrite (PR #227624)
Kareem Ergawy via flang-commits
flang-commits at lists.llvm.org
Wed Sep 30 06:17:49 PDT 2026
https://github.com/ergawy updated https://github.com/llvm/llvm-project/pull/227624
>From ccc16f4e3ecddef7b3696ba1964070fcedb02e40 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at gmail.com>
Date: Wed, 30 Sep 2026 01:53:16 -0700
Subject: [PATCH] [flang][codegen] Rebuild compile-time-only block arguments in
cg-rewrite
A fir.shape, fir.shape_shift, fir.shift or fir.slice only describes an
array at compile time. cg-rewrite reads one through its defining op and
folds it into the code-gen form, so none of these types has an LLVM
lowering and none is expected to reach codegen.
A pass that merges two blocks differing only in such a value passes it
as a block argument instead, and a block argument has no defining op.
The fold then finds nothing to read: a slice is dropped, leaving a
descriptor for the whole array rather than the section, and a shape
reaches a cast on a null defining op.
Rebuild the value in the block. The operands behind it are integers,
which can be block arguments, so take those as arguments, forward them
along each branch, and build the value from them at the top of the
block. The merge is kept and no block is duplicated.
---
flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp | 116 ++++++++++++++++++
...g-rewrite-compile-time-only-block-args.fir | 73 +++++++++++
2 files changed, 189 insertions(+)
create mode 100644 flang/test/Fir/cg-rewrite-compile-time-only-block-args.fir
diff --git a/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp b/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
index f4ae6841afc12..36430c440f234 100644
--- a/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
+++ b/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
@@ -20,6 +20,7 @@
#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/IR/Iterators.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Transforms/DialectConversion.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/Debug.h"
@@ -434,6 +435,119 @@ static void simpleDCE(mlir::RewriterBase &rewriter, mlir::Operation *op) {
});
}
+/// A fir.shape, fir.shape_shift, fir.shift or fir.slice only describes an
+/// array at compile time. The rewrites above read it through its defining op
+/// and fold it into the code-gen form, so such a value never reaches codegen
+/// and none of these types has an LLVM lowering.
+static bool isCompileTimeOnly(mlir::Type type) {
+ return mlir::isa<fir::ShapeType, fir::ShapeShiftType, fir::ShiftType,
+ fir::SliceType>(type);
+}
+
+/// Return the op \p value is built by, if this pass can rebuild it elsewhere
+/// from its operands. A fir.slice naming components is excluded: its fields
+/// are fir.field values, which are themselves compile-time only and so cannot
+/// be passed along a branch.
+static mlir::Operation *getRebuildableDefiningOp(mlir::Value value) {
+ mlir::Operation *def = value.getDefiningOp();
+ if (!def)
+ return nullptr;
+ if (auto slice = mlir::dyn_cast<fir::SliceOp>(def))
+ return slice.getFields().empty() ? def : nullptr;
+ if (mlir::isa<fir::ShapeOp, fir::ShapeShiftOp, fir::ShiftOp>(def))
+ return def;
+ return nullptr;
+}
+
+/// True if \p lhs and \p rhs can be rebuilt by one op taking one operand list.
+static bool haveSameShape(mlir::Operation *lhs, mlir::Operation *rhs) {
+ if (lhs->getName() != rhs->getName() ||
+ lhs->getNumOperands() != rhs->getNumOperands())
+ return false;
+ return llvm::equal(lhs->getOperandTypes(), rhs->getOperandTypes());
+}
+
+/// Rebuild the compile-time-only block argument \p argIndex of \p block from
+/// its operands, so the rewrites above find a defining op again.
+///
+/// Every branch into the block passes a value of one of those types. The
+/// operands behind it are integers, which can be block arguments, so they are
+/// appended to the block and forwarded along each branch instead, and the
+/// value is rebuilt from them at the top of the block.
+static void rebuildBlockArgument(mlir::Block *block, unsigned argIndex) {
+ mlir::BlockArgument arg = block->getArgument(argIndex);
+
+ // Collect one incoming edge per branch, keeping the operands each carries.
+ // A branch can name the block twice, so each edge is held separately.
+ struct Edge {
+ mlir::BranchOpInterface branch;
+ unsigned successorIndex;
+ };
+ llvm::SmallVector<Edge> edges;
+ mlir::Operation *model = nullptr;
+ for (auto it = block->pred_begin(), e = block->pred_end(); it != e; ++it) {
+ auto branch =
+ mlir::dyn_cast<mlir::BranchOpInterface>((*it)->getTerminator());
+ if (!branch)
+ return;
+ unsigned successorIndex = it.getSuccessorIndex();
+ mlir::SuccessorOperands operands =
+ branch.getSuccessorOperands(successorIndex);
+ // An operand the branch itself produces has no value to read here.
+ if (operands.isOperandProduced(argIndex))
+ return;
+ mlir::Operation *def = getRebuildableDefiningOp(operands[argIndex]);
+ if (!def)
+ return;
+ if (!model)
+ model = def;
+ else if (!haveSameShape(model, def))
+ return;
+ edges.push_back({branch, successorIndex});
+ }
+ if (edges.empty())
+ return;
+
+ // Take the operands as block arguments and forward them along each branch.
+ unsigned firstNewArg = block->getNumArguments();
+ for (auto [type, value] :
+ llvm::zip_equal(model->getOperandTypes(), model->getOperands()))
+ block->addArgument(type, value.getLoc());
+ for (Edge &edge : edges) {
+ mlir::SuccessorOperands operands =
+ edge.branch.getSuccessorOperands(edge.successorIndex);
+ operands.append(
+ getRebuildableDefiningOp(operands[argIndex])->getOperands());
+ }
+
+ // Rebuild the value at the top of the block and drop the old argument.
+ mlir::OpBuilder builder(block, block->begin());
+ mlir::Operation *rebuilt =
+ builder.create(model->getLoc(), model->getName().getIdentifier(),
+ block->getArguments().drop_front(firstNewArg),
+ model->getResultTypes(), model->getAttrs());
+ arg.replaceAllUsesWith(rebuilt->getResult(0));
+ for (Edge &edge : edges)
+ edge.branch.getSuccessorOperands(edge.successorIndex).erase(argIndex);
+ block->eraseArgument(argIndex);
+}
+
+/// Rebuild every compile-time-only block argument in \p op.
+static void rebuildCompileTimeOnlyBlockArguments(mlir::Operation *op) {
+ llvm::SmallVector<mlir::Block *> blocks;
+ op->walk([&](mlir::Block *block) {
+ if (!block->isEntryBlock() &&
+ llvm::any_of(block->getArgumentTypes(), isCompileTimeOnly))
+ blocks.push_back(block);
+ });
+ for (mlir::Block *block : blocks)
+ // Rebuilding drops an argument, so work back to front to keep the
+ // indices of the arguments still to be looked at.
+ for (unsigned i = block->getNumArguments(); i > 0; --i)
+ if (isCompileTimeOnly(block->getArgument(i - 1).getType()))
+ rebuildBlockArgument(block, i - 1);
+}
+
class CodeGenRewrite : public fir::impl::CodeGenRewriteBase<CodeGenRewrite> {
public:
using CodeGenRewriteBase<CodeGenRewrite>::CodeGenRewriteBase;
@@ -441,6 +555,8 @@ class CodeGenRewrite : public fir::impl::CodeGenRewriteBase<CodeGenRewrite> {
void runOnOperation() override final {
mlir::ModuleOp mod = getOperation();
+ rebuildCompileTimeOnlyBlockArguments(mod.getOperation());
+
auto &context = getContext();
mlir::ConversionTarget target(context);
target.addLegalDialect<mlir::arith::ArithDialect, fir::FIROpsDialect,
diff --git a/flang/test/Fir/cg-rewrite-compile-time-only-block-args.fir b/flang/test/Fir/cg-rewrite-compile-time-only-block-args.fir
new file mode 100644
index 0000000000000..f7e8994149277
--- /dev/null
+++ b/flang/test/Fir/cg-rewrite-compile-time-only-block-args.fir
@@ -0,0 +1,73 @@
+// Merging two blocks that differ only in a fir.slice or a fir.shape_shift
+// passes that value as a block argument, leaving it with no defining op for
+// cg-rewrite to read. Check that cg-rewrite takes the operands as block
+// arguments instead and rebuilds the value in the block.
+
+// RUN: fir-opt %s -pass-pipeline='builtin.module(func.func(canonicalize{region-simplify=aggressive}),cg-rewrite)' | FileCheck %s
+// RUN: fir-opt %s -pass-pipeline='builtin.module(func.func(canonicalize{region-simplify=aggressive}),cg-rewrite,fir-to-llvm-ir)' -o /dev/null
+
+func.func private @use(!fir.box<!fir.array<?xf64>>)
+
+// CHECK-LABEL: func.func @slice_block_argument
+// CHECK: cf.cond_br %{{.*}}, ^[[BB1:.*]](%{{.*}}, %{{.*}}, %{{.*}} : index, index, index), ^[[BB2:.*]]
+// CHECK: ^[[BB1]](%[[T0:.*]]: index, %[[T1:.*]]: index, %[[T2:.*]]: index):
+// CHECK: fircg.ext_embox %{{.*}}[%[[T0]], %[[T1]], %[[T2]]]
+// CHECK: ^[[BB2]]:
+// CHECK: cf.cond_br %{{.*}}, ^[[BB1]](%{{.*}}, %{{.*}}, %{{.*}} : index, index, index)
+// CHECK-NOT: !fir.slice
+func.func @slice_block_argument(%cond1: i1, %cond2: i1) {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ %c4 = arith.constant 4 : index
+ %shift = fir.shape_shift %c1, %c4 : (index, index) -> !fir.shapeshift<1>
+ %heap = fir.allocmem !fir.array<?xf64>, %c4
+ %slice1 = fir.slice %c1, %c2, %c1 : (index, index, index) -> !fir.slice<1>
+ %slice2 = fir.slice %c1, %c3, %c1 : (index, index, index) -> !fir.slice<1>
+ cf.cond_br %cond1, ^bb1, ^between
+^bb1:
+ %emb1 = fir.embox %heap(%shift) [%slice1] : (!fir.heap<!fir.array<?xf64>>, !fir.shapeshift<1>, !fir.slice<1>) -> !fir.box<!fir.array<?xf64>>
+ func.call @use(%emb1) : (!fir.box<!fir.array<?xf64>>) -> ()
+ cf.br ^exit
+^between:
+ cf.cond_br %cond2, ^bb2, ^exit
+^bb2:
+ %emb2 = fir.embox %heap(%shift) [%slice2] : (!fir.heap<!fir.array<?xf64>>, !fir.shapeshift<1>, !fir.slice<1>) -> !fir.box<!fir.array<?xf64>>
+ func.call @use(%emb2) : (!fir.box<!fir.array<?xf64>>) -> ()
+ cf.br ^exit
+^exit:
+ fir.freemem %heap : !fir.heap<!fir.array<?xf64>>
+ return
+}
+
+// A fir.shape_shift arriving as a block argument used to reach a dyn_cast on a
+// null defining op.
+
+// CHECK-LABEL: func.func @shapeshift_block_argument
+// CHECK: cf.cond_br %{{.*}}, ^[[SBB1:.*]](%{{.*}}, %{{.*}} : index, index), ^[[SBB2:.*]]
+// CHECK: ^[[SBB1]](%[[ORIGIN:.*]]: index, %[[EXTENT:.*]]: index):
+// CHECK: fircg.ext_embox %{{.*}}(%[[EXTENT]]) origin %[[ORIGIN]]
+// CHECK-NOT: !fir.shapeshift
+func.func @shapeshift_block_argument(%cond1: i1, %cond2: i1) {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ %c4 = arith.constant 4 : index
+ %heap = fir.allocmem !fir.array<?xf64>, %c4
+ %shift1 = fir.shape_shift %c1, %c2 : (index, index) -> !fir.shapeshift<1>
+ %shift2 = fir.shape_shift %c1, %c3 : (index, index) -> !fir.shapeshift<1>
+ cf.cond_br %cond1, ^bb1, ^between
+^bb1:
+ %emb1 = fir.embox %heap(%shift1) : (!fir.heap<!fir.array<?xf64>>, !fir.shapeshift<1>) -> !fir.box<!fir.array<?xf64>>
+ func.call @use(%emb1) : (!fir.box<!fir.array<?xf64>>) -> ()
+ cf.br ^exit
+^between:
+ cf.cond_br %cond2, ^bb2, ^exit
+^bb2:
+ %emb2 = fir.embox %heap(%shift2) : (!fir.heap<!fir.array<?xf64>>, !fir.shapeshift<1>) -> !fir.box<!fir.array<?xf64>>
+ func.call @use(%emb2) : (!fir.box<!fir.array<?xf64>>) -> ()
+ cf.br ^exit
+^exit:
+ fir.freemem %heap : !fir.heap<!fir.array<?xf64>>
+ return
+}
More information about the flang-commits
mailing list