[flang] Support non-index shape/shift/slice for CG box operations. (#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.
This commit is contained in:
Slava Zakharin
2025-01-28 09:38:33 -08:00
committed by GitHub
parent 1b729c3d70
commit 0b80491cd5
2 changed files with 92 additions and 22 deletions

View File

@@ -1675,22 +1675,26 @@ struct EmboxCommonConversion : public fir::FIROpConversion<OP> {
this->attachTBAATag(storeOp, boxTy, boxTy, nullptr);
return storage;
}
};
/// Compute the extent of a triplet slice (lb:ub:step).
static mlir::Value
computeTripletExtent(mlir::ConversionPatternRewriter &rewriter,
mlir::Location loc, mlir::Value lb, mlir::Value ub,
mlir::Value step, mlir::Value zero, mlir::Type type) {
mlir::Value extent = rewriter.create<mlir::LLVM::SubOp>(loc, type, ub, lb);
extent = rewriter.create<mlir::LLVM::AddOp>(loc, type, extent, step);
extent = rewriter.create<mlir::LLVM::SDivOp>(loc, type, extent, step);
// If the resulting extent is negative (`ub-lb` and `step` have different
// signs), zero must be returned instead.
auto cmp = rewriter.create<mlir::LLVM::ICmpOp>(
loc, mlir::LLVM::ICmpPredicate::sgt, extent, zero);
return rewriter.create<mlir::LLVM::SelectOp>(loc, cmp, extent, zero);
}
/// Compute the extent of a triplet slice (lb:ub:step).
mlir::Value computeTripletExtent(mlir::ConversionPatternRewriter &rewriter,
mlir::Location loc, mlir::Value lb,
mlir::Value ub, mlir::Value step,
mlir::Value zero, mlir::Type type) const {
lb = this->integerCast(loc, rewriter, type, lb);
ub = this->integerCast(loc, rewriter, type, ub);
step = this->integerCast(loc, rewriter, type, step);
zero = this->integerCast(loc, rewriter, type, zero);
mlir::Value extent = rewriter.create<mlir::LLVM::SubOp>(loc, type, ub, lb);
extent = rewriter.create<mlir::LLVM::AddOp>(loc, type, extent, step);
extent = rewriter.create<mlir::LLVM::SDivOp>(loc, type, extent, step);
// If the resulting extent is negative (`ub-lb` and `step` have different
// signs), zero must be returned instead.
auto cmp = rewriter.create<mlir::LLVM::ICmpOp>(
loc, mlir::LLVM::ICmpPredicate::sgt, extent, zero);
return rewriter.create<mlir::LLVM::SelectOp>(loc, cmp, extent, zero);
}
};
/// Create a generic box on a memory reference. This conversions lowers the
/// abstract box to the appropriate, initialized descriptor.
@@ -1851,14 +1855,16 @@ struct XEmboxOpConversion : public EmboxCommonConversion<fir::cg::XEmboxOp> {
// translating everything to values in the descriptor wherever the entity
// has a dynamic array dimension.
for (unsigned di = 0, descIdx = 0; di < rank; ++di) {
mlir::Value extent = operands[shapeOffset];
mlir::Value extent =
integerCast(loc, rewriter, i64Ty, operands[shapeOffset]);
mlir::Value outerExtent = extent;
bool skipNext = false;
if (hasSlice) {
mlir::Value off = operands[sliceOffset];
mlir::Value off =
integerCast(loc, rewriter, i64Ty, operands[sliceOffset]);
mlir::Value adj = one;
if (hasShift)
adj = operands[shiftOffset];
adj = integerCast(loc, rewriter, i64Ty, operands[shiftOffset]);
auto ao = rewriter.create<mlir::LLVM::SubOp>(loc, i64Ty, off, adj);
if (constRows > 0) {
cstInteriorIndices.push_back(ao);
@@ -1895,7 +1901,7 @@ struct XEmboxOpConversion : public EmboxCommonConversion<fir::cg::XEmboxOp> {
// the lower bound.
if (hasShift && !(hasSlice || hasSubcomp || hasSubstr) &&
(isaPointerOrAllocatable || !normalizedLowerBound(xbox))) {
lb = operands[shiftOffset];
lb = integerCast(loc, rewriter, i64Ty, operands[shiftOffset]);
auto extentIsEmpty = rewriter.create<mlir::LLVM::ICmpOp>(
loc, mlir::LLVM::ICmpPredicate::eq, extent, zero);
lb = rewriter.create<mlir::LLVM::SelectOp>(loc, extentIsEmpty, one,
@@ -1907,9 +1913,12 @@ struct XEmboxOpConversion : public EmboxCommonConversion<fir::cg::XEmboxOp> {
// store step (scaled by shaped extent)
mlir::Value step = prevDimByteStride;
if (hasSlice)
step = rewriter.create<mlir::LLVM::MulOp>(loc, i64Ty, step,
operands[sliceOffset + 2]);
if (hasSlice) {
mlir::Value sliceStep =
integerCast(loc, rewriter, i64Ty, operands[sliceOffset + 2]);
step =
rewriter.create<mlir::LLVM::MulOp>(loc, i64Ty, step, sliceStep);
}
dest = insertStride(rewriter, loc, dest, descIdx, step);
++descIdx;
}