[flang-commits] [flang] [flang] Canonicalize fir.array_coor by pulling in embox/rebox. (PR #92858)
Slava Zakharin via flang-commits
flang-commits at lists.llvm.org
Wed May 29 22:14:10 PDT 2024
https://github.com/vzakhari updated https://github.com/llvm/llvm-project/pull/92858
>From 98bfd5237a4a41673f778eba8f6fa52c6901768e Mon Sep 17 00:00:00 2001
From: Slava Zakharin <szakharin at nvidia.com>
Date: Mon, 20 May 2024 21:11:56 -0700
Subject: [PATCH 1/5] [flang] Canonicalize fir.array_coor by pulling in
embox/rebox.
In a simple case like this:
```
program test
integer :: u(120, 2)
u(1:120,1:2) = u(1:120,1:2) + 2
end program
```
Flang is creating a copy loop with fir.array_coor using
a result of fir.embox inserted before the loop. This results in split
address computations before and inside the loop, which can be seen
as many more arithmetic operations than required after converting
FIR to LLVM dialect. Even though LLVM SROA/mem2reg are able
to optimize the temporary descriptor, and then LICM is able to hoist
the invariant computations, we seem to get better mix of LLVM dialect
operations after FIR-to-LLVM codegen. This may also slightly reduce
the compilation time taken by LLVM to optimize the generate LLVM IR.
This may also slightly reduce the time spent by FIR AliasAnalysis
to reach the memory reference source.
This patch also includes one change in the CodeGen to fix what I believe
is a bug: the indices used in fir.array_coor with a slice are
always ranging from 1, while the current code assumed they were
in the range of the slice bounds.
---
.../include/flang/Optimizer/Dialect/FIROps.td | 1 +
flang/lib/Optimizer/CodeGen/CodeGen.cpp | 6 +-
flang/lib/Optimizer/Dialect/FIROps.cpp | 93 ++++++++++++++
.../test/Fir/array-coor-canonicalization.fir | 121 ++++++++++++++++++
flang/test/Fir/convert-to-llvm.fir | 4 +-
5 files changed, 222 insertions(+), 3 deletions(-)
create mode 100644 flang/test/Fir/array-coor-canonicalization.fir
diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td
index d9c1149040066..7ffa0145072d6 100644
--- a/flang/include/flang/Optimizer/Dialect/FIROps.td
+++ b/flang/include/flang/Optimizer/Dialect/FIROps.td
@@ -1650,6 +1650,7 @@ def fir_ArrayCoorOp : fir_Op<"array_coor",
}];
let hasVerifier = 1;
+ let hasCanonicalizer = 1;
}
def fir_CoordinateOp : fir_Op<"coordinate_of", [NoMemoryEffect]> {
diff --git a/flang/lib/Optimizer/CodeGen/CodeGen.cpp b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
index 72172f63888e1..8d70ceb31bce1 100644
--- a/flang/lib/Optimizer/CodeGen/CodeGen.cpp
+++ b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
@@ -2162,7 +2162,11 @@ struct XArrayCoorOpConversion
if (normalSlice)
step = integerCast(loc, rewriter, idxTy, operands[sliceOffset + 2]);
}
- auto idx = rewriter.create<mlir::LLVM::SubOp>(loc, idxTy, index, lb, nsw);
+ // The array_coor indices are always 1-based if slicing
+ // is in effect. The non-default lower bounds only
+ // apply to the indices of the slice itself.
+ auto idx = rewriter.create<mlir::LLVM::SubOp>(loc, idxTy, index,
+ isSliced ? one : lb, nsw);
mlir::Value diff =
rewriter.create<mlir::LLVM::MulOp>(loc, idxTy, idx, step, nsw);
if (normalSlice) {
diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index 94113da9a46cf..67320d4e6e7eb 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -394,11 +394,17 @@ mlir::LogicalResult fir::ArrayCoorOp::verify() {
} else {
auto s = mlir::cast<fir::ShiftType>(shapeTy);
shapeTyRank = s.getRank();
+ // TODO: it looks like PreCGRewrite and CodeGen can support
+ // fir.shift with plain array reference, so we may consider
+ // removing this check.
if (!mlir::isa<fir::BaseBoxType>(getMemref().getType()))
return emitOpError("shift can only be provided with fir.box memref");
}
if (arrDim && arrDim != shapeTyRank)
return emitOpError("rank of dimension mismatched");
+ // TODO: support slicing with changing the numbder of dimensions,
+ // e.g. when array_coord represents an element access to array(:,1,:)
+ // slice: the shape is 3D and the number of indices is 2 in this case.
if (shapeTyRank != getIndices().size())
return emitOpError("number of indices do not match dim rank");
}
@@ -417,6 +423,93 @@ mlir::LogicalResult fir::ArrayCoorOp::verify() {
return mlir::success();
}
+// Pull in fir.embox and fir.rebox into fir.array_coor when possible.
+struct SimplifyArrayCoorOp : public mlir::OpRewritePattern<fir::ArrayCoorOp> {
+ using mlir::OpRewritePattern<fir::ArrayCoorOp>::OpRewritePattern;
+ mlir::LogicalResult
+ matchAndRewrite(fir::ArrayCoorOp op,
+ mlir::PatternRewriter &rewriter) const override {
+ mlir::Value memref = op.getMemref();
+ if (!mlir::isa<fir::BaseBoxType>(memref.getType()))
+ return mlir::failure();
+
+ mlir::Value boxedMemref, boxedShape, boxedSlice;
+ if (auto emboxOp =
+ mlir::dyn_cast_or_null<fir::EmboxOp>(memref.getDefiningOp())) {
+ boxedMemref = emboxOp.getMemref();
+ boxedShape = emboxOp.getShape();
+ boxedSlice = emboxOp.getSlice();
+ // If any of operands, that are not currently supported for migration
+ // to ArrayCoorOp, is present, don't rewrite.
+ if (!emboxOp.getTypeparams().empty() || emboxOp.getSourceBox() ||
+ emboxOp.getAccessMap())
+ return mlir::failure();
+ } else if (auto reboxOp = mlir::dyn_cast_or_null<fir::ReboxOp>(
+ memref.getDefiningOp())) {
+ boxedMemref = reboxOp.getBox();
+ boxedShape = reboxOp.getShape();
+ boxedSlice = reboxOp.getSlice();
+ } else {
+ return mlir::failure();
+ }
+
+ // Slices changing the number of dimensions are not supported
+ // for array_coor yet.
+ unsigned origBoxRank;
+ if (mlir::isa<fir::BaseBoxType>(boxedMemref.getType()))
+ origBoxRank = fir::getBoxRank(boxedMemref.getType());
+ else if (auto arrTy = mlir::dyn_cast<fir::SequenceType>(
+ fir::unwrapRefType(boxedMemref.getType())))
+ origBoxRank = arrTy.getDimension();
+ else
+ return mlir::failure();
+
+ if (fir::getBoxRank(memref.getType()) != origBoxRank)
+ return mlir::failure();
+
+ // Slices with substring are not supported by array_coor.
+ if (boxedSlice)
+ if (auto sliceOp =
+ mlir::dyn_cast_or_null<fir::SliceOp>(boxedSlice.getDefiningOp()))
+ if (!sliceOp.getSubstr().empty())
+ return mlir::failure();
+
+ // If embox/rebox and array_coor have conflicting shapes or slices,
+ // do nothing.
+ if (op.getShape() && boxedShape && boxedShape != op.getShape())
+ return mlir::failure();
+ if (op.getSlice() && boxedSlice && boxedSlice != op.getSlice())
+ return mlir::failure();
+
+ // TODO: temporarily avoid producing array_coor with the shape shift
+ // and plain array reference (it seems to be a limitation of
+ // ArrayCoorOp verifier).
+ if (!mlir::isa<fir::BaseBoxType>(boxedMemref.getType())) {
+ if (boxedShape) {
+ if (mlir::isa<fir::ShiftType>(boxedShape.getType()))
+ return mlir::failure();
+ } else if (op.getShape() &&
+ mlir::isa<fir::ShiftType>(op.getShape().getType())) {
+ return mlir::failure();
+ }
+ }
+
+ rewriter.modifyOpInPlace(op, [&]() {
+ op.getMemrefMutable().assign(boxedMemref);
+ if (boxedShape)
+ op.getShapeMutable().assign(boxedShape);
+ if (boxedSlice)
+ op.getSliceMutable().assign(boxedSlice);
+ });
+ return mlir::success();
+ }
+};
+
+void fir::ArrayCoorOp::getCanonicalizationPatterns(
+ mlir::RewritePatternSet &patterns, mlir::MLIRContext *context) {
+ patterns.add<SimplifyArrayCoorOp>(context);
+}
+
//===----------------------------------------------------------------------===//
// ArrayLoadOp
//===----------------------------------------------------------------------===//
diff --git a/flang/test/Fir/array-coor-canonicalization.fir b/flang/test/Fir/array-coor-canonicalization.fir
new file mode 100644
index 0000000000000..ec36b2e57b9f9
--- /dev/null
+++ b/flang/test/Fir/array-coor-canonicalization.fir
@@ -0,0 +1,121 @@
+// RUN: fir-opt --canonicalize %s | FileCheck %s
+
+// CHECK-LABEL: func.func @_QPtest1(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<120x2xi32>> {fir.bindc_name = "u"}) {
+// CHECK: %[[VAL_6:.*]] = fir.shape
+// CHECK: %[[VAL_7:.*]] = fir.declare %[[VAL_0]](%[[VAL_6]])
+// CHECK: %[[VAL_8:.*]] = fir.slice
+// CHECK: fir.do_loop
+// CHECK: fir.do_loop
+// CHECK: %[[VAL_11:.*]] = fir.array_coor %[[VAL_7]](%[[VAL_6]]) {{\[}}%[[VAL_8]]]
+func.func @_QPtest1(%arg0: !fir.ref<!fir.array<120x2xi32>> {fir.bindc_name = "u"}) {
+ %c1 = arith.constant 1 : index
+ %c2_i32 = arith.constant 2 : i32
+ %c2 = arith.constant 2 : index
+ %c120 = arith.constant 120 : index
+ %0 = fir.dummy_scope : !fir.dscope
+ %1 = fir.shape %c120, %c2 : (index, index) -> !fir.shape<2>
+ %2 = fir.declare %arg0(%1) dummy_scope %0 {uniq_name = "_QFtest1Eu"} : (!fir.ref<!fir.array<120x2xi32>>, !fir.shape<2>, !fir.dscope) -> !fir.ref<!fir.array<120x2xi32>>
+ %3 = fir.slice %c1, %c120, %c1, %c1, %c2, %c1 : (index, index, index, index, index, index) -> !fir.slice<2>
+ %4 = fir.embox %2(%1) [%3] : (!fir.ref<!fir.array<120x2xi32>>, !fir.shape<2>, !fir.slice<2>) -> !fir.box<!fir.array<120x2xi32>>
+ fir.do_loop %arg1 = %c1 to %c2 step %c1 unordered {
+ fir.do_loop %arg2 = %c1 to %c120 step %c1 unordered {
+ %5 = fir.array_coor %4 %arg2, %arg1 : (!fir.box<!fir.array<120x2xi32>>, index, index) -> !fir.ref<i32>
+ fir.store %c2_i32 to %5 : !fir.ref<i32>
+ }
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @_QPtest2(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.box<!fir.array<?x?xi32>> {fir.bindc_name = "u"}) {
+// CHECK: %[[VAL_8:.*]] = fir.shift
+// CHECK: %[[VAL_9:.*]] = fir.declare %[[VAL_0]](%[[VAL_8]])
+// CHECK: fir.do_loop
+// CHECK: fir.do_loop
+// CHECK: %[[VAL_17:.*]] = fir.array_coor %[[VAL_9]](%[[VAL_8]])
+func.func @_QPtest2(%arg0: !fir.box<!fir.array<?x?xi32>> {fir.bindc_name = "u"}) {
+ %c9 = arith.constant 9 : index
+ %c1 = arith.constant 1 : index
+ %c0 = arith.constant 0 : index
+ %c11 = arith.constant 11 : index
+ %c10 = arith.constant 10 : index
+ %c2_i32 = arith.constant 2 : i32
+ %0 = fir.dummy_scope : !fir.dscope
+ %1 = fir.shift %c10, %c11 : (index, index) -> !fir.shift<2>
+ %2 = fir.declare %arg0(%1) dummy_scope %0 {uniq_name = "_QFtest2Eu"} : (!fir.box<!fir.array<?x?xi32>>, !fir.shift<2>, !fir.dscope) -> !fir.box<!fir.array<?x?xi32>>
+ %3 = fir.rebox %2(%1) : (!fir.box<!fir.array<?x?xi32>>, !fir.shift<2>) -> !fir.box<!fir.array<?x?xi32>>
+ %4:3 = fir.box_dims %3, %c0 : (!fir.box<!fir.array<?x?xi32>>, index) -> (index, index, index)
+ %5:3 = fir.box_dims %3, %c1 : (!fir.box<!fir.array<?x?xi32>>, index) -> (index, index, index)
+ fir.do_loop %arg1 = %c1 to %5#1 step %c1 unordered {
+ fir.do_loop %arg2 = %c1 to %4#1 step %c1 unordered {
+ %6 = arith.addi %arg2, %c9 : index
+ %7 = arith.addi %arg1, %c10 : index
+ %8 = fir.array_coor %3(%1) %6, %7 : (!fir.box<!fir.array<?x?xi32>>, !fir.shift<2>, index, index) -> !fir.ref<i32>
+ fir.store %c2_i32 to %8 : !fir.ref<i32>
+ }
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @_QPtest3(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.box<!fir.array<?x?xi32>> {fir.bindc_name = "u"}) {
+// CHECK: %[[VAL_10:.*]] = fir.shift
+// CHECK: %[[VAL_11:.*]] = fir.declare %[[VAL_0]](%[[VAL_10]])
+// CHECK: %[[VAL_12:.*]] = fir.slice
+// CHECK: fir.do_loop
+// CHECK: fir.do_loop
+// CHECK: %[[VAL_15:.*]] = fir.array_coor %[[VAL_11]](%[[VAL_10]]) {{\[}}%[[VAL_12]]]
+func.func @_QPtest3(%arg0: !fir.box<!fir.array<?x?xi32>> {fir.bindc_name = "u"}) {
+ %c2 = arith.constant 2 : index
+ %c12 = arith.constant 12 : index
+ %c11 = arith.constant 11 : index
+ %c111 = arith.constant 111 : index
+ %c1 = arith.constant 1 : index
+ %c120 = arith.constant 120 : index
+ %c10 = arith.constant 10 : index
+ %c2_i32 = arith.constant 2 : i32
+ %0 = fir.dummy_scope : !fir.dscope
+ %1 = fir.shift %c10, %c11 : (index, index) -> !fir.shift<2>
+ %2 = fir.declare %arg0(%1) dummy_scope %0 {uniq_name = "_QFtest3Eu"} : (!fir.box<!fir.array<?x?xi32>>, !fir.shift<2>, !fir.dscope) -> !fir.box<!fir.array<?x?xi32>>
+ %3 = fir.rebox %2(%1) : (!fir.box<!fir.array<?x?xi32>>, !fir.shift<2>) -> !fir.box<!fir.array<?x?xi32>>
+ %4 = fir.slice %c10, %c120, %c1, %c11, %c12, %c1 : (index, index, index, index, index, index) -> !fir.slice<2>
+ %5 = fir.rebox %3(%1) [%4] : (!fir.box<!fir.array<?x?xi32>>, !fir.shift<2>, !fir.slice<2>) -> !fir.box<!fir.array<111x2xi32>>
+ fir.do_loop %arg1 = %c1 to %c2 step %c1 unordered {
+ fir.do_loop %arg2 = %c1 to %c111 step %c1 unordered {
+ %6 = fir.array_coor %5 %arg2, %arg1 : (!fir.box<!fir.array<111x2xi32>>, index, index) -> !fir.ref<i32>
+ fir.store %c2_i32 to %6 : !fir.ref<i32>
+ }
+ }
+ return
+}
+
+// TODO: fir.array_coor with slices changing the number of dimensions
+// is not supported yet.
+// CHECK-LABEL: func.func @_QPtest4() {
+// CHECK: %[[VAL_3:.*]] = fir.alloca !fir.array<100x100x100xi32> {bindc_name = "u", uniq_name = "_QFtest4Eu"}
+// CHECK: %[[VAL_4:.*]] = fir.shape
+// CHECK: %[[VAL_5:.*]] = fir.declare %[[VAL_3]](%[[VAL_4]]) {uniq_name = "_QFtest4Eu"} : (!fir.ref<!fir.array<100x100x100xi32>>, !fir.shape<3>) -> !fir.ref<!fir.array<100x100x100xi32>>
+// CHECK: %[[VAL_7:.*]] = fir.slice
+// CHECK: %[[VAL_8:.*]] = fir.embox %[[VAL_5]](%[[VAL_4]]) {{\[}}%[[VAL_7]]] : (!fir.ref<!fir.array<100x100x100xi32>>, !fir.shape<3>, !fir.slice<3>) -> !fir.box<!fir.array<100x100xi32>>
+// CHECK: fir.do_loop
+// CHECK: fir.do_loop
+// CHECK: %[[VAL_11:.*]] = fir.array_coor %[[VAL_8]]
+func.func @_QPtest4() {
+ %c1 = arith.constant 1 : index
+ %c2_i32 = arith.constant 2 : i32
+ %c100 = arith.constant 100 : index
+ %0 = fir.alloca !fir.array<100x100x100xi32> {bindc_name = "u", uniq_name = "_QFtest4Eu"}
+ %1 = fir.shape %c100, %c100, %c100 : (index, index, index) -> !fir.shape<3>
+ %2 = fir.declare %0(%1) {uniq_name = "_QFtest4Eu"} : (!fir.ref<!fir.array<100x100x100xi32>>, !fir.shape<3>) -> !fir.ref<!fir.array<100x100x100xi32>>
+ %3 = fir.undefined index
+ %4 = fir.slice %c1, %c100, %c1, %c1, %3, %3, %c1, %c100, %c1 : (index, index, index, index, index, index, index, index, index) -> !fir.slice<3>
+ %5 = fir.embox %2(%1) [%4] : (!fir.ref<!fir.array<100x100x100xi32>>, !fir.shape<3>, !fir.slice<3>) -> !fir.box<!fir.array<100x100xi32>>
+ fir.do_loop %arg0 = %c1 to %c100 step %c1 unordered {
+ fir.do_loop %arg1 = %c1 to %c100 step %c1 unordered {
+ %6 = fir.array_coor %5 %arg1, %arg0 : (!fir.box<!fir.array<100x100xi32>>, index, index) -> !fir.ref<i32>
+ fir.store %c2_i32 to %6 : !fir.ref<i32>
+ }
+ }
+ return
+}
diff --git a/flang/test/Fir/convert-to-llvm.fir b/flang/test/Fir/convert-to-llvm.fir
index 21323a5e657c9..80aebd5f8500d 100644
--- a/flang/test/Fir/convert-to-llvm.fir
+++ b/flang/test/Fir/convert-to-llvm.fir
@@ -2084,7 +2084,7 @@ func.func @ext_array_coor1(%arg0: !fir.ref<!fir.array<?xi32>>) {
// CHECK: %[[C0:.*]] = llvm.mlir.constant(0 : i64) : i64
// CHECK: %[[C1:.*]] = llvm.mlir.constant(1 : i64) : i64
// CHECK: %[[C0_1:.*]] = llvm.mlir.constant(0 : i64) : i64
-// CHECK: %[[IDX:.*]] = llvm.sub %[[C0]], %[[C0]] overflow<nsw> : i64
+// CHECK: %[[IDX:.*]] = llvm.sub %[[C0]], %[[C1]] overflow<nsw> : i64
// CHECK: %[[DIFF0:.*]] = llvm.mul %[[IDX]], %[[C0]] overflow<nsw> : i64
// CHECK: %[[ADJ:.*]] = llvm.sub %[[C0]], %[[C0]] overflow<nsw> : i64
// CHECK: %[[DIFF1:.*]] = llvm.add %[[DIFF0]], %[[ADJ]] overflow<nsw> : i64
@@ -2153,7 +2153,7 @@ func.func @ext_array_coor4(%arg0: !fir.ref<!fir.array<100xi32>>) {
// CHECK: %[[C1:.*]] = llvm.mlir.constant(1 : i64) : i64
// CHECK: %[[C1_1:.*]] = llvm.mlir.constant(1 : i64) : i64
// CHECK: %[[C0_1:.*]] = llvm.mlir.constant(0 : i64) : i64
-// CHECK: %[[IDX:.*]] = llvm.sub %[[C1]], %[[C0]] overflow<nsw> : i64
+// CHECK: %[[IDX:.*]] = llvm.sub %[[C1]], %[[C1_1]] overflow<nsw> : i64
// CHECK: %[[DIFF0:.*]] = llvm.mul %[[IDX]], %[[C1]] overflow<nsw> : i64
// CHECK: %[[ADJ:.*]] = llvm.sub %[[C10]], %[[C0]] overflow<nsw> : i64
// CHECK: %[[DIFF1:.*]] = llvm.add %[[DIFF0]], %[[ADJ]] overflow<nsw> : i64
>From 28e6c49bd6508942cdb82b2edb15d97d772f25a3 Mon Sep 17 00:00:00 2001
From: Slava Zakharin <szakharin at nvidia.com>
Date: Tue, 21 May 2024 17:54:31 -0700
Subject: [PATCH 2/5] Addressed review comments.
---
flang/lib/Optimizer/Dialect/FIROps.cpp | 123 +++++++
.../test/Fir/array-coor-canonicalization.fir | 305 ++++++++++++++++++
2 files changed, 428 insertions(+)
diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index 67320d4e6e7eb..5c5123273da8a 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -481,6 +481,129 @@ struct SimplifyArrayCoorOp : public mlir::OpRewritePattern<fir::ArrayCoorOp> {
if (op.getSlice() && boxedSlice && boxedSlice != op.getSlice())
return mlir::failure();
+ // The embox/rebox and array_coor either have compatible
+ // shape/slice at this point or shape/slice is null
+ // in one of them but not in the other.
+ // The compatibility means they are equal or both null.
+ if (!op.getShape()) {
+ if (boxedShape) {
+ if (op.getSlice()) {
+ if (!boxedSlice) {
+ // %0 = fir.rebox %arg(%shape)
+ // %1 = fir.array_coor %0 [%slice] %idx
+ // Both the slice indices and %idx are 1-based, so the rebox
+ // may be pulled in as:
+ // %1 = fir.array_coor %arg [%slice] %idx
+ boxedShape = nullptr;
+ } else {
+ // %0 = fir.rebox %arg(%shape) [%slice]
+ // %1 = fir.array_coor %0 [%slice] %idx
+ // I believe this FIR may only be valid if the shape specifies
+ // that all lower bounds are 1s and the slice's start indices
+ // are all 1s.
+ // We could pull in the rebox as:
+ // %1 = fir.array_coor %arg [%slice] %idx
+ // Do not do anything for the time being.
+ return mlir::failure();
+ }
+ } else { // !op.getSlice()
+ if (!boxedSlice) {
+ // %0 = fir.rebox %arg(%shape)
+ // %1 = fir.array_coor %0 %idx
+ // Pull in as:
+ // %1 = fir.array_coor %arg %idx
+ boxedShape = nullptr;
+ } else {
+ // %0 = fir.rebox %arg(%shape) [%slice]
+ // %1 = fir.array_coor %0 %idx
+ // Pull in as:
+ // %1 = fir.array_coor %arg(%shape) %[slice] %idx
+ }
+ }
+ } else { // !boxedShape
+ if (op.getSlice()) {
+ if (!boxedSlice) {
+ // %0 = fir.rebox %arg
+ // %1 = fir.array_coor %0 [%slice] %idx
+ // Pull in as:
+ // %1 = fir.array_coor %arg [%slice] %idx
+ } else {
+ // %0 = fir.rebox %arg [%slice]
+ // %1 = fir.array_coor %0 [%slice] %idx
+ // Pull in as:
+ // %1 = fir.array_coor %arg [%slice] %idx
+ }
+ } else { // !op.getSlice()
+ if (!boxedSlice) {
+ // %0 = fir.rebox %arg
+ // %1 = fir.array_coor %0 %idx
+ // Pull in as:
+ // %1 = fir.array_coor %arg %idx
+ } else {
+ // %0 = fir.rebox %arg [%slice]
+ // %1 = fir.array_coor %0 %idx
+ // Pull in as:
+ // %1 = fir.array_coor %arg [%slice] %idx
+ }
+ }
+ }
+ } else { // op.getShape()
+ if (boxedShape) {
+ // Check if pulling in non-default shape is correct.
+ if (op.getSlice()) {
+ if (!boxedSlice) {
+ // %0 = fir.rebox %arg(%shape)
+ // %1 = fir.array_coor %0(%shape) [%slice] %idx
+ // Pull in as:
+ // %1 = fir.array_coor %arg(%shape) [%slice] %idx
+ } else {
+ // %0 = fir.rebox %arg(%shape) [%slice]
+ // %1 = fir.array_coor %0(%shape) [%slice] %idx
+ // Pull in as:
+ // %1 = fir.array_coor %arg(%shape) [%slice] %idx
+ }
+ } else { // !op.getSlice()
+ if (!boxedSlice) {
+ // %0 = fir.rebox %arg(%shape)
+ // %1 = fir.array_coor %0(%shape) %idx
+ // Pull in as:
+ // %1 = fir.array_coor %arg(%shape) %idx
+ } else {
+ // %0 = fir.rebox %arg(%shape) [%slice]
+ // %1 = fir.array_coor %0(%shape) %idx
+ // Cannot pull in without adjusting the array_coor indices.
+ return mlir::failure();
+ }
+ }
+ } else { // !boxedShape
+ if (op.getSlice()) {
+ if (!boxedSlice) {
+ // %0 = fir.rebox %arg
+ // %1 = fir.array_coor %0(%shape) [%slice] %idx
+ // Pull in as:
+ // %1 = fir.array_coor %arg(%shape) [%slice] %idx
+ } else {
+ // %0 = fir.rebox %arg [%slice]
+ // %1 = fir.array_coor %0(%shape) [%slice] %idx
+ return mlir::failure();
+ }
+ } else { // !op.getSlice()
+ if (!boxedSlice) {
+ // %0 = fir.rebox %arg
+ // %1 = fir.array_coor %0(%shape) %idx
+ // Pull in as:
+ // %1 = fir.array_coor %arg(%shape) %idx
+ } else {
+ // %0 = fir.rebox %arg [%slice]
+ // %1 = fir.array_coor %0(%shape) %idx
+ // Cannot pull in without adjusting the slice and array_coor
+ // indices.
+ return mlir::failure();
+ }
+ }
+ }
+ }
+
// TODO: temporarily avoid producing array_coor with the shape shift
// and plain array reference (it seems to be a limitation of
// ArrayCoorOp verifier).
diff --git a/flang/test/Fir/array-coor-canonicalization.fir b/flang/test/Fir/array-coor-canonicalization.fir
index ec36b2e57b9f9..86ca06ddb100f 100644
--- a/flang/test/Fir/array-coor-canonicalization.fir
+++ b/flang/test/Fir/array-coor-canonicalization.fir
@@ -119,3 +119,308 @@ func.func @_QPtest4() {
}
return
}
+
+// CHECK-LABEL: func.func @test5(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 10 : index
+// CHECK: %[[VAL_3:.*]] = fir.slice %[[VAL_1]], %[[VAL_2]], %[[VAL_1]] : (index, index, index) -> !fir.slice<1>
+// CHECK: %[[VAL_4:.*]] = fir.array_coor %[[VAL_0]] {{\[}}%[[VAL_3]]] %[[VAL_1]] : (!fir.ref<!fir.array<10xi32>>, !fir.slice<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_4]] : !fir.ref<i32>
+// CHECK: }
+func.func @test5(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %1 = fir.shape_shift %c10, %c10 : (index, index) -> !fir.shapeshift<1>
+ %s = fir.slice %c1, %c10, %c1 : (index, index, index) -> !fir.slice<1>
+ %2 = fir.embox %arg0(%1) : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>) -> !fir.box<!fir.array<10xi32>>
+ %3 = fir.array_coor %2 [%s] %c1 : (!fir.box<!fir.array<10xi32>>, !fir.slice<1>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test6(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 10 : index
+// CHECK: %[[VAL_3:.*]] = arith.constant 19 : index
+// CHECK: %[[VAL_4:.*]] = fir.shape_shift %[[VAL_2]], %[[VAL_2]] : (index, index) -> !fir.shapeshift<1>
+// CHECK: %[[VAL_5:.*]] = fir.slice %[[VAL_2]], %[[VAL_3]], %[[VAL_1]] : (index, index, index) -> !fir.slice<1>
+// CHECK: %[[VAL_6:.*]] = fir.embox %[[VAL_0]](%[[VAL_4]]) {{\[}}%[[VAL_5]]] : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>, !fir.slice<1>) -> !fir.box<!fir.array<10xi32>>
+// CHECK: %[[VAL_7:.*]] = fir.array_coor %[[VAL_6]] {{\[}}%[[VAL_5]]] %[[VAL_1]] : (!fir.box<!fir.array<10xi32>>, !fir.slice<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_7]] : !fir.ref<i32>
+// CHECK: }
+func.func @test6(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c19 = arith.constant 19 : index
+ %1 = fir.shape_shift %c10, %c10 : (index, index) -> !fir.shapeshift<1>
+ %s = fir.slice %c10, %c19, %c1 : (index, index, index) -> !fir.slice<1>
+ %2 = fir.embox %arg0(%1) [%s] : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>, !fir.slice<1>) -> !fir.box<!fir.array<10xi32>>
+ %3 = fir.array_coor %2 [%s] %c1 : (!fir.box<!fir.array<10xi32>>, !fir.slice<1>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test7(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = fir.array_coor %[[VAL_0]] %[[VAL_1]] : (!fir.ref<!fir.array<10xi32>>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_2]] : !fir.ref<i32>
+// CHECK: }
+func.func @test7(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %1 = fir.shape_shift %c10, %c10 : (index, index) -> !fir.shapeshift<1>
+ %2 = fir.embox %arg0(%1) : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>) -> !fir.box<!fir.array<10xi32>>
+ %3 = fir.array_coor %2 %c1 : (!fir.box<!fir.array<10xi32>>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test8(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 10 : index
+// CHECK: %[[VAL_3:.*]] = arith.constant 19 : index
+// CHECK: %[[VAL_4:.*]] = fir.shape_shift %[[VAL_2]], %[[VAL_2]] : (index, index) -> !fir.shapeshift<1>
+// CHECK: %[[VAL_5:.*]] = fir.slice %[[VAL_2]], %[[VAL_3]], %[[VAL_1]] : (index, index, index) -> !fir.slice<1>
+// CHECK: %[[VAL_6:.*]] = fir.array_coor %[[VAL_0]](%[[VAL_4]]) {{\[}}%[[VAL_5]]] %[[VAL_1]] : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>, !fir.slice<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_6]] : !fir.ref<i32>
+// CHECK: }
+func.func @test8(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c19 = arith.constant 19 : index
+ %1 = fir.shape_shift %c10, %c10 : (index, index) -> !fir.shapeshift<1>
+ %s = fir.slice %c10, %c19, %c1 : (index, index, index) -> !fir.slice<1>
+ %2 = fir.embox %arg0(%1) [%s] : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>, !fir.slice<1>) -> !fir.box<!fir.array<10xi32>>
+ %3 = fir.array_coor %2 %c1 : (!fir.box<!fir.array<10xi32>>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test9(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.box<!fir.array<?xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 9 : index
+// CHECK: %[[VAL_3:.*]] = fir.slice %[[VAL_1]], %[[VAL_2]], %[[VAL_1]] : (index, index, index) -> !fir.slice<1>
+// CHECK: %[[VAL_4:.*]] = fir.array_coor %[[VAL_0]] {{\[}}%[[VAL_3]]] %[[VAL_1]] : (!fir.box<!fir.array<?xi32>>, !fir.slice<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_4]] : !fir.ref<i32>
+// CHECK: }
+func.func @test9(%arg0: !fir.box<!fir.array<?xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c9 = arith.constant 9 : index
+ %s = fir.slice %c1, %c9, %c1 : (index, index, index) -> !fir.slice<1>
+ %2 = fir.rebox %arg0 : (!fir.box<!fir.array<?xi32>>) -> !fir.box<!fir.array<?xi32>>
+ %3 = fir.array_coor %2 [%s] %c1 : (!fir.box<!fir.array<?xi32>>, !fir.slice<1>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test10(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.box<!fir.array<?xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 9 : index
+// CHECK: %[[VAL_3:.*]] = fir.slice %[[VAL_1]], %[[VAL_2]], %[[VAL_1]] : (index, index, index) -> !fir.slice<1>
+// CHECK: %[[VAL_4:.*]] = fir.array_coor %[[VAL_0]] {{\[}}%[[VAL_3]]] %[[VAL_1]] : (!fir.box<!fir.array<?xi32>>, !fir.slice<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_4]] : !fir.ref<i32>
+// CHECK: }
+func.func @test10(%arg0: !fir.box<!fir.array<?xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c9 = arith.constant 9 : index
+ %s = fir.slice %c1, %c9, %c1 : (index, index, index) -> !fir.slice<1>
+ %2 = fir.rebox %arg0[%s] : (!fir.box<!fir.array<?xi32>>, !fir.slice<1>) -> !fir.box<!fir.array<?xi32>>
+ %3 = fir.array_coor %2 [%s] %c1 : (!fir.box<!fir.array<?xi32>>, !fir.slice<1>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test11(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.box<!fir.array<?xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = fir.array_coor %[[VAL_0]] %[[VAL_1]] : (!fir.box<!fir.array<?xi32>>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_2]] : !fir.ref<i32>
+// CHECK: }
+func.func @test11(%arg0: !fir.box<!fir.array<?xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %2 = fir.rebox %arg0 : (!fir.box<!fir.array<?xi32>>) -> !fir.box<!fir.array<?xi32>>
+ %3 = fir.array_coor %2 %c1 : (!fir.box<!fir.array<?xi32>>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test12(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.box<!fir.array<?xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 9 : index
+// CHECK: %[[VAL_3:.*]] = fir.slice %[[VAL_1]], %[[VAL_2]], %[[VAL_1]] : (index, index, index) -> !fir.slice<1>
+// CHECK: %[[VAL_4:.*]] = fir.array_coor %[[VAL_0]] {{\[}}%[[VAL_3]]] %[[VAL_1]] : (!fir.box<!fir.array<?xi32>>, !fir.slice<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_4]] : !fir.ref<i32>
+// CHECK: }
+func.func @test12(%arg0: !fir.box<!fir.array<?xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c9 = arith.constant 9 : index
+ %s = fir.slice %c1, %c9, %c1 : (index, index, index) -> !fir.slice<1>
+ %2 = fir.rebox %arg0[%s] : (!fir.box<!fir.array<?xi32>>, !fir.slice<1>) -> !fir.box<!fir.array<?xi32>>
+ %3 = fir.array_coor %2 %c1 : (!fir.box<!fir.array<?xi32>>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test13(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 10 : index
+// CHECK: %[[VAL_3:.*]] = arith.constant 19 : index
+// CHECK: %[[VAL_4:.*]] = fir.shape_shift %[[VAL_2]], %[[VAL_2]] : (index, index) -> !fir.shapeshift<1>
+// CHECK: %[[VAL_5:.*]] = fir.slice %[[VAL_2]], %[[VAL_3]], %[[VAL_1]] : (index, index, index) -> !fir.slice<1>
+// CHECK: %[[VAL_6:.*]] = fir.array_coor %[[VAL_0]](%[[VAL_4]]) {{\[}}%[[VAL_5]]] %[[VAL_1]] : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>, !fir.slice<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_6]] : !fir.ref<i32>
+// CHECK: }
+func.func @test13(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c19 = arith.constant 19 : index
+ %1 = fir.shape_shift %c10, %c10 : (index, index) -> !fir.shapeshift<1>
+ %s = fir.slice %c10, %c19, %c1 : (index, index, index) -> !fir.slice<1>
+ %2 = fir.embox %arg0(%1) : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>) -> !fir.box<!fir.array<10xi32>>
+ %3 = fir.array_coor %2(%1) [%s] %c1 : (!fir.box<!fir.array<10xi32>>, !fir.shapeshift<1>, !fir.slice<1>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test14(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 10 : index
+// CHECK: %[[VAL_3:.*]] = arith.constant 19 : index
+// CHECK: %[[VAL_4:.*]] = fir.shape_shift %[[VAL_2]], %[[VAL_2]] : (index, index) -> !fir.shapeshift<1>
+// CHECK: %[[VAL_5:.*]] = fir.slice %[[VAL_2]], %[[VAL_3]], %[[VAL_1]] : (index, index, index) -> !fir.slice<1>
+// CHECK: %[[VAL_6:.*]] = fir.array_coor %[[VAL_0]](%[[VAL_4]]) {{\[}}%[[VAL_5]]] %[[VAL_1]] : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>, !fir.slice<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_6]] : !fir.ref<i32>
+// CHECK: }
+func.func @test14(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c19 = arith.constant 19 : index
+ %1 = fir.shape_shift %c10, %c10 : (index, index) -> !fir.shapeshift<1>
+ %s = fir.slice %c10, %c19, %c1 : (index, index, index) -> !fir.slice<1>
+ %2 = fir.embox %arg0(%1) [%s] : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>, !fir.slice<1>) -> !fir.box<!fir.array<10xi32>>
+ %3 = fir.array_coor %2(%1) [%s] %c1 : (!fir.box<!fir.array<10xi32>>, !fir.shapeshift<1>, !fir.slice<1>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test15(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 10 : index
+// CHECK: %[[VAL_2:.*]] = fir.shape_shift %[[VAL_1]], %[[VAL_1]] : (index, index) -> !fir.shapeshift<1>
+// CHECK: %[[VAL_3:.*]] = fir.array_coor %[[VAL_0]](%[[VAL_2]]) %[[VAL_1]] : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_3]] : !fir.ref<i32>
+// CHECK: }
+func.func @test15(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+ %c10 = arith.constant 10 : index
+ %1 = fir.shape_shift %c10, %c10 : (index, index) -> !fir.shapeshift<1>
+ %2 = fir.embox %arg0(%1) : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>) -> !fir.box<!fir.array<10xi32>>
+ %3 = fir.array_coor %2(%1) %c10 : (!fir.box<!fir.array<10xi32>>, !fir.shapeshift<1>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test16(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 10 : index
+// CHECK: %[[VAL_3:.*]] = arith.constant 19 : index
+// CHECK: %[[VAL_4:.*]] = fir.shape_shift %[[VAL_2]], %[[VAL_2]] : (index, index) -> !fir.shapeshift<1>
+// CHECK: %[[VAL_5:.*]] = fir.slice %[[VAL_2]], %[[VAL_3]], %[[VAL_1]] : (index, index, index) -> !fir.slice<1>
+// CHECK: %[[VAL_6:.*]] = fir.embox %[[VAL_0]](%[[VAL_4]]) {{\[}}%[[VAL_5]]] : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>, !fir.slice<1>) -> !fir.box<!fir.array<10xi32>>
+// CHECK: %[[VAL_7:.*]] = fir.array_coor %[[VAL_6]](%[[VAL_4]]) %[[VAL_2]] : (!fir.box<!fir.array<10xi32>>, !fir.shapeshift<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_7]] : !fir.ref<i32>
+// CHECK: }
+func.func @test16(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c19 = arith.constant 19 : index
+ %1 = fir.shape_shift %c10, %c10 : (index, index) -> !fir.shapeshift<1>
+ %s = fir.slice %c10, %c19, %c1 : (index, index, index) -> !fir.slice<1>
+ %2 = fir.embox %arg0(%1) [%s] : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>, !fir.slice<1>) -> !fir.box<!fir.array<10xi32>>
+ %3 = fir.array_coor %2(%1) %c10 : (!fir.box<!fir.array<10xi32>>, !fir.shapeshift<1>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test17(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.box<!fir.array<?xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 10 : index
+// CHECK: %[[VAL_3:.*]] = arith.constant 19 : index
+// CHECK: %[[VAL_4:.*]] = fir.shape_shift %[[VAL_2]], %[[VAL_2]] : (index, index) -> !fir.shapeshift<1>
+// CHECK: %[[VAL_5:.*]] = fir.slice %[[VAL_2]], %[[VAL_3]], %[[VAL_1]] : (index, index, index) -> !fir.slice<1>
+// CHECK: %[[VAL_6:.*]] = fir.array_coor %[[VAL_0]](%[[VAL_4]]) {{\[}}%[[VAL_5]]] %[[VAL_1]] : (!fir.box<!fir.array<?xi32>>, !fir.shapeshift<1>, !fir.slice<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_6]] : !fir.ref<i32>
+// CHECK: }
+func.func @test17(%arg0: !fir.box<!fir.array<?xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c19 = arith.constant 19 : index
+ %1 = fir.shape_shift %c10, %c10 : (index, index) -> !fir.shapeshift<1>
+ %s = fir.slice %c10, %c19, %c1 : (index, index, index) -> !fir.slice<1>
+ %2 = fir.rebox %arg0 : (!fir.box<!fir.array<?xi32>>) -> !fir.box<!fir.array<?xi32>>
+ %3 = fir.array_coor %2(%1) [%s] %c1 : (!fir.box<!fir.array<?xi32>>, !fir.shapeshift<1>, !fir.slice<1>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test18(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 10 : index
+// CHECK: %[[VAL_3:.*]] = arith.constant 9 : index
+// CHECK: %[[VAL_4:.*]] = fir.shape_shift %[[VAL_1]], %[[VAL_2]] : (index, index) -> !fir.shapeshift<1>
+// CHECK: %[[VAL_5:.*]] = fir.slice %[[VAL_1]], %[[VAL_3]], %[[VAL_1]] : (index, index, index) -> !fir.slice<1>
+// CHECK: %[[VAL_6:.*]] = fir.embox %[[VAL_0]] {{\[}}%[[VAL_5]]] : (!fir.ref<!fir.array<10xi32>>, !fir.slice<1>) -> !fir.box<!fir.array<10xi32>>
+// CHECK: %[[VAL_7:.*]] = fir.array_coor %[[VAL_6]](%[[VAL_4]]) {{\[}}%[[VAL_5]]] %[[VAL_1]] : (!fir.box<!fir.array<10xi32>>, !fir.shapeshift<1>, !fir.slice<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_7]] : !fir.ref<i32>
+// CHECK: }
+func.func @test18(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c9 = arith.constant 9 : index
+ %1 = fir.shape_shift %c1, %c10 : (index, index) -> !fir.shapeshift<1>
+ %s = fir.slice %c1, %c9, %c1 : (index, index, index) -> !fir.slice<1>
+ %2 = fir.embox %arg0 [%s] : (!fir.ref<!fir.array<10xi32>>, !fir.slice<1>) -> !fir.box<!fir.array<10xi32>>
+ %3 = fir.array_coor %2(%1) [%s] %c1 : (!fir.box<!fir.array<10xi32>>, !fir.shapeshift<1>, !fir.slice<1>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test19(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.box<!fir.array<?xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 10 : index
+// CHECK: %[[VAL_2:.*]] = fir.shape_shift %[[VAL_1]], %[[VAL_1]] : (index, index) -> !fir.shapeshift<1>
+// CHECK: %[[VAL_3:.*]] = fir.array_coor %[[VAL_0]](%[[VAL_2]]) %[[VAL_1]] : (!fir.box<!fir.array<?xi32>>, !fir.shapeshift<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_3]] : !fir.ref<i32>
+// CHECK: }
+func.func @test19(%arg0: !fir.box<!fir.array<?xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c19 = arith.constant 19 : index
+ %1 = fir.shape_shift %c10, %c10 : (index, index) -> !fir.shapeshift<1>
+ %s = fir.slice %c10, %c19, %c1 : (index, index, index) -> !fir.slice<1>
+ %2 = fir.rebox %arg0 : (!fir.box<!fir.array<?xi32>>) -> !fir.box<!fir.array<?xi32>>
+ %3 = fir.array_coor %2(%1) %c10 : (!fir.box<!fir.array<?xi32>>, !fir.shapeshift<1>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test20(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 10 : index
+// CHECK: %[[VAL_3:.*]] = arith.constant 9 : index
+// CHECK: %[[VAL_4:.*]] = fir.shape_shift %[[VAL_2]], %[[VAL_2]] : (index, index) -> !fir.shapeshift<1>
+// CHECK: %[[VAL_5:.*]] = fir.slice %[[VAL_1]], %[[VAL_3]], %[[VAL_1]] : (index, index, index) -> !fir.slice<1>
+// CHECK: %[[VAL_6:.*]] = fir.embox %[[VAL_0]] {{\[}}%[[VAL_5]]] : (!fir.ref<!fir.array<10xi32>>, !fir.slice<1>) -> !fir.box<!fir.array<10xi32>>
+// CHECK: %[[VAL_7:.*]] = fir.array_coor %[[VAL_6]](%[[VAL_4]]) %[[VAL_2]] : (!fir.box<!fir.array<10xi32>>, !fir.shapeshift<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_7]] : !fir.ref<i32>
+// CHECK: }
+func.func @test20(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c9 = arith.constant 9 : index
+ %1 = fir.shape_shift %c10, %c10 : (index, index) -> !fir.shapeshift<1>
+ %s = fir.slice %c1, %c9, %c1 : (index, index, index) -> !fir.slice<1>
+ %2 = fir.embox %arg0 [%s] : (!fir.ref<!fir.array<10xi32>>, !fir.slice<1>) -> !fir.box<!fir.array<10xi32>>
+ %3 = fir.array_coor %2(%1) %c10 : (!fir.box<!fir.array<10xi32>>, !fir.shapeshift<1>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
>From 43a70e6cb598fd86ad906ad60eaaf9a9becb487d Mon Sep 17 00:00:00 2001
From: Slava Zakharin <szakharin at nvidia.com>
Date: Wed, 22 May 2024 19:48:47 -0700
Subject: [PATCH 3/5] Fixed more cases.
---
flang/lib/Optimizer/Dialect/FIROps.cpp | 137 +++++++++++++----
.../test/Fir/array-coor-canonicalization.fir | 139 ++++++++++++++++--
2 files changed, 240 insertions(+), 36 deletions(-)
diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index 5c5123273da8a..afe5893616cac 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -448,11 +448,22 @@ struct SimplifyArrayCoorOp : public mlir::OpRewritePattern<fir::ArrayCoorOp> {
memref.getDefiningOp())) {
boxedMemref = reboxOp.getBox();
boxedShape = reboxOp.getShape();
+ // Avoid pulling in rebox that performs reshaping.
+ // There is no way to represent box reshaping with array_coor.
+ if (boxedShape && !mlir::isa<fir::ShiftType>(boxedShape.getType()))
+ return mlir::failure();
boxedSlice = reboxOp.getSlice();
} else {
return mlir::failure();
}
+ bool boxedShapeIsShift =
+ boxedShape && mlir::isa<fir::ShiftType>(boxedShape.getType());
+ bool boxedShapeIsShape =
+ boxedShape && mlir::isa<fir::ShapeType>(boxedShape.getType());
+ bool boxedShapeIsShapeShift =
+ boxedShape && mlir::isa<fir::ShapeShiftType>(boxedShape.getType());
+
// Slices changing the number of dimensions are not supported
// for array_coor yet.
unsigned origBoxRank;
@@ -481,6 +492,21 @@ struct SimplifyArrayCoorOp : public mlir::OpRewritePattern<fir::ArrayCoorOp> {
if (op.getSlice() && boxedSlice && boxedSlice != op.getSlice())
return mlir::failure();
+ // If v is a shape_shift operation:
+ // fir.shape_shift %l1, %e1, %l2, %e2, ...
+ // create:
+ // fir.shape %e1, %e2, ...
+ auto getShapeFromShapeShift = [&](mlir::Value v) -> mlir::Value {
+ auto shapeShiftOp =
+ mlir::dyn_cast_or_null<fir::ShapeShiftOp>(v.getDefiningOp());
+ if (!shapeShiftOp)
+ return nullptr;
+ mlir::OpBuilder::InsertionGuard guard(rewriter);
+ rewriter.setInsertionPointAfter(shapeShiftOp);
+ return rewriter.create<fir::ShapeOp>(shapeShiftOp.getLoc(),
+ shapeShiftOp.getExtents());
+ };
+
// The embox/rebox and array_coor either have compatible
// shape/slice at this point or shape/slice is null
// in one of them but not in the other.
@@ -489,32 +515,90 @@ struct SimplifyArrayCoorOp : public mlir::OpRewritePattern<fir::ArrayCoorOp> {
if (boxedShape) {
if (op.getSlice()) {
if (!boxedSlice) {
- // %0 = fir.rebox %arg(%shape)
- // %1 = fir.array_coor %0 [%slice] %idx
- // Both the slice indices and %idx are 1-based, so the rebox
- // may be pulled in as:
- // %1 = fir.array_coor %arg [%slice] %idx
- boxedShape = nullptr;
+ if (boxedShapeIsShift) {
+ // %0 = fir.rebox %arg(%shift)
+ // %1 = fir.array_coor %0 [%slice] %idx
+ // Both the slice indices and %idx are 1-based, so the rebox
+ // may be pulled in as:
+ // %1 = fir.array_coor %arg [%slice] %idx
+ boxedShape = nullptr;
+ } else if (boxedShapeIsShape) {
+ // %0 = fir.embox %arg(%shape)
+ // %1 = fir.array_coor %0 [%slice] %idx
+ // Pull in as:
+ // %1 = fir.array_coor %arg(%shape) [%slice] %idx
+ } else if (boxedShapeIsShapeShift) {
+ // %0 = fir.embox %arg(%shapeshift)
+ // %1 = fir.array_coor %0 [%slice] %idx
+ // Pull in as:
+ // %shape = fir.shape <extents from the %shapeshift>
+ // %1 = fir.array_coor %arg(%shape) [%slice] %idx
+ boxedShape = getShapeFromShapeShift(boxedShape);
+ if (!boxedShape)
+ return mlir::failure();
+ } else {
+ return mlir::failure();
+ }
} else {
- // %0 = fir.rebox %arg(%shape) [%slice]
- // %1 = fir.array_coor %0 [%slice] %idx
- // I believe this FIR may only be valid if the shape specifies
- // that all lower bounds are 1s and the slice's start indices
- // are all 1s.
- // We could pull in the rebox as:
- // %1 = fir.array_coor %arg [%slice] %idx
- // Do not do anything for the time being.
- return mlir::failure();
+ if (boxedShapeIsShift) {
+ // %0 = fir.rebox %arg(%shift) [%slice]
+ // %1 = fir.array_coor %0 [%slice] %idx
+ // This FIR may only be valid if the shape specifies
+ // that all lower bounds are 1s and the slice's start indices
+ // and strides are all 1s.
+ // We could pull in the rebox as:
+ // %1 = fir.array_coor %arg [%slice] %idx
+ // Do not do anything for the time being.
+ return mlir::failure();
+ } else if (boxedShapeIsShape) {
+ // %0 = fir.embox %arg(%shape) [%slice]
+ // %1 = fir.array_coor %0 [%slice] %idx
+ // This FIR may only be valid if the slice's start indices
+ // and strides are all 1s.
+ // We could pull in the embox as:
+ // %1 = fir.array_coor %arg(%shape) [%slice] %idx
+ return mlir::failure();
+ } else if (boxedShapeIsShapeShift) {
+ // %0 = fir.embox %arg(%shapeshift) [%slice]
+ // %1 = fir.array_coor %0 [%slice] %idx
+ // This FIR may only be valid if the shape specifies
+ // that all lower bounds are 1s and the slice's start indices
+ // and strides are all 1s.
+ // We could pull in the embox as:
+ // %shape = fir.shape <extents from the %shapeshift>
+ // %1 = fir.array_coor %arg(%shape) [%slice] %idx
+ return mlir::failure();
+ } else {
+ return mlir::failure();
+ }
}
} else { // !op.getSlice()
if (!boxedSlice) {
- // %0 = fir.rebox %arg(%shape)
- // %1 = fir.array_coor %0 %idx
- // Pull in as:
- // %1 = fir.array_coor %arg %idx
- boxedShape = nullptr;
+ if (boxedShapeIsShift) {
+ // %0 = fir.rebox %arg(%shift)
+ // %1 = fir.array_coor %0 %idx
+ // Pull in as:
+ // %1 = fir.array_coor %arg %idx
+ boxedShape = nullptr;
+ } else if (boxedShapeIsShape) {
+ // %0 = fir.embox %arg(%shape)
+ // %1 = fir.array_coor %0 %idx
+ // Pull in as:
+ // %1 = fir.array_coor %arg(%shape) %idx
+ } else if (boxedShapeIsShapeShift) {
+ // %0 = fir.embox %arg(%shapeshift)
+ // %1 = fir.array_coor %0 %idx
+ // Pull in as:
+ // %shape = fir.shape <extents from the %shapeshift>
+ // %1 = fir.array_coor %arg(%shape) %idx
+ boxedShape = getShapeFromShapeShift(boxedShape);
+ if (!boxedShape)
+ return mlir::failure();
+ } else {
+ return mlir::failure();
+ }
} else {
- // %0 = fir.rebox %arg(%shape) [%slice]
+ // %0 = fir.embox %arg(%shape) [%slice]
// %1 = fir.array_coor %0 %idx
// Pull in as:
// %1 = fir.array_coor %arg(%shape) %[slice] %idx
@@ -530,6 +614,8 @@ struct SimplifyArrayCoorOp : public mlir::OpRewritePattern<fir::ArrayCoorOp> {
} else {
// %0 = fir.rebox %arg [%slice]
// %1 = fir.array_coor %0 [%slice] %idx
+ // This is a valid FIR iff the slice's lower bounds
+ // and strides are all 1s.
// Pull in as:
// %1 = fir.array_coor %arg [%slice] %idx
}
@@ -552,24 +638,24 @@ struct SimplifyArrayCoorOp : public mlir::OpRewritePattern<fir::ArrayCoorOp> {
// Check if pulling in non-default shape is correct.
if (op.getSlice()) {
if (!boxedSlice) {
- // %0 = fir.rebox %arg(%shape)
+ // %0 = fir.embox %arg(%shape)
// %1 = fir.array_coor %0(%shape) [%slice] %idx
// Pull in as:
// %1 = fir.array_coor %arg(%shape) [%slice] %idx
} else {
- // %0 = fir.rebox %arg(%shape) [%slice]
+ // %0 = fir.embox %arg(%shape) [%slice]
// %1 = fir.array_coor %0(%shape) [%slice] %idx
// Pull in as:
// %1 = fir.array_coor %arg(%shape) [%slice] %idx
}
} else { // !op.getSlice()
if (!boxedSlice) {
- // %0 = fir.rebox %arg(%shape)
+ // %0 = fir.embox %arg(%shape)
// %1 = fir.array_coor %0(%shape) %idx
// Pull in as:
// %1 = fir.array_coor %arg(%shape) %idx
} else {
- // %0 = fir.rebox %arg(%shape) [%slice]
+ // %0 = fir.embox %arg(%shape) [%slice]
// %1 = fir.array_coor %0(%shape) %idx
// Cannot pull in without adjusting the array_coor indices.
return mlir::failure();
@@ -630,6 +716,7 @@ struct SimplifyArrayCoorOp : public mlir::OpRewritePattern<fir::ArrayCoorOp> {
void fir::ArrayCoorOp::getCanonicalizationPatterns(
mlir::RewritePatternSet &patterns, mlir::MLIRContext *context) {
+ // TODO: !fir.shape<1> operand may be removed from array_coor always.
patterns.add<SimplifyArrayCoorOp>(context);
}
diff --git a/flang/test/Fir/array-coor-canonicalization.fir b/flang/test/Fir/array-coor-canonicalization.fir
index 86ca06ddb100f..63c3f98672557 100644
--- a/flang/test/Fir/array-coor-canonicalization.fir
+++ b/flang/test/Fir/array-coor-canonicalization.fir
@@ -120,15 +120,54 @@ func.func @_QPtest4() {
return
}
-// CHECK-LABEL: func.func @test5(
-// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+// CHECK-LABEL: func.func @test5_1(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.box<!fir.array<?xi32>>) -> !fir.ref<i32> {
// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
-// CHECK: %[[VAL_2:.*]] = arith.constant 10 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 9 : index
// CHECK: %[[VAL_3:.*]] = fir.slice %[[VAL_1]], %[[VAL_2]], %[[VAL_1]] : (index, index, index) -> !fir.slice<1>
-// CHECK: %[[VAL_4:.*]] = fir.array_coor %[[VAL_0]] {{\[}}%[[VAL_3]]] %[[VAL_1]] : (!fir.ref<!fir.array<10xi32>>, !fir.slice<1>, index) -> !fir.ref<i32>
+// CHECK: %[[VAL_4:.*]] = fir.array_coor %[[VAL_0]] {{\[}}%[[VAL_3]]] %[[VAL_1]] : (!fir.box<!fir.array<?xi32>>, !fir.slice<1>, index) -> !fir.ref<i32>
// CHECK: return %[[VAL_4]] : !fir.ref<i32>
// CHECK: }
-func.func @test5(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+func.func @test5_1(%arg0: !fir.box<!fir.array<?xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c9 = arith.constant 9 : index
+ %sh = fir.shift %c10 : (index) -> !fir.shift<1>
+ %2 = fir.rebox %arg0(%sh) : (!fir.box<!fir.array<?xi32>>, !fir.shift<1>) -> !fir.box<!fir.array<?xi32>>
+ %s = fir.slice %c1, %c9, %c1 : (index, index, index) -> !fir.slice<1>
+ %3 = fir.array_coor %2 [%s] %c1 : (!fir.box<!fir.array<?xi32>>, !fir.slice<1>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test5_2(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 10 : index
+// CHECK: %[[VAL_3:.*]] = fir.shape %[[VAL_2]] : (index) -> !fir.shape<1>
+// CHECK: %[[VAL_4:.*]] = fir.slice %[[VAL_1]], %[[VAL_2]], %[[VAL_1]] : (index, index, index) -> !fir.slice<1>
+// CHECK: %[[VAL_5:.*]] = fir.array_coor %[[VAL_0]](%[[VAL_3]]) {{\[}}%[[VAL_4]]] %[[VAL_1]] : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>, !fir.slice<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_5]] : !fir.ref<i32>
+// CHECK: }
+func.func @test5_2(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %1 = fir.shape %c10 : (index) -> !fir.shape<1>
+ %s = fir.slice %c1, %c10, %c1 : (index, index, index) -> !fir.slice<1>
+ %2 = fir.embox %arg0(%1) : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>) -> !fir.box<!fir.array<10xi32>>
+ %3 = fir.array_coor %2 [%s] %c1 : (!fir.box<!fir.array<10xi32>>, !fir.slice<1>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test5_3(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 10 : index
+// CHECK: %[[VAL_3:.*]] = fir.shape %[[VAL_2]] : (index) -> !fir.shape<1>
+// CHECK: %[[VAL_4:.*]] = fir.slice %[[VAL_1]], %[[VAL_2]], %[[VAL_1]] : (index, index, index) -> !fir.slice<1>
+// CHECK: %[[VAL_5:.*]] = fir.array_coor %[[VAL_0]](%[[VAL_3]]) {{\[}}%[[VAL_4]]] %[[VAL_1]] : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>, !fir.slice<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_5]] : !fir.ref<i32>
+// CHECK: }
+func.func @test5_3(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
%c1 = arith.constant 1 : index
%c10 = arith.constant 10 : index
%1 = fir.shape_shift %c10, %c10 : (index, index) -> !fir.shapeshift<1>
@@ -138,7 +177,50 @@ func.func @test5(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
return %3 : !fir.ref<i32>
}
-// CHECK-LABEL: func.func @test6(
+// CHECK-LABEL: func.func @test6_1(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.box<!fir.array<?xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 9 : index
+// CHECK: %[[VAL_3:.*]] = fir.shift %[[VAL_1]] : (index) -> !fir.shift<1>
+// CHECK: %[[VAL_4:.*]] = fir.slice %[[VAL_1]], %[[VAL_2]], %[[VAL_1]] : (index, index, index) -> !fir.slice<1>
+// CHECK: %[[VAL_5:.*]] = fir.rebox %[[VAL_0]](%[[VAL_3]]) {{\[}}%[[VAL_4]]] : (!fir.box<!fir.array<?xi32>>, !fir.shift<1>, !fir.slice<1>) -> !fir.box<!fir.array<?xi32>>
+// CHECK: %[[VAL_6:.*]] = fir.array_coor %[[VAL_5]] {{\[}}%[[VAL_4]]] %[[VAL_1]] : (!fir.box<!fir.array<?xi32>>, !fir.slice<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_6]] : !fir.ref<i32>
+// CHECK: }
+func.func @test6_1(%arg0: !fir.box<!fir.array<?xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c9 = arith.constant 9 : index
+ %sh = fir.shift %c1 : (index) -> !fir.shift<1>
+ %s = fir.slice %c1, %c9, %c1 : (index, index, index) -> !fir.slice<1>
+ %2 = fir.rebox %arg0(%sh) [%s] : (!fir.box<!fir.array<?xi32>>, !fir.shift<1>, !fir.slice<1>) -> !fir.box<!fir.array<?xi32>>
+ %3 = fir.array_coor %2 [%s] %c1 : (!fir.box<!fir.array<?xi32>>, !fir.slice<1>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test6_2(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 10 : index
+// CHECK: %[[VAL_3:.*]] = arith.constant 19 : index
+// CHECK: %[[VAL_4:.*]] = fir.shape %[[VAL_2]] : (index) -> !fir.shape<1>
+// CHECK: %[[VAL_5:.*]] = fir.slice %[[VAL_2]], %[[VAL_3]], %[[VAL_1]] : (index, index, index) -> !fir.slice<1>
+// CHECK: %[[VAL_6:.*]] = fir.embox %[[VAL_0]](%[[VAL_4]]) {{\[}}%[[VAL_5]]] : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>, !fir.slice<1>) -> !fir.box<!fir.array<10xi32>>
+// CHECK: %[[VAL_7:.*]] = fir.array_coor %[[VAL_6]] {{\[}}%[[VAL_5]]] %[[VAL_1]] : (!fir.box<!fir.array<10xi32>>, !fir.slice<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_7]] : !fir.ref<i32>
+// CHECK: }
+func.func @test6_2(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c19 = arith.constant 19 : index
+ %1 = fir.shape %c10 : (index) -> !fir.shape<1>
+ %s = fir.slice %c10, %c19, %c1 : (index, index, index) -> !fir.slice<1>
+ %2 = fir.embox %arg0(%1) [%s] : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>, !fir.slice<1>) -> !fir.box<!fir.array<10xi32>>
+ %3 = fir.array_coor %2 [%s] %c1 : (!fir.box<!fir.array<10xi32>>, !fir.slice<1>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test6_3(
// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
// CHECK: %[[VAL_2:.*]] = arith.constant 10 : index
@@ -149,7 +231,7 @@ func.func @test5(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
// CHECK: %[[VAL_7:.*]] = fir.array_coor %[[VAL_6]] {{\[}}%[[VAL_5]]] %[[VAL_1]] : (!fir.box<!fir.array<10xi32>>, !fir.slice<1>, index) -> !fir.ref<i32>
// CHECK: return %[[VAL_7]] : !fir.ref<i32>
// CHECK: }
-func.func @test6(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+func.func @test6_3(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
%c1 = arith.constant 1 : index
%c10 = arith.constant 10 : index
%c19 = arith.constant 19 : index
@@ -160,13 +242,48 @@ func.func @test6(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
return %3 : !fir.ref<i32>
}
-// CHECK-LABEL: func.func @test7(
-// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+// CHECK-LABEL: func.func @test7_1(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.box<!fir.array<?xi32>>) -> !fir.ref<i32> {
// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
-// CHECK: %[[VAL_2:.*]] = fir.array_coor %[[VAL_0]] %[[VAL_1]] : (!fir.ref<!fir.array<10xi32>>, index) -> !fir.ref<i32>
+// CHECK: %[[VAL_2:.*]] = fir.array_coor %[[VAL_0]] %[[VAL_1]] : (!fir.box<!fir.array<?xi32>>, index) -> !fir.ref<i32>
// CHECK: return %[[VAL_2]] : !fir.ref<i32>
// CHECK: }
-func.func @test7(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+func.func @test7_1(%arg0: !fir.box<!fir.array<?xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c9 = arith.constant 9 : index
+ %sh = fir.shift %c10 : (index) -> !fir.shift<1>
+ %2 = fir.rebox %arg0(%sh) : (!fir.box<!fir.array<?xi32>>, !fir.shift<1>) -> !fir.box<!fir.array<?xi32>>
+ %3 = fir.array_coor %2 %c1 : (!fir.box<!fir.array<?xi32>>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test7_2(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 10 : index
+// CHECK: %[[VAL_3:.*]] = fir.shape %[[VAL_2]] : (index) -> !fir.shape<1>
+// CHECK: %[[VAL_4:.*]] = fir.array_coor %[[VAL_0]](%[[VAL_3]]) %[[VAL_1]] : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_4]] : !fir.ref<i32>
+// CHECK: }
+func.func @test7_2(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %1 = fir.shape %c10 : (index) -> !fir.shape<1>
+ %2 = fir.embox %arg0(%1) : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>) -> !fir.box<!fir.array<10xi32>>
+ %3 = fir.array_coor %2 %c1 : (!fir.box<!fir.array<10xi32>>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test7_3(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 10 : index
+// CHECK: %[[VAL_3:.*]] = fir.shape %[[VAL_2]] : (index) -> !fir.shape<1>
+// CHECK: %[[VAL_4:.*]] = fir.array_coor %[[VAL_0]](%[[VAL_3]]) %[[VAL_1]] : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_4]] : !fir.ref<i32>
+// CHECK: }
+func.func @test7_3(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
%c1 = arith.constant 1 : index
%c10 = arith.constant 10 : index
%1 = fir.shape_shift %c10, %c10 : (index, index) -> !fir.shapeshift<1>
>From 4f3a0040a0663be0067ecb88ccefc1ff4032e209 Mon Sep 17 00:00:00 2001
From: Slava Zakharin <szakharin at nvidia.com>
Date: Tue, 28 May 2024 15:42:09 -0700
Subject: [PATCH 4/5] Fixed typos.
---
flang/lib/Optimizer/Dialect/FIROps.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index afe5893616cac..9f996c36f9be4 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -402,8 +402,8 @@ mlir::LogicalResult fir::ArrayCoorOp::verify() {
}
if (arrDim && arrDim != shapeTyRank)
return emitOpError("rank of dimension mismatched");
- // TODO: support slicing with changing the numbder of dimensions,
- // e.g. when array_coord represents an element access to array(:,1,:)
+ // TODO: support slicing with changing the number of dimensions,
+ // e.g. when array_coor represents an element access to array(:,1,:)
// slice: the shape is 3D and the number of indices is 2 in this case.
if (shapeTyRank != getIndices().size())
return emitOpError("number of indices do not match dim rank");
>From c3bb9485b1768c9ec288df3482ab83dd3887087b Mon Sep 17 00:00:00 2001
From: Slava Zakharin <szakharin at nvidia.com>
Date: Wed, 29 May 2024 17:51:56 -0700
Subject: [PATCH 5/5] Fixed the patterns to follow array_coor specification.
---
flang/lib/Optimizer/CodeGen/CodeGen.cpp | 6 +-
flang/lib/Optimizer/Dialect/FIROps.cpp | 124 ++++++++++++++----
.../test/Fir/array-coor-canonicalization.fir | 73 +++++++++--
flang/test/Fir/convert-to-llvm.fir | 4 +-
4 files changed, 163 insertions(+), 44 deletions(-)
diff --git a/flang/lib/Optimizer/CodeGen/CodeGen.cpp b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
index 8d70ceb31bce1..72172f63888e1 100644
--- a/flang/lib/Optimizer/CodeGen/CodeGen.cpp
+++ b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
@@ -2162,11 +2162,7 @@ struct XArrayCoorOpConversion
if (normalSlice)
step = integerCast(loc, rewriter, idxTy, operands[sliceOffset + 2]);
}
- // The array_coor indices are always 1-based if slicing
- // is in effect. The non-default lower bounds only
- // apply to the indices of the slice itself.
- auto idx = rewriter.create<mlir::LLVM::SubOp>(loc, idxTy, index,
- isSliced ? one : lb, nsw);
+ auto idx = rewriter.create<mlir::LLVM::SubOp>(loc, idxTy, index, lb, nsw);
mlir::Value diff =
rewriter.create<mlir::LLVM::MulOp>(loc, idxTy, idx, step, nsw);
if (normalSlice) {
diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index 9f996c36f9be4..69e13006645c0 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -492,21 +492,7 @@ struct SimplifyArrayCoorOp : public mlir::OpRewritePattern<fir::ArrayCoorOp> {
if (op.getSlice() && boxedSlice && boxedSlice != op.getSlice())
return mlir::failure();
- // If v is a shape_shift operation:
- // fir.shape_shift %l1, %e1, %l2, %e2, ...
- // create:
- // fir.shape %e1, %e2, ...
- auto getShapeFromShapeShift = [&](mlir::Value v) -> mlir::Value {
- auto shapeShiftOp =
- mlir::dyn_cast_or_null<fir::ShapeShiftOp>(v.getDefiningOp());
- if (!shapeShiftOp)
- return nullptr;
- mlir::OpBuilder::InsertionGuard guard(rewriter);
- rewriter.setInsertionPointAfter(shapeShiftOp);
- return rewriter.create<fir::ShapeOp>(shapeShiftOp.getLoc(),
- shapeShiftOp.getExtents());
- };
-
+ std::optional<IndicesVectorTy> shiftedIndices;
// The embox/rebox and array_coor either have compatible
// shape/slice at this point or shape/slice is null
// in one of them but not in the other.
@@ -533,7 +519,7 @@ struct SimplifyArrayCoorOp : public mlir::OpRewritePattern<fir::ArrayCoorOp> {
// Pull in as:
// %shape = fir.shape <extents from the %shapeshift>
// %1 = fir.array_coor %arg(%shape) [%slice] %idx
- boxedShape = getShapeFromShapeShift(boxedShape);
+ boxedShape = getShapeFromShapeShift(boxedShape, rewriter);
if (!boxedShape)
return mlir::failure();
} else {
@@ -591,17 +577,43 @@ struct SimplifyArrayCoorOp : public mlir::OpRewritePattern<fir::ArrayCoorOp> {
// Pull in as:
// %shape = fir.shape <extents from the %shapeshift>
// %1 = fir.array_coor %arg(%shape) %idx
- boxedShape = getShapeFromShapeShift(boxedShape);
+ boxedShape = getShapeFromShapeShift(boxedShape, rewriter);
if (!boxedShape)
return mlir::failure();
} else {
return mlir::failure();
}
} else {
- // %0 = fir.embox %arg(%shape) [%slice]
- // %1 = fir.array_coor %0 %idx
- // Pull in as:
- // %1 = fir.array_coor %arg(%shape) %[slice] %idx
+ if (boxedShapeIsShift) {
+ // %0 = fir.embox %arg(%shift) [%slice]
+ // %1 = fir.array_coor %0 %idx
+ // Pull in as:
+ // %tmp = arith.addi %idx, %shift.origin
+ // %idx_shifted = arith.subi %tmp, 1
+ // %1 = fir.array_coor %arg(%shift) %[slice] %idx_shifted
+ shiftedIndices =
+ getShiftedIndices(boxedShape, op.getIndices(), rewriter);
+ if (!shiftedIndices)
+ return mlir::failure();
+ } else if (boxedShapeIsShape) {
+ // %0 = fir.embox %arg(%shape) [%slice]
+ // %1 = fir.array_coor %0 %idx
+ // Pull in as:
+ // %1 = fir.array_coor %arg(%shape) %[slice] %idx
+ } else if (boxedShapeIsShapeShift) {
+ // %0 = fir.embox %arg(%shapeshift) [%slice]
+ // %1 = fir.array_coor %0 %idx
+ // Pull in as:
+ // %tmp = arith.addi %idx, %shapeshift.lb
+ // %idx_shifted = arith.subi %tmp, 1
+ // %1 = fir.array_coor %arg(%shapeshift) %[slice] %idx_shifted
+ shiftedIndices =
+ getShiftedIndices(boxedShape, op.getIndices(), rewriter);
+ if (!shiftedIndices)
+ return mlir::failure();
+ } else {
+ return mlir::failure();
+ }
}
}
} else { // !boxedShape
@@ -657,8 +669,8 @@ struct SimplifyArrayCoorOp : public mlir::OpRewritePattern<fir::ArrayCoorOp> {
} else {
// %0 = fir.embox %arg(%shape) [%slice]
// %1 = fir.array_coor %0(%shape) %idx
- // Cannot pull in without adjusting the array_coor indices.
- return mlir::failure();
+ // Pull in as:
+ // %1 = fir.array_coor %arg(%shape) [%slice] %idx
}
}
} else { // !boxedShape
@@ -682,8 +694,7 @@ struct SimplifyArrayCoorOp : public mlir::OpRewritePattern<fir::ArrayCoorOp> {
} else {
// %0 = fir.rebox %arg [%slice]
// %1 = fir.array_coor %0(%shape) %idx
- // Cannot pull in without adjusting the slice and array_coor
- // indices.
+ // Cannot pull in without adjusting the slice indices.
return mlir::failure();
}
}
@@ -709,9 +720,72 @@ struct SimplifyArrayCoorOp : public mlir::OpRewritePattern<fir::ArrayCoorOp> {
op.getShapeMutable().assign(boxedShape);
if (boxedSlice)
op.getSliceMutable().assign(boxedSlice);
+ if (shiftedIndices)
+ op.getIndicesMutable().assign(*shiftedIndices);
});
return mlir::success();
}
+
+private:
+ using IndicesVectorTy = std::vector<mlir::Value>;
+
+ // If v is a shape_shift operation:
+ // fir.shape_shift %l1, %e1, %l2, %e2, ...
+ // create:
+ // fir.shape %e1, %e2, ...
+ static mlir::Value getShapeFromShapeShift(mlir::Value v,
+ mlir::PatternRewriter &rewriter) {
+ auto shapeShiftOp =
+ mlir::dyn_cast_or_null<fir::ShapeShiftOp>(v.getDefiningOp());
+ if (!shapeShiftOp)
+ return nullptr;
+ mlir::OpBuilder::InsertionGuard guard(rewriter);
+ rewriter.setInsertionPoint(shapeShiftOp);
+ return rewriter.create<fir::ShapeOp>(shapeShiftOp.getLoc(),
+ shapeShiftOp.getExtents());
+ }
+
+ static std::optional<IndicesVectorTy>
+ getShiftedIndices(mlir::Value v, mlir::ValueRange indices,
+ mlir::PatternRewriter &rewriter) {
+ auto insertAdjustments = [&](mlir::Operation *op, mlir::ValueRange lbs) {
+ // Compute the shifted indices using the extended type.
+ // Note that this can probably result in less efficient
+ // MLIR and further LLVM IR due to the extra conversions.
+ mlir::OpBuilder::InsertPoint savedIP = rewriter.saveInsertionPoint();
+ rewriter.setInsertionPoint(op);
+ mlir::Location loc = op->getLoc();
+ mlir::Type idxTy = rewriter.getIndexType();
+ mlir::Value one = rewriter.create<mlir::arith::ConstantOp>(
+ loc, idxTy, rewriter.getIndexAttr(1));
+ rewriter.restoreInsertionPoint(savedIP);
+ auto nsw = mlir::arith::IntegerOverflowFlags::nsw;
+
+ IndicesVectorTy shiftedIndices;
+ for (auto [lb, idx] : llvm::zip(lbs, indices)) {
+ mlir::Value extLb = rewriter.create<fir::ConvertOp>(loc, idxTy, lb);
+ mlir::Value extIdx = rewriter.create<fir::ConvertOp>(loc, idxTy, idx);
+ mlir::Value add =
+ rewriter.create<mlir::arith::AddIOp>(loc, extIdx, extLb, nsw);
+ mlir::Value sub =
+ rewriter.create<mlir::arith::SubIOp>(loc, add, one, nsw);
+ shiftedIndices.push_back(sub);
+ }
+
+ return std::move(shiftedIndices);
+ };
+
+ if (auto shiftOp =
+ mlir::dyn_cast_or_null<fir::ShiftOp>(v.getDefiningOp())) {
+ return insertAdjustments(shiftOp.getOperation(), shiftOp.getOrigins());
+ } else if (auto shapeShiftOp = mlir::dyn_cast_or_null<fir::ShapeShiftOp>(
+ v.getDefiningOp())) {
+ return insertAdjustments(shapeShiftOp.getOperation(),
+ shapeShiftOp.getOrigins());
+ }
+
+ return std::nullopt;
+ }
};
void fir::ArrayCoorOp::getCanonicalizationPatterns(
diff --git a/flang/test/Fir/array-coor-canonicalization.fir b/flang/test/Fir/array-coor-canonicalization.fir
index 63c3f98672557..bc58347bdf392 100644
--- a/flang/test/Fir/array-coor-canonicalization.fir
+++ b/flang/test/Fir/array-coor-canonicalization.fir
@@ -292,24 +292,74 @@ func.func @test7_3(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
return %3 : !fir.ref<i32>
}
-// CHECK-LABEL: func.func @test8(
-// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+// CHECK-LABEL: func.func @test8_1(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.box<!fir.array<?xi32>>,
+// CHECK-SAME: %[[VAL_1:.*]]: index,
+// CHECK-SAME: %[[VAL_2:.*]]: index) -> !fir.ref<i32> {
+// CHECK: %[[VAL_3:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_4:.*]] = arith.constant 10 : index
+// CHECK: %[[VAL_5:.*]] = arith.constant 19 : index
+// CHECK: %[[VAL_6:.*]] = fir.shift %[[VAL_1]] : (index) -> !fir.shift<1>
+// CHECK: %[[VAL_7:.*]] = fir.slice %[[VAL_4]], %[[VAL_5]], %[[VAL_3]] : (index, index, index) -> !fir.slice<1>
+// CHECK: %[[VAL_8:.*]] = arith.addi %[[VAL_2]], %[[VAL_1]] overflow<nsw> : index
+// CHECK: %[[VAL_9:.*]] = arith.subi %[[VAL_8]], %[[VAL_3]] overflow<nsw> : index
+// CHECK: %[[VAL_10:.*]] = fir.array_coor %[[VAL_0]](%[[VAL_6]]) {{\[}}%[[VAL_7]]] %[[VAL_9]] : (!fir.box<!fir.array<?xi32>>, !fir.shift<1>, !fir.slice<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_10]] : !fir.ref<i32>
+// CHECK: }
+func.func @test8_1(%arg0: !fir.box<!fir.array<?xi32>>, %lb: index, %idx: index) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c19 = arith.constant 19 : index
+ %sh = fir.shift %lb : (index) -> !fir.shift<1>
+ %s = fir.slice %c10, %c19, %c1 : (index, index, index) -> !fir.slice<1>
+ %2 = fir.rebox %arg0(%sh) [%s] : (!fir.box<!fir.array<?xi32>>, !fir.shift<1>, !fir.slice<1>) -> !fir.box<!fir.array<?xi32>>
+ %3 = fir.array_coor %2 %idx : (!fir.box<!fir.array<?xi32>>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test8_2(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
// CHECK: %[[VAL_2:.*]] = arith.constant 10 : index
-// CHECK: %[[VAL_3:.*]] = arith.constant 19 : index
-// CHECK: %[[VAL_4:.*]] = fir.shape_shift %[[VAL_2]], %[[VAL_2]] : (index, index) -> !fir.shapeshift<1>
-// CHECK: %[[VAL_5:.*]] = fir.slice %[[VAL_2]], %[[VAL_3]], %[[VAL_1]] : (index, index, index) -> !fir.slice<1>
-// CHECK: %[[VAL_6:.*]] = fir.array_coor %[[VAL_0]](%[[VAL_4]]) {{\[}}%[[VAL_5]]] %[[VAL_1]] : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>, !fir.slice<1>, index) -> !fir.ref<i32>
+// CHECK: %[[VAL_3:.*]] = arith.constant 9 : index
+// CHECK: %[[VAL_4:.*]] = fir.shape %[[VAL_2]] : (index) -> !fir.shape<1>
+// CHECK: %[[VAL_5:.*]] = fir.slice %[[VAL_1]], %[[VAL_3]], %[[VAL_1]] : (index, index, index) -> !fir.slice<1>
+// CHECK: %[[VAL_6:.*]] = fir.array_coor %[[VAL_0]](%[[VAL_4]]) {{\[}}%[[VAL_5]]] %[[VAL_1]] : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>, !fir.slice<1>, index) -> !fir.ref<i32>
// CHECK: return %[[VAL_6]] : !fir.ref<i32>
// CHECK: }
-func.func @test8(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+func.func @test8_2(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c9 = arith.constant 9 : index
+ %1 = fir.shape %c10 : (index) -> !fir.shape<1>
+ %s = fir.slice %c1, %c9, %c1 : (index, index, index) -> !fir.slice<1>
+ %2 = fir.embox %arg0(%1) [%s] : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>, !fir.slice<1>) -> !fir.box<!fir.array<10xi32>>
+ %3 = fir.array_coor %2 %c1 : (!fir.box<!fir.array<10xi32>>, index) -> !fir.ref<i32>
+ return %3 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: func.func @test8_3(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<10xi32>>,
+// CHECK-SAME: %[[VAL_1:.*]]: index,
+// CHECK-SAME: %[[VAL_2:.*]]: index) -> !fir.ref<i32> {
+// CHECK: %[[VAL_3:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_4:.*]] = arith.constant 10 : index
+// CHECK: %[[VAL_5:.*]] = arith.constant 19 : index
+// CHECK: %[[VAL_6:.*]] = fir.shape_shift %[[VAL_1]], %[[VAL_4]] : (index, index) -> !fir.shapeshift<1>
+// CHECK: %[[VAL_7:.*]] = fir.slice %[[VAL_4]], %[[VAL_5]], %[[VAL_3]] : (index, index, index) -> !fir.slice<1>
+// CHECK: %[[VAL_8:.*]] = arith.addi %[[VAL_2]], %[[VAL_1]] overflow<nsw> : index
+// CHECK: %[[VAL_9:.*]] = arith.subi %[[VAL_8]], %[[VAL_3]] overflow<nsw> : index
+// CHECK: %[[VAL_10:.*]] = fir.array_coor %[[VAL_0]](%[[VAL_6]]) {{\[}}%[[VAL_7]]] %[[VAL_9]] : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>, !fir.slice<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_10]] : !fir.ref<i32>
+// CHECK: }
+func.func @test8_3(%arg0: !fir.ref<!fir.array<10xi32>>, %lb: index, %idx: index) -> !fir.ref<i32> {
%c1 = arith.constant 1 : index
%c10 = arith.constant 10 : index
%c19 = arith.constant 19 : index
- %1 = fir.shape_shift %c10, %c10 : (index, index) -> !fir.shapeshift<1>
+ %1 = fir.shape_shift %lb, %c10 : (index, index) -> !fir.shapeshift<1>
%s = fir.slice %c10, %c19, %c1 : (index, index, index) -> !fir.slice<1>
%2 = fir.embox %arg0(%1) [%s] : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>, !fir.slice<1>) -> !fir.box<!fir.array<10xi32>>
- %3 = fir.array_coor %2 %c1 : (!fir.box<!fir.array<10xi32>>, index) -> !fir.ref<i32>
+ %3 = fir.array_coor %2 %idx : (!fir.box<!fir.array<10xi32>>, index) -> !fir.ref<i32>
return %3 : !fir.ref<i32>
}
@@ -444,9 +494,8 @@ func.func @test15(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
// CHECK: %[[VAL_3:.*]] = arith.constant 19 : index
// CHECK: %[[VAL_4:.*]] = fir.shape_shift %[[VAL_2]], %[[VAL_2]] : (index, index) -> !fir.shapeshift<1>
// CHECK: %[[VAL_5:.*]] = fir.slice %[[VAL_2]], %[[VAL_3]], %[[VAL_1]] : (index, index, index) -> !fir.slice<1>
-// CHECK: %[[VAL_6:.*]] = fir.embox %[[VAL_0]](%[[VAL_4]]) {{\[}}%[[VAL_5]]] : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>, !fir.slice<1>) -> !fir.box<!fir.array<10xi32>>
-// CHECK: %[[VAL_7:.*]] = fir.array_coor %[[VAL_6]](%[[VAL_4]]) %[[VAL_2]] : (!fir.box<!fir.array<10xi32>>, !fir.shapeshift<1>, index) -> !fir.ref<i32>
-// CHECK: return %[[VAL_7]] : !fir.ref<i32>
+// CHECK: %[[VAL_6:.*]] = fir.array_coor %[[VAL_0]](%[[VAL_4]]) {{\[}}%[[VAL_5]]] %[[VAL_2]] : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>, !fir.slice<1>, index) -> !fir.ref<i32>
+// CHECK: return %[[VAL_6]] : !fir.ref<i32>
// CHECK: }
func.func @test16(%arg0: !fir.ref<!fir.array<10xi32>>) -> !fir.ref<i32> {
%c1 = arith.constant 1 : index
diff --git a/flang/test/Fir/convert-to-llvm.fir b/flang/test/Fir/convert-to-llvm.fir
index 80aebd5f8500d..21323a5e657c9 100644
--- a/flang/test/Fir/convert-to-llvm.fir
+++ b/flang/test/Fir/convert-to-llvm.fir
@@ -2084,7 +2084,7 @@ func.func @ext_array_coor1(%arg0: !fir.ref<!fir.array<?xi32>>) {
// CHECK: %[[C0:.*]] = llvm.mlir.constant(0 : i64) : i64
// CHECK: %[[C1:.*]] = llvm.mlir.constant(1 : i64) : i64
// CHECK: %[[C0_1:.*]] = llvm.mlir.constant(0 : i64) : i64
-// CHECK: %[[IDX:.*]] = llvm.sub %[[C0]], %[[C1]] overflow<nsw> : i64
+// CHECK: %[[IDX:.*]] = llvm.sub %[[C0]], %[[C0]] overflow<nsw> : i64
// CHECK: %[[DIFF0:.*]] = llvm.mul %[[IDX]], %[[C0]] overflow<nsw> : i64
// CHECK: %[[ADJ:.*]] = llvm.sub %[[C0]], %[[C0]] overflow<nsw> : i64
// CHECK: %[[DIFF1:.*]] = llvm.add %[[DIFF0]], %[[ADJ]] overflow<nsw> : i64
@@ -2153,7 +2153,7 @@ func.func @ext_array_coor4(%arg0: !fir.ref<!fir.array<100xi32>>) {
// CHECK: %[[C1:.*]] = llvm.mlir.constant(1 : i64) : i64
// CHECK: %[[C1_1:.*]] = llvm.mlir.constant(1 : i64) : i64
// CHECK: %[[C0_1:.*]] = llvm.mlir.constant(0 : i64) : i64
-// CHECK: %[[IDX:.*]] = llvm.sub %[[C1]], %[[C1_1]] overflow<nsw> : i64
+// CHECK: %[[IDX:.*]] = llvm.sub %[[C1]], %[[C0]] overflow<nsw> : i64
// CHECK: %[[DIFF0:.*]] = llvm.mul %[[IDX]], %[[C1]] overflow<nsw> : i64
// CHECK: %[[ADJ:.*]] = llvm.sub %[[C10]], %[[C0]] overflow<nsw> : i64
// CHECK: %[[DIFF1:.*]] = llvm.add %[[DIFF0]], %[[ADJ]] overflow<nsw> : i64
More information about the flang-commits
mailing list