Skip to content

Commit 0b80491

Browse files
authored
[flang] Support non-index shape/shift/slice for CG box operations. (llvm#124625)
That is another problem uncovered during hlfir.reshape inlining, where the shape bits could be any integer type. This patch adds explicit convertions to `index` type where needed.
1 parent 1b729c3 commit 0b80491

File tree

2 files changed

+92
-22
lines changed

2 files changed

+92
-22
lines changed

flang/lib/Optimizer/CodeGen/CodeGen.cpp

Lines changed: 31 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1675,22 +1675,26 @@ struct EmboxCommonConversion : public fir::FIROpConversion<OP> {
16751675
this->attachTBAATag(storeOp, boxTy, boxTy, nullptr);
16761676
return storage;
16771677
}
1678-
};
16791678

1680-
/// Compute the extent of a triplet slice (lb:ub:step).
1681-
static mlir::Value
1682-
computeTripletExtent(mlir::ConversionPatternRewriter &rewriter,
1683-
mlir::Location loc, mlir::Value lb, mlir::Value ub,
1684-
mlir::Value step, mlir::Value zero, mlir::Type type) {
1685-
mlir::Value extent = rewriter.create<mlir::LLVM::SubOp>(loc, type, ub, lb);
1686-
extent = rewriter.create<mlir::LLVM::AddOp>(loc, type, extent, step);
1687-
extent = rewriter.create<mlir::LLVM::SDivOp>(loc, type, extent, step);
1688-
// If the resulting extent is negative (`ub-lb` and `step` have different
1689-
// signs), zero must be returned instead.
1690-
auto cmp = rewriter.create<mlir::LLVM::ICmpOp>(
1691-
loc, mlir::LLVM::ICmpPredicate::sgt, extent, zero);
1692-
return rewriter.create<mlir::LLVM::SelectOp>(loc, cmp, extent, zero);
1693-
}
1679+
/// Compute the extent of a triplet slice (lb:ub:step).
1680+
mlir::Value computeTripletExtent(mlir::ConversionPatternRewriter &rewriter,
1681+
mlir::Location loc, mlir::Value lb,
1682+
mlir::Value ub, mlir::Value step,
1683+
mlir::Value zero, mlir::Type type) const {
1684+
lb = this->integerCast(loc, rewriter, type, lb);
1685+
ub = this->integerCast(loc, rewriter, type, ub);
1686+
step = this->integerCast(loc, rewriter, type, step);
1687+
zero = this->integerCast(loc, rewriter, type, zero);
1688+
mlir::Value extent = rewriter.create<mlir::LLVM::SubOp>(loc, type, ub, lb);
1689+
extent = rewriter.create<mlir::LLVM::AddOp>(loc, type, extent, step);
1690+
extent = rewriter.create<mlir::LLVM::SDivOp>(loc, type, extent, step);
1691+
// If the resulting extent is negative (`ub-lb` and `step` have different
1692+
// signs), zero must be returned instead.
1693+
auto cmp = rewriter.create<mlir::LLVM::ICmpOp>(
1694+
loc, mlir::LLVM::ICmpPredicate::sgt, extent, zero);
1695+
return rewriter.create<mlir::LLVM::SelectOp>(loc, cmp, extent, zero);
1696+
}
1697+
};
16941698

16951699
/// Create a generic box on a memory reference. This conversions lowers the
16961700
/// abstract box to the appropriate, initialized descriptor.
@@ -1851,14 +1855,16 @@ struct XEmboxOpConversion : public EmboxCommonConversion<fir::cg::XEmboxOp> {
18511855
// translating everything to values in the descriptor wherever the entity
18521856
// has a dynamic array dimension.
18531857
for (unsigned di = 0, descIdx = 0; di < rank; ++di) {
1854-
mlir::Value extent = operands[shapeOffset];
1858+
mlir::Value extent =
1859+
integerCast(loc, rewriter, i64Ty, operands[shapeOffset]);
18551860
mlir::Value outerExtent = extent;
18561861
bool skipNext = false;
18571862
if (hasSlice) {
1858-
mlir::Value off = operands[sliceOffset];
1863+
mlir::Value off =
1864+
integerCast(loc, rewriter, i64Ty, operands[sliceOffset]);
18591865
mlir::Value adj = one;
18601866
if (hasShift)
1861-
adj = operands[shiftOffset];
1867+
adj = integerCast(loc, rewriter, i64Ty, operands[shiftOffset]);
18621868
auto ao = rewriter.create<mlir::LLVM::SubOp>(loc, i64Ty, off, adj);
18631869
if (constRows > 0) {
18641870
cstInteriorIndices.push_back(ao);
@@ -1895,7 +1901,7 @@ struct XEmboxOpConversion : public EmboxCommonConversion<fir::cg::XEmboxOp> {
18951901
// the lower bound.
18961902
if (hasShift && !(hasSlice || hasSubcomp || hasSubstr) &&
18971903
(isaPointerOrAllocatable || !normalizedLowerBound(xbox))) {
1898-
lb = operands[shiftOffset];
1904+
lb = integerCast(loc, rewriter, i64Ty, operands[shiftOffset]);
18991905
auto extentIsEmpty = rewriter.create<mlir::LLVM::ICmpOp>(
19001906
loc, mlir::LLVM::ICmpPredicate::eq, extent, zero);
19011907
lb = rewriter.create<mlir::LLVM::SelectOp>(loc, extentIsEmpty, one,
@@ -1907,9 +1913,12 @@ struct XEmboxOpConversion : public EmboxCommonConversion<fir::cg::XEmboxOp> {
19071913

19081914
// store step (scaled by shaped extent)
19091915
mlir::Value step = prevDimByteStride;
1910-
if (hasSlice)
1911-
step = rewriter.create<mlir::LLVM::MulOp>(loc, i64Ty, step,
1912-
operands[sliceOffset + 2]);
1916+
if (hasSlice) {
1917+
mlir::Value sliceStep =
1918+
integerCast(loc, rewriter, i64Ty, operands[sliceOffset + 2]);
1919+
step =
1920+
rewriter.create<mlir::LLVM::MulOp>(loc, i64Ty, step, sliceStep);
1921+
}
19131922
dest = insertStride(rewriter, loc, dest, descIdx, step);
19141923
++descIdx;
19151924
}

flang/test/Fir/convert-to-llvm.fir

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1909,6 +1909,67 @@ func.func @xembox0(%arg0: !fir.ref<!fir.array<?xi32>>) {
19091909
// CHECK: %[[BOX10:.*]] = llvm.insertvalue %[[BASE_PTR]], %[[BOX9]][0] : !llvm.struct<(ptr, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<1 x array<3 x i64>>)>
19101910
// CHECK: llvm.store %[[BOX10]], %[[ALLOCA]] : !llvm.struct<(ptr, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<1 x array<3 x i64>>)>, !llvm.ptr
19111911

1912+
// Test i32 shape/shift/slice:
1913+
func.func @xembox0_i32(%arg0: !fir.ref<!fir.array<?xi32>>) {
1914+
%c0 = arith.constant 0 : i32
1915+
%c0_i64 = arith.constant 0 : i64
1916+
%0 = fircg.ext_embox %arg0(%c0) origin %c0[%c0, %c0, %c0] : (!fir.ref<!fir.array<?xi32>>, i32, i32, i32, i32, i32) -> !fir.box<!fir.array<?xi32>>
1917+
return
1918+
}
1919+
// CHECK-LABEL: llvm.func @xembox0_i32(
1920+
// CHECK-SAME: %[[ARG0:.*]]: !llvm.ptr
1921+
// CHECK: %[[ALLOCA_SIZE:.*]] = llvm.mlir.constant(1 : i32) : i32
1922+
// GENERIC: %[[ALLOCA:.*]] = llvm.alloca %[[ALLOCA_SIZE]] x !llvm.struct<(ptr, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<1 x array<3 x i64>>)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
1923+
// AMDGPU: %[[AA:.*]] = llvm.alloca %[[ALLOCA_SIZE]] x !llvm.struct<(ptr, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<1 x array<3 x i64>>)> {alignment = 8 : i64} : (i32) -> !llvm.ptr<5>
1924+
// AMDGPU: %[[ALLOCA:.*]] = llvm.addrspacecast %[[AA]] : !llvm.ptr<5> to !llvm.ptr
1925+
// CHECK: %[[C0_I32:.*]] = llvm.mlir.constant(0 : i32) : i32
1926+
// CHECK: %[[C0:.*]] = llvm.mlir.constant(0 : i64) : i64
1927+
// CHECK: %[[TYPE:.*]] = llvm.mlir.constant(9 : i32) : i32
1928+
// CHECK: %[[NULL:.*]] = llvm.mlir.zero : !llvm.ptr
1929+
// CHECK: %[[GEP:.*]] = llvm.getelementptr %[[NULL]][1]
1930+
// CHECK: %[[ELEM_LEN_I64:.*]] = llvm.ptrtoint %[[GEP]] : !llvm.ptr to i64
1931+
// CHECK: %[[BOX0:.*]] = llvm.mlir.undef : !llvm.struct<(ptr, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<1 x array<3 x i64>>)>
1932+
// CHECK: %[[BOX1:.*]] = llvm.insertvalue %[[ELEM_LEN_I64]], %[[BOX0]][1] : !llvm.struct<(ptr, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<1 x array<3 x i64>>)>
1933+
// CHECK: %[[VERSION:.*]] = llvm.mlir.constant(20240719 : i32) : i32
1934+
// CHECK: %[[BOX2:.*]] = llvm.insertvalue %[[VERSION]], %[[BOX1]][2] : !llvm.struct<(ptr, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<1 x array<3 x i64>>)>
1935+
// CHECK: %[[RANK:.*]] = llvm.mlir.constant(1 : i32) : i32
1936+
// CHECK: %[[RANK_I8:.*]] = llvm.trunc %[[RANK]] : i32 to i8
1937+
// CHECK: %[[BOX3:.*]] = llvm.insertvalue %[[RANK_I8]], %[[BOX2]][3] : !llvm.struct<(ptr, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<1 x array<3 x i64>>)>
1938+
// CHECK: %[[TYPE_I8:.*]] = llvm.trunc %[[TYPE]] : i32 to i8
1939+
// CHECK: %[[BOX4:.*]] = llvm.insertvalue %[[TYPE_I8]], %[[BOX3]][4] : !llvm.struct<(ptr, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<1 x array<3 x i64>>)>
1940+
// CHECK: %[[ATTR:.*]] = llvm.mlir.constant(0 : i32) : i32
1941+
// CHECK: %[[ATTR_I8:.*]] = llvm.trunc %[[ATTR]] : i32 to i8
1942+
// CHECK: %[[BOX5:.*]] = llvm.insertvalue %[[ATTR_I8]], %[[BOX4]][5] : !llvm.struct<(ptr, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<1 x array<3 x i64>>)>
1943+
// CHECK: %[[F18ADDENDUM:.*]] = llvm.mlir.constant(0 : i32) : i32
1944+
// CHECK: %[[F18ADDENDUM_I8:.*]] = llvm.trunc %[[F18ADDENDUM]] : i32 to i8
1945+
// CHECK: %[[BOX6:.*]] = llvm.insertvalue %[[F18ADDENDUM_I8]], %[[BOX5]][6] : !llvm.struct<(ptr, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<1 x array<3 x i64>>)>
1946+
// CHECK: %[[ZERO:.*]] = llvm.mlir.constant(0 : i64) : i64
1947+
// CHECK: %[[ONE:.*]] = llvm.mlir.constant(1 : i64) : i64
1948+
// CHECK: %[[C0_1:.*]] = llvm.sext %[[C0_I32]] : i32 to i64
1949+
// CHECK: %[[C0_2:.*]] = llvm.sext %[[C0_I32]] : i32 to i64
1950+
// CHECK: %[[C0_3:.*]] = llvm.sext %[[C0_I32]] : i32 to i64
1951+
// CHECK: %[[ADJUSTED_OFFSET:.*]] = llvm.sub %[[C0_2]], %[[C0_3]] : i64
1952+
// CHECK: %[[DIM_OFFSET:.*]] = llvm.mul %[[ADJUSTED_OFFSET]], %[[ONE]] : i64
1953+
// CHECK: %[[PTR_OFFSET:.*]] = llvm.add %[[DIM_OFFSET]], %[[ZERO]] : i64
1954+
// CHECK: %[[C0_4:.*]] = llvm.sext %[[C0_I32]] : i32 to i64
1955+
// CHECK: %[[C0_5:.*]] = llvm.sext %[[C0_I32]] : i32 to i64
1956+
// CHECK: %[[C0_6:.*]] = llvm.sext %[[C0_I32]] : i32 to i64
1957+
// CHECK: %[[EXTENT0:.*]] = llvm.sub %[[C0_5]], %[[C0_4]] : i64
1958+
// CHECK: %[[EXTENT1:.*]] = llvm.add %[[EXTENT0]], %[[C0_6]] : i64
1959+
// CHECK: %[[EXTENT2:.*]] = llvm.sdiv %[[EXTENT1]], %[[C0_6]] : i64
1960+
// CHECK: %[[EXTENT_CMP:.*]] = llvm.icmp "sgt" %[[EXTENT2]], %[[ZERO]] : i64
1961+
// CHECK: %[[EXTENT:.*]] = llvm.select %[[EXTENT_CMP]], %[[EXTENT2]], %[[ZERO]] : i1, i64
1962+
// CHECK: %[[BOX7:.*]] = llvm.insertvalue %[[ONE]], %[[BOX6]][7, 0, 0] : !llvm.struct<(ptr, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<1 x array<3 x i64>>)>
1963+
// CHECK: %[[BOX8:.*]] = llvm.insertvalue %[[EXTENT]], %[[BOX7]][7, 0, 1] : !llvm.struct<(ptr, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<1 x array<3 x i64>>)>
1964+
// CHECK: %[[C0_7:.*]] = llvm.sext %[[C0_I32]] : i32 to i64
1965+
// CHECK: %[[STRIDE:.*]] = llvm.mul %[[ELEM_LEN_I64]], %[[C0_7]] : i64
1966+
// CHECK: %[[BOX9:.*]] = llvm.insertvalue %[[STRIDE]], %[[BOX8]][7, 0, 2] : !llvm.struct<(ptr, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<1 x array<3 x i64>>)>
1967+
// CHECK: %[[PREV_DIM:.*]] = llvm.mul %[[ELEM_LEN_I64]], %[[C0_1]] : i64
1968+
// CHECK: %[[PREV_PTROFF:.*]] = llvm.mul %[[ONE]], %[[C0_1]] : i64
1969+
// CHECK: %[[BASE_PTR:.*]] = llvm.getelementptr %[[ARG0]][%[[PTR_OFFSET]]] : (!llvm.ptr, i64) -> !llvm.ptr, i32
1970+
// CHECK: %[[BOX10:.*]] = llvm.insertvalue %[[BASE_PTR]], %[[BOX9]][0] : !llvm.struct<(ptr, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<1 x array<3 x i64>>)>
1971+
// CHECK: llvm.store %[[BOX10]], %[[ALLOCA]] : !llvm.struct<(ptr, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<1 x array<3 x i64>>)>, !llvm.ptr
1972+
19121973
// Check adjustment of element scaling factor.
19131974

19141975
func.func @xembox1(%arg0: !fir.ref<!fir.array<?x!fir.char<1, 10>>>) {

0 commit comments

Comments
 (0)