Skip to content

[TritonToLinalg](fix) lower zero-stride block pointer loads #1508

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
zhanwei33 wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

[TritonToLinalg](fix) lower zero-stride block pointer loads #1508

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
49e20c2
Fix zero-stride block pointer load boundary check
zhanwei33 Aug 11, 2026
869b4d6
Lower zero-stride block loads through fallback
zhanwei33 Aug 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 68 additions & 8 deletions third_party/ascend/lib/TritonToLinalg/LoadStoreConverter.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,56 @@ using namespace triton;
const std::string MayImplicitTransposeWithLastAxisTAG =
"MayImplicitTransposeWithLastAxis";

namespace {

Value getRemappedOrOriginal(Value value, ConversionPatternRewriter &rewriter) {
if (Value remapped = rewriter.getRemappedValue(value))
return remapped;
return value;
}

bool hasStaticZeroStride(triton::MakeTensorPtrOp makeTensorPtrOp) {
return llvm::any_of(makeTensorPtrOp.getStrides(), [](Value stride) {
auto constantStride = getConstantIntValue(stride);
return constantStride.has_value() && constantStride.value() == 0;
});
}

SmallVector<OpFoldResult> getBoundarySizesFromMakeTensorPtr(
triton::MakeTensorPtrOp makeTensorPtrOp,
llvm::ArrayRef<int32_t> boundaryCheck, llvm::ArrayRef<int64_t> tileShape,
const Location &loc, ConversionPatternRewriter &rewriter) {
assert(makeTensorPtrOp.getShape().size() == tileShape.size());
assert(makeTensorPtrOp.getOffsets().size() == tileShape.size());

SmallVector<OpFoldResult> boundarySizes =
getAsIndexOpFoldResult(rewriter.getContext(), tileShape);
const OpFoldResult zero = rewriter.getIndexAttr(0);

for (size_t i = 0; i < tileShape.size(); ++i) {
if (llvm::find(boundaryCheck, i) == boundaryCheck.end())
continue;

OpFoldResult shape = getOpFoldResultOfLayoutInfo(
getRemappedOrOriginal(makeTensorPtrOp.getShape()[i], rewriter),
rewriter);
OpFoldResult offset = getOpFoldResultOfLayoutInfo(
getRemappedOrOriginal(makeTensorPtrOp.getOffsets()[i], rewriter),
rewriter);
OpFoldResult nonNegativeOffset =
maxOpFoldResult(offset, zero, loc, rewriter);
OpFoldResult remaining = maxOpFoldResult(
subOpFoldResult(shape, nonNegativeOffset, loc, rewriter), zero, loc,
rewriter);
boundarySizes[i] =
minOpFoldResult(boundarySizes[i], remaining, loc, rewriter);
}

return boundarySizes;
}

} // namespace

LogicalResult
AddPtrConverter::matchAndRewrite(triton::AddPtrOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const {
Expand Down Expand Up @@ -325,13 +375,19 @@ LogicalResult LoadConverter::replaceMaskedLoadWithTensorOther(
LogicalResult
LoadConverter::matchAndRewrite(triton::LoadOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const {
auto makeTensorPtrOp = op.getPtr().getDefiningOp<triton::MakeTensorPtrOp>();
const bool hasZeroStrideMTP =
makeTensorPtrOp && hasStaticZeroStride(makeTensorPtrOp);

// Check if tt.load is modified by AddPtrConverter to a specified state.
if (checkModifiedByAddPtrConverter(op).succeeded()) {
if (!hasZeroStrideMTP && checkModifiedByAddPtrConverter(op).succeeded()) {
return continueModifyFromAddPtrConverter(op, adaptor, rewriter);
}

auto ptr = adaptor.getPtr();
Value ptr = hasZeroStrideMTP ? rewriter.getRemappedValue(op.getPtr())
: adaptor.getPtr();
if (!ptr)
return rewriter.notifyMatchFailure(op, "missing remapped tensor pointer");
auto mask = op.getMask();
auto other = op.getOther();
auto loc = op.getLoc();
Expand Down Expand Up @@ -367,7 +423,7 @@ LoadConverter::matchAndRewrite(triton::LoadOp op, OpAdaptor adaptor,
}

int64_t lastStride = -1;
if (isa<BlockArgument>(ptr)) {
if (!hasZeroStrideMTP && isa<BlockArgument>(ptr)) {
auto u = ptr;
while (auto blkArg = dyn_cast<BlockArgument>(u)) {
if (auto forOp = dyn_cast<scf::ForOp>(blkArg.getOwner()->getParentOp())) {
Expand All @@ -392,14 +448,15 @@ LoadConverter::matchAndRewrite(triton::LoadOp op, OpAdaptor adaptor,
return rewriter.notifyMatchFailure(
op, "LoadOp expects a memref, not a memref of pointers");
}
if (!op->hasAttr(ConverterUtils::GeneratedByMakeTensorPtrTAG)) {
if (!hasZeroStrideMTP &&
!op->hasAttr(ConverterUtils::GeneratedByMakeTensorPtrTAG)) {
auto memrefOp = dyn_cast<memref::ReinterpretCastOp>(ptr.getDefiningOp());
auto ret = mlir::ConverterUtils::getLastStrideOfReinterpretCastOp(memrefOp);
if (ret.has_value())
lastStride = *ret;
}
bool mayImplicitTransposeWithLastAxis =
(existDotFlag) &&
(!hasZeroStrideMTP) && (existDotFlag) &&
(!op->hasAttr(ConverterUtils::GeneratedByMakeTensorPtrTAG)) &&
(lastStride != 1 &&
mlir::ConverterUtils::isaPermutedMemRefType(memRefType));
Expand Down Expand Up @@ -464,9 +521,12 @@ LoadConverter::matchAndRewrite(triton::LoadOp op, OpAdaptor adaptor,
// boundary check
auto boundaryCheck = op.getBoundaryCheck();
if (!boundaryCheck.empty()) {
auto makeTensorPtrOp = op.getPtr().getDefiningOp<triton::MakeTensorPtrOp>();
auto boundarySizes = mlir::ConverterUtils::getBoundarySizes(
boundaryCheck, /*remapped*/ ptr, loc, rewriter);
auto boundarySizes =
hasZeroStrideMTP
? getBoundarySizesFromMakeTensorPtr(makeTensorPtrOp, boundaryCheck,
memRefShape, loc, rewriter)
: mlir::ConverterUtils::getBoundarySizes(
boundaryCheck, /*remapped*/ ptr, loc, rewriter);
// handle the padding
auto padding = op.getPadding();
SmallVector<OpFoldResult> srcOffsets(boundarySizes.size(),
Expand Down
76 changes: 43 additions & 33 deletions third_party/ascend/lib/TritonToLinalg/StridedLoadStoreRewrite.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,13 @@ static std::optional<int64_t> getStaticConstInt(Value v) {
return std::nullopt;
}

static bool hasStaticZeroStride(ValueRange strides) {
return llvm::any_of(strides, [](Value stride) {
auto value = getStaticConstInt(stride);
return value.has_value() && value.value() == 0;
});
}

static std::optional<int64_t> getStaticMaskUpperBound(Value mask) {
if (!mask)
return std::nullopt;
Expand Down Expand Up @@ -817,40 +824,43 @@ static LogicalResult tryRewriteBlockPtrLoad(triton::LoadOp op,
ArrayRef<int64_t> shape = resultType.getShape();
int64_t rank = static_cast<int64_t>(shape.size());

// ---- order must match the "non-permuted" layout: order[i] == rank-1-i,
// i.e. innermost (fastest-changing) is the last dim of the tensor.
// ImplicitPermute handles permuted layouts via tt.trans, so we
// leave anything non-canonical alone.
auto order = mtpt.getOrder();
if (static_cast<int64_t>(order.size()) != rank)
return failure();
for (int64_t i = 0; i < rank; ++i) {
if (order[i] != rank - 1 - i) {
return failure();
}
}

// ---- stride check ----
auto strides = mtpt.getStrides();
if (strides.empty() || static_cast<int64_t>(strides.size()) != rank)
return failure();
// Stride dispatch: strided DMA on the MTE engine only supports power-of-two
// strides; a non-power-of-two stride would degrade to a slow scalar access.
// Dynamic strides stay on the structured SIMD path because they may be
// runtime stride 1 or power-of-two, where SIMT stride is slower. So we
// only rewrite to SIMT stride for *static non-power-of-two* strides:
// stride 1 -> contiguous; stride 2 (even dim) -> deinterleave;
// stride >= 4 (power of two) -> (compact) strided DMA.
APInt lastStrideC;
if (!matchPattern(strides.back(), m_ConstantInt(&lastStrideC)))
return failure();
int64_t lastStride = std::abs(lastStrideC.getSExtValue());
if (lastStride <= 1)
return failure();
if (lastStride == 2)
return failure(); // even -> deinterleave; odd -> strided DMA
if ((lastStride & (lastStride - 1)) == 0)
return failure(); // power-of-two >= 4 -> strided DMA
const bool zeroStrideBroadcast = hasStaticZeroStride(strides);
int64_t lastStride = 0;
if (!zeroStrideBroadcast) {
Comment on lines +831 to +833

@github-actions github-actions Bot Aug 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[other · low]
lastStride is declared at function scope with initial value 0, but is only meaningfully assigned inside the if (!zeroStrideBroadcast) block (line 856). When zeroStrideBroadcast=true, the code path reaches line 1022 (LLVM_DEBUG printing last_stride), where lastStride remains 0 rather than reflecting the actual last stride value. Since hasStaticZeroStride only checks if any stride is 0, the last stride may well be non-zero, making the debug output misleading.

Suggestion: move lastStride inside the if (!zeroStrideBroadcast) block to scope it correctly, and for the debug output at line 1022 either extract the last stride on-the-fly from strides.back() or conditionally skip printing it for the broadcast case.

Suggestion:

Suggested change
const bool zeroStrideBroadcast = hasStaticZeroStride(strides);
int64_t lastStride = 0;
if (!zeroStrideBroadcast) {
const bool zeroStrideBroadcast = hasStaticZeroStride(strides);
if (!zeroStrideBroadcast) {
int64_t lastStride = 0;
// ... order check, stride check ...
APInt lastStrideC;
if (!matchPattern(strides.back(), m_ConstantInt(&lastStrideC)))
return failure();
lastStride = std::abs(lastStrideC.getSExtValue());
// ...

// ---- order must match the "non-permuted" layout: order[i] == rank-1-i,
// i.e. innermost (fastest-changing) is the last dim of the tensor.
// ImplicitPermute handles permuted layouts via tt.trans, so we
// leave anything non-canonical alone.
auto order = mtpt.getOrder();
if (static_cast<int64_t>(order.size()) != rank)
return failure();
for (int64_t i = 0; i < rank; ++i) {
if (order[i] != rank - 1 - i)
return failure();
}

// Stride dispatch: strided DMA on the MTE engine only supports power-of-two
// strides; a non-power-of-two stride would degrade to a slow scalar access.
// Dynamic strides stay on the structured SIMD path because they may be
// runtime stride 1 or power-of-two, where SIMT stride is slower. So we
// only rewrite to SIMT stride for *static non-power-of-two* strides:
// stride 1 -> contiguous; stride 2 (even dim) -> deinterleave;
// stride >= 4 (power of two) -> (compact) strided DMA.
APInt lastStrideC;
if (!matchPattern(strides.back(), m_ConstantInt(&lastStrideC)))
return failure();
lastStride = std::abs(lastStrideC.getSExtValue());
if (lastStride <= 1)
return failure();
if (lastStride == 2)
return failure(); // even -> deinterleave; odd -> strided DMA
if ((lastStride & (lastStride - 1)) == 0)
return failure(); // power-of-two >= 4 -> strided DMA
}

// ---- Compute per-axis effective base offsets: mtpt.offsets[d] +
// (advance.offsets[d] if present)
Expand Down Expand Up @@ -915,8 +925,8 @@ static LogicalResult tryRewriteBlockPtrLoad(triton::LoadOp op,
scalarBase = rewriter.create<arith::AddIOp>(loc, scalarBase, prod);
}

if (resultType.getRank() >= 1 && resultType.getRank() <= 3 &&
resultType.hasStaticShape()) {
if (!zeroStrideBroadcast && resultType.getRank() >= 1 &&
resultType.getRank() <= 3 && resultType.hasStaticShape()) {
FailureOr<SmallVector<Value>> numels = getBlockPtrStrideNumels(
op.getOperation(), op.getMask(), op.getBoundaryCheck(), mtpt,
resultType, logicalOffsets, rewriter);
Expand Down Expand Up @@ -944,7 +954,7 @@ static LogicalResult tryRewriteBlockPtrLoad(triton::LoadOp op,
}
}

if (rank <= 3)
if (!zeroStrideBroadcast && rank <= 3)
return failure();

auto i64TensorTy = RankedTensorType::get(shape, i64Ty);
Expand Down
104 changes: 98 additions & 6 deletions third_party/ascend/lib/TritonToUnstructure/UnstructureConversionPass.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,45 @@ static Type getResultElementType(RankedTensorType ptrType) {
return resultElementType;
}

static triton::MakeTensorPtrOp getBaseMakeTensorPtr(Value ptr) {
while (auto advance = ptr.getDefiningOp<triton::AdvanceOp>())
ptr = advance.getPtr();
return ptr.getDefiningOp<triton::MakeTensorPtrOp>();
}

static bool hasStaticZeroStride(triton::MakeTensorPtrOp makeTensorPtr) {
return makeTensorPtr &&
llvm::any_of(makeTensorPtr.getStrides(), [](Value stride) {
auto constantStride = getConstantIntValue(stride);
return constantStride.has_value() && constantStride.value() == 0;
});
}

static Value castToIndex(Value value, Location loc, PatternRewriter &rewriter) {
if (value.getType().isIndex())
return value;
return rewriter.create<arith::IndexCastOp>(loc, rewriter.getIndexType(),
value);
}

static Value getPaddingValue(triton::LoadOp op, Type elementType,
PatternRewriter &rewriter) {
auto padding = op.getPadding();
if (!padding.has_value())
return Value();

auto loc = op.getLoc();
TypedAttr padAttr = rewriter.getZeroAttr(elementType);
if (padding.value() == triton::PaddingOption::PAD_NAN) {
auto floatType = dyn_cast<FloatType>(elementType);
if (!floatType)
return Value();
auto nan = llvm::APFloat::getNaN(floatType.getFloatSemantics());
padAttr = rewriter.getFloatAttr(elementType, nan);
}
return rewriter.create<arith::ConstantOp>(loc, padAttr);
}

static int64_t getTypeSizeInByte(Type type) {
if (auto intType = dyn_cast<IntegerType>(type)) {
return intType.getWidth() / kBitsPerByte;
Expand Down Expand Up @@ -479,6 +518,25 @@ LogicalResult UnstructuredMemAccessConverter<MemAccOpTy>::matchAndRewrite(

auto ptrOffsetInfo = offsetMap.at(ptr);

// BishengIR does not accept a memref with a static zero stride. A
// make_tensor_ptr load nevertheless has well-defined broadcast semantics:
// the physical offset does not advance on that axis, while boundary_check
// remains in the logical shape/offset coordinate system. Force only this
// load through the existing scalar-loop fallback so it never materializes a
// zero-strided memref. Other accesses keep their original structured path.
triton::MakeTensorPtrOp zeroStrideMakeTensorPtr;
if constexpr (std::is_same_v<MemAccOpTy, triton::LoadOp>) {
zeroStrideMakeTensorPtr = getBaseMakeTensorPtr(ptr);
// A5 pure-SIMT has a dedicated earlier rewrite to tt.indirect_load. Keep
// that efficient route intact; every other target/mode takes this generic
// fallback.
if (hasStaticZeroStride(zeroStrideMakeTensorPtr) &&
!(compileOn91095Flag && forceSimtTemplateFlag))
ptrOffsetInfo.setUnstructured(ptrOffsetInfo.getRank());
else
zeroStrideMakeTensorPtr = nullptr;
}

if (checkUnstructureAnnotated(op, rewriter))
ptrOffsetInfo.setUnstructured(ptrOffsetInfo.getRank());

Expand Down Expand Up @@ -582,6 +640,14 @@ LogicalResult UnstructuredMemAccessConverter<MemAccOpTy>::matchAndRewrite(
if (isLoadLike) {
iterArg =
rewriter.create<tensor::EmptyOp>(loc, resultShape, resultElementType);
if constexpr (std::is_same_v<MemAccOpTy, triton::LoadOp>) {
if (zeroStrideMakeTensorPtr && !op.getBoundaryCheck().empty()) {
if (Value padding = getPaddingValue(op, resultElementType, rewriter)) {
iterArg = rewriter.create<linalg::FillOp>(loc, padding, iterArg)
.getResult(0);
}
}
}
}
Value newOpResult = nullptr;

Expand All @@ -606,8 +672,11 @@ LogicalResult UnstructuredMemAccessConverter<MemAccOpTy>::matchAndRewrite(
extractedShape.push_back(size);
} else {
scf::ForOp forOp;
if (auto mtptOp =
srcPtr.template getDefiningOp<triton::MakeTensorPtrOp>()) {
auto mtptOp =
zeroStrideMakeTensorPtr
? zeroStrideMakeTensorPtr
: srcPtr.template getDefiningOp<triton::MakeTensorPtrOp>();
if (mtptOp) {
auto tptShape = mtptOp.getShape()[i];
if (tptShape.getType() != rewriter.getIndexType()) {
tptShape = rewriter.create<arith::IndexCastOp>(
Expand All @@ -618,18 +687,38 @@ LogicalResult UnstructuredMemAccessConverter<MemAccOpTy>::matchAndRewrite(

Value loopLower = zeroIdx;
Value loopUpper = sizeVal;
if constexpr (std::is_same_v<MemAccOpTy, triton::LoadOp>) {
if (zeroStrideMakeTensorPtr && llvm::find(op.getBoundaryCheck(), i) !=
op.getBoundaryCheck().end()) {
Value logicalOffset =
castToIndex(ptrOffsetInfo.getOffsets()[i], loc, rewriter);
Value logicalShape =
castToIndex(zeroStrideMakeTensorPtr.getShape()[i], loc, rewriter);
Value firstValidIndex =
rewriter.create<arith::SubIOp>(loc, zeroIdx, logicalOffset);
firstValidIndex =
rewriter.create<arith::MaxSIOp>(loc, firstValidIndex, zeroIdx);
loopLower =
rewriter.create<arith::MinSIOp>(loc, firstValidIndex, sizeVal);

Value remaining =
rewriter.create<arith::SubIOp>(loc, logicalShape, logicalOffset);
remaining = rewriter.create<arith::MaxSIOp>(loc, remaining, zeroIdx);
loopUpper = rewriter.create<arith::MinSIOp>(loc, remaining, sizeVal);
}
}
if (mstate && i < mstate->dims.size() && i < mstate->offsets.size()) {
Value maskOffset =
getValueOrCreateConstantIndexOp(rewriter, loc, mstate->offsets[i]);
maskOffset = rewriter.create<arith::MaxSIOp>(loc, maskOffset, zeroIdx);
maskOffset = rewriter.create<arith::MinSIOp>(loc, maskOffset, sizeVal);
loopLower = maskOffset;
loopLower = rewriter.create<arith::MaxSIOp>(loc, loopLower, maskOffset);

Value maskDim =
getValueOrCreateConstantIndexOp(rewriter, loc, mstate->dims[i]);
maskDim = rewriter.create<arith::AddIOp>(loc, maskOffset, maskDim);
maskDim = rewriter.create<arith::MinSIOp>(loc, maskDim, sizeVal);
loopUpper = maskDim;
loopUpper = rewriter.create<arith::MinSIOp>(loc, loopUpper, maskDim);
}

if (isLoadLike) {
Expand Down Expand Up @@ -658,8 +747,11 @@ LogicalResult UnstructuredMemAccessConverter<MemAccOpTy>::matchAndRewrite(

Value extractedOffset;
if (fullyUnstructured) {
if (auto mtptOp =
srcPtr.template getDefiningOp<triton::MakeTensorPtrOp>()) {
auto mtptOp =
zeroStrideMakeTensorPtr
? zeroStrideMakeTensorPtr
: srcPtr.template getDefiningOp<triton::MakeTensorPtrOp>();
if (mtptOp) {
auto I64Type = rewriter.getIntegerType(64);
srcPtr = mtptOp.getBase();
extractedOffset = rewriter.create<arith::ConstantIntOp>(loc, 0, 64);
Expand Down
Loading
Loading