[flang-commits] [flang] [flang] Fold reads of fir.create_box components (PR #210388)
Slava Zakharin via flang-commits
flang-commits at lists.llvm.org
Fri Jul 17 10:55:08 PDT 2026
https://github.com/vzakhari created https://github.com/llvm/llvm-project/pull/210388
Fold fir.box_addr and constant-dimension fir.box_dims of a descriptor
built by fir.create_box back to the operands that built it, so the
create_box can die when nothing consumes it as a real descriptor.
fir.array_coor is intentionally not folded through create_box, since
that would drop its explicit byte strides.
Assisted-by: Cursor
>From 08732fe2b02a0a3597f406547aad9dfc77d44086 Mon Sep 17 00:00:00 2001
From: Slava Zakharin <szakharin at nvidia.com>
Date: Fri, 17 Jul 2026 10:05:20 -0700
Subject: [PATCH] [flang] Fold reads of fir.create_box components
Fold fir.box_addr and constant-dimension fir.box_dims of a descriptor
built by fir.create_box back to the operands that built it, so the
create_box can die when nothing consumes it as a real descriptor.
fir.array_coor is intentionally not folded through create_box, since
that would drop its explicit byte strides.
Assisted-by: Cursor
---
.../include/flang/Optimizer/Dialect/FIROps.td | 2 +
flang/lib/Optimizer/Dialect/FIROps.cpp | 50 ++++++++++++++++
.../test/Fir/array-coor-canonicalization.fir | 15 +++++
flang/test/Fir/box-dims-folding.fir | 60 +++++++++++++++++++
flang/test/Fir/boxaddr-folding.fir | 37 ++++++++++++
5 files changed, 164 insertions(+)
create mode 100644 flang/test/Fir/box-dims-folding.fir
diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td
index df7fdb27586d9..23aee26e0914a 100644
--- a/flang/include/flang/Optimizer/Dialect/FIROps.td
+++ b/flang/include/flang/Optimizer/Dialect/FIROps.td
@@ -1314,6 +1314,8 @@ def fir_BoxDimsOp
$val `,` $dim attr-dict `:` functional-type(operands, results)
}];
+ let hasCanonicalizer = 1;
+
let extraClassDeclaration = [{
mlir::Type getTupleType();
mlir::Value getLowerBound() {return getResult(0);};
diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index 38c662fb23a99..d0e280dbf9712 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -1505,6 +1505,15 @@ mlir::OpFoldResult fir::BoxAddrOp::fold(FoldAdaptor adaptor) {
return box.getMemref();
}
}
+ // fir.create_box always describes the whole array (no slice), so the
+ // base address is exactly its memref operand. The verifier guarantees the
+ // result box preserves the memref storage kind, so the types match.
+ if (auto box = mlir::dyn_cast<fir::CreateBoxOp>(v)) {
+ if (box.getMemref().getType() == getType()) {
+ propagateAttributes(getOperation(), box.getMemref().getDefiningOp());
+ return box.getMemref();
+ }
+ }
if (auto box = mlir::dyn_cast<fir::EmboxCharOp>(v))
if (box.getMemref().getType() == getType())
return box.getMemref();
@@ -1604,6 +1613,47 @@ void fir::BoxEleSizeOp::getCanonicalizationPatterns(
// BoxDimsOp
//===----------------------------------------------------------------------===//
+namespace {
+/// fir.box_dims of a descriptor built by fir.create_box with a statically known
+/// dimension forwards the lower bound, extent, and byte stride operands that
+/// built that dimension. These are plain SSA operands of the create_box, so
+/// replacing the box_dims results with them is a value forwarding; the
+/// create_box dominates the box_dims, so the operands are already available.
+struct SimplifyBoxDimsFromCreateBox
+ : public mlir::OpRewritePattern<fir::BoxDimsOp> {
+ using mlir::OpRewritePattern<fir::BoxDimsOp>::OpRewritePattern;
+
+ mlir::LogicalResult
+ matchAndRewrite(fir::BoxDimsOp op,
+ mlir::PatternRewriter &rewriter) const override {
+ auto createBox = op.getVal().getDefiningOp<fir::CreateBoxOp>();
+ if (!createBox)
+ return mlir::failure();
+
+ std::optional<std::int64_t> dim = fir::getIntIfConstant(op.getDim());
+ if (!dim || *dim < 0 || *dim >= createBox.getRank())
+ return mlir::failure();
+
+ // The create_box operands are AnyIntegerType and may differ from the
+ // index-typed box_dims results, so convert them.
+ mlir::Location loc = op.getLoc();
+ rewriter.replaceOp(op,
+ {rewriter.createOrFold<fir::ConvertOp>(
+ loc, op.getType(0), createBox.getLbounds()[*dim]),
+ rewriter.createOrFold<fir::ConvertOp>(
+ loc, op.getType(1), createBox.getExtents()[*dim]),
+ rewriter.createOrFold<fir::ConvertOp>(
+ loc, op.getType(2), createBox.getStrides()[*dim])});
+ return mlir::success();
+ }
+};
+} // namespace
+
+void fir::BoxDimsOp::getCanonicalizationPatterns(
+ mlir::RewritePatternSet &results, mlir::MLIRContext *context) {
+ results.insert<SimplifyBoxDimsFromCreateBox>(context);
+}
+
/// Get the result types packed in a tuple tuple
mlir::Type fir::BoxDimsOp::getTupleType() {
// note: triple, but 4 is nearest power of 2
diff --git a/flang/test/Fir/array-coor-canonicalization.fir b/flang/test/Fir/array-coor-canonicalization.fir
index cb2efa007294f..d2b8bc06a0c76 100644
--- a/flang/test/Fir/array-coor-canonicalization.fir
+++ b/flang/test/Fir/array-coor-canonicalization.fir
@@ -1059,3 +1059,18 @@ func.func @test37_merge_no_fold_inner_rank_too_large(%base: !fir.ref<!fir.array<
%inner = fir.array_coor %view(%inner_shape) %i, %j : (!fir.ref<!fir.array<3x4xf64>>, !fir.shape<2>, index, index) -> !fir.ref<f64>
return %inner : !fir.ref<f64>
}
+
+// No fold: a fir.array_coor on a descriptor built by fir.create_box must not be
+// rewritten into an array_coor on the base pointer. create_box carries
+// arbitrary byte strides that an array_coor on a ref+shape cannot express, so
+// keeping the array_coor on the box is what preserves those strides.
+// CHECK-LABEL: func.func @test38_no_fold_array_coor_create_box(
+// CHECK: %[[BOX:.*]] = fir.create_box
+// CHECK: fir.array_coor %[[BOX]]
+func.func @test38_no_fold_array_coor_create_box(%arg0: !fir.ref<!fir.array<?xf32>>,
+ %lb: index, %ext: index,
+ %sm: index, %idx: index) -> !fir.ref<f32> {
+ %0 = fir.create_box %arg0 lbs(%lb) extents(%ext) strides(%sm) : (!fir.ref<!fir.array<?xf32>>, index, index, index) -> !fir.box<!fir.array<?xf32>>
+ %1 = fir.array_coor %0 %idx : (!fir.box<!fir.array<?xf32>>, index) -> !fir.ref<f32>
+ return %1 : !fir.ref<f32>
+}
diff --git a/flang/test/Fir/box-dims-folding.fir b/flang/test/Fir/box-dims-folding.fir
new file mode 100644
index 0000000000000..e566df81c3f9a
--- /dev/null
+++ b/flang/test/Fir/box-dims-folding.fir
@@ -0,0 +1,60 @@
+// RUN: fir-opt --canonicalize %s -split-input-file | FileCheck %s
+
+// fir.box_dims of a fir.create_box with a constant dimension forwards the lower
+// bound, extent, and byte stride operands that built that dimension.
+// CHECK-LABEL: func.func @fold_box_dims_create_box(
+// CHECK-SAME: %[[BASE:[^:]+]]: !fir.ref<!fir.array<?xf32>>, %[[LB:[^:]+]]: index, %[[EXT:[^:]+]]: index, %[[SM:[^:]+]]: index)
+// CHECK-NOT: fir.box_dims
+// CHECK: return %[[LB]], %[[EXT]], %[[SM]]
+func.func @fold_box_dims_create_box(%base: !fir.ref<!fir.array<?xf32>>, %lb: index, %ext: index, %sm: index) -> (index, index, index) {
+ %c0 = arith.constant 0 : index
+ %0 = fir.create_box %base lbs(%lb) extents(%ext) strides(%sm) : (!fir.ref<!fir.array<?xf32>>, index, index, index) -> !fir.box<!fir.array<?xf32>>
+ %1:3 = fir.box_dims %0, %c0 : (!fir.box<!fir.array<?xf32>>, index) -> (index, index, index)
+ return %1#0, %1#1, %1#2 : index, index, index
+}
+
+// -----
+
+// The forwarded operands are selected per dimension: dim 1 forwards the second
+// lower bound / extent / stride triple.
+// CHECK-LABEL: func.func @fold_box_dims_create_box_2d(
+// CHECK-SAME: %[[BASE:[^:]+]]: !fir.ref<!fir.array<?x?xf32>>, %[[LB0:[^:]+]]: index, %[[EXT0:[^:]+]]: index, %[[SM0:[^:]+]]: index, %[[LB1:[^:]+]]: index, %[[EXT1:[^:]+]]: index, %[[SM1:[^:]+]]: index)
+// CHECK-NOT: fir.box_dims
+// CHECK: return %[[LB1]], %[[EXT1]], %[[SM1]]
+func.func @fold_box_dims_create_box_2d(%base: !fir.ref<!fir.array<?x?xf32>>, %lb0: index, %ext0: index, %sm0: index, %lb1: index, %ext1: index, %sm1: index) -> (index, index, index) {
+ %c1 = arith.constant 1 : index
+ %0 = fir.create_box %base lbs(%lb0, %lb1) extents(%ext0, %ext1) strides(%sm0, %sm1) : (!fir.ref<!fir.array<?x?xf32>>, index, index, index, index, index, index) -> !fir.box<!fir.array<?x?xf32>>
+ %1:3 = fir.box_dims %0, %c1 : (!fir.box<!fir.array<?x?xf32>>, index) -> (index, index, index)
+ return %1#0, %1#1, %1#2 : index, index, index
+}
+
+// -----
+
+// The create_box operands are AnyIntegerType; when they are not index-typed the
+// fold inserts fir.convert so the forwarded values match the box_dims results.
+// CHECK-LABEL: func.func @fold_box_dims_create_box_i64_operands(
+// CHECK-SAME: %[[BASE:[^:]+]]: !fir.ref<!fir.array<?xf32>>, %[[LB:[^:]+]]: i64, %[[EXT:[^:]+]]: i64, %[[SM:[^:]+]]: i64)
+// CHECK-NOT: fir.box_dims
+// CHECK-DAG: %[[LBC:.*]] = fir.convert %[[LB]] : (i64) -> index
+// CHECK-DAG: %[[EXTC:.*]] = fir.convert %[[EXT]] : (i64) -> index
+// CHECK-DAG: %[[SMC:.*]] = fir.convert %[[SM]] : (i64) -> index
+// CHECK: return %[[LBC]], %[[EXTC]], %[[SMC]]
+func.func @fold_box_dims_create_box_i64_operands(%base: !fir.ref<!fir.array<?xf32>>, %lb: i64, %ext: i64, %sm: i64) -> (index, index, index) {
+ %c0 = arith.constant 0 : index
+ %0 = fir.create_box %base lbs(%lb) extents(%ext) strides(%sm) : (!fir.ref<!fir.array<?xf32>>, i64, i64, i64) -> !fir.box<!fir.array<?xf32>>
+ %1:3 = fir.box_dims %0, %c0 : (!fir.box<!fir.array<?xf32>>, index) -> (index, index, index)
+ return %1#0, %1#1, %1#2 : index, index, index
+}
+
+// -----
+
+// No fold: the queried dimension is not a compile-time constant, so the right
+// operand triple cannot be selected.
+// CHECK-LABEL: func.func @no_fold_box_dims_create_box_dynamic_dim(
+// CHECK: %[[BOX:.*]] = fir.create_box
+// CHECK: fir.box_dims %[[BOX]]
+func.func @no_fold_box_dims_create_box_dynamic_dim(%base: !fir.ref<!fir.array<?x?xf32>>, %lb0: index, %ext0: index, %sm0: index, %lb1: index, %ext1: index, %sm1: index, %d: index) -> (index, index, index) {
+ %0 = fir.create_box %base lbs(%lb0, %lb1) extents(%ext0, %ext1) strides(%sm0, %sm1) : (!fir.ref<!fir.array<?x?xf32>>, index, index, index, index, index, index) -> !fir.box<!fir.array<?x?xf32>>
+ %1:3 = fir.box_dims %0, %d : (!fir.box<!fir.array<?x?xf32>>, index) -> (index, index, index)
+ return %1#0, %1#1, %1#2 : index, index, index
+}
diff --git a/flang/test/Fir/boxaddr-folding.fir b/flang/test/Fir/boxaddr-folding.fir
index 4fd23864ae49c..f7bbbdb71556a 100644
--- a/flang/test/Fir/boxaddr-folding.fir
+++ b/flang/test/Fir/boxaddr-folding.fir
@@ -44,3 +44,40 @@ func.func @check_folding(%arg0 : !fir.ref<!fir.array<?xi32>>) {
func.func @check(%arg0: !fir.ref<!fir.array<?xi32>>) {
return
}
+
+// -----
+
+// The base address of a descriptor built by fir.create_box is exactly its
+// memref operand (create_box never slices), so box_addr folds to that operand
+// when the result type preserves the memref storage kind.
+// CHECK-LABEL: func @check_create_box_folding
+func.func @check_create_box_folding(%arg0: !fir.ref<!fir.array<?xf32>>, %lb: index, %ext: index, %sm: index) {
+ %0 = fir.create_box %arg0 lbs(%lb) extents(%ext) strides(%sm) : (!fir.ref<!fir.array<?xf32>>, index, index, index) -> !fir.box<!fir.array<?xf32>>
+ %1 = fir.box_addr %0 : (!fir.box<!fir.array<?xf32>>) -> !fir.ref<!fir.array<?xf32>>
+ // CHECK-NOT: fir.box_addr
+ // CHECK: fir.call @check_ref(%arg0)
+ fir.call @check_ref(%1) : (!fir.ref<!fir.array<?xf32>>) -> ()
+ return
+}
+
+func.func @check_ref(%arg0: !fir.ref<!fir.array<?xf32>>) {
+ return
+}
+
+// -----
+
+// When the box_addr result type does not match the create_box memref type the
+// storage kind would change, so the fold must not fire.
+// CHECK-LABEL: func @check_create_box_no_folding
+func.func @check_create_box_no_folding(%arg0: !fir.heap<!fir.array<?xf32>>, %lb: index, %ext: index, %sm: index) {
+ %0 = fir.create_box %arg0 lbs(%lb) extents(%ext) strides(%sm) : (!fir.heap<!fir.array<?xf32>>, index, index, index) -> !fir.box<!fir.heap<!fir.array<?xf32>>>
+ %1 = fir.box_addr %0 : (!fir.box<!fir.heap<!fir.array<?xf32>>>) -> !fir.ref<!fir.array<?xf32>>
+ // CHECK: %[[ADDR:.*]] = fir.box_addr
+ // CHECK: fir.call @check_ref(%[[ADDR]])
+ fir.call @check_ref(%1) : (!fir.ref<!fir.array<?xf32>>) -> ()
+ return
+}
+
+func.func @check_ref(%arg0: !fir.ref<!fir.array<?xf32>>) {
+ return
+}
More information about the flang-commits
mailing list