Skip to content

Add support for unstructured masks in MaskAnalysis #333

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

Merged
Brandon Myers (bmyerz0) merged 4 commits into from
Sep 3, 2025
Merged

Add support for unstructured masks in MaskAnalysis #333

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
287a9f4
Add support for generic masks in MaskAnalysis
Aug 22, 2025
45ddfc0
Update per comment.
Aug 29, 2025
b06f5f3
Update test.
Aug 29, 2025
76eb4e0
Update per comment.
Sep 3, 2025
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
5 changes: 4 additions & 1 deletion include/triton-shared/Analysis/MaskAnalysis.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,9 @@ namespace triton {
// It will in pattern of cmp -> expandDims -> broadcast
// 3. scalar_mask[:, None] where scalar mask is scalar bool.
// It will in pattern of splat -> expandDims -> broadcast
Comment thread
bmyerz0 marked this conversation as resolved.
// These 3 patterns are only about how a bool tensor was created from 1D or
// scalar bool. How the 1D and scalar bool were created is not important for the
// unstructured mask.
// Only one tensor mask is allowed. If multiple dimensions have failed
// MaskAnalysis, then MaskAnalysis will still fail on the current operation.
struct MaskState {
Expand All @@ -69,7 +72,7 @@ struct MaskState {

MaskState(bool useUnsafeMask = false) : useUnsafeMask(useUnsafeMask) {}

SmallVector<std::pair<unsigned, Value>> getGenericMasks();
SmallVector<std::pair<unsigned, Value>> getUnstructuredMasks();

int64_t getRank() const { return dims.size(); }

Expand Down
91 changes: 70 additions & 21 deletions lib/Analysis/MaskAnalysis.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -371,18 +371,63 @@ LogicalResult MaskState::parseAnd(arith::AndIOp andOp, const Location loc,

// merge the masks.
if (lhsState.masks.size() == rhsState.masks.size()) {
auto shapedType = cast<ShapedType>(andOp.getType());
assert(shapedType.hasStaticShape());
for (size_t i = 0; i < lhsState.masks.size(); i++) {
if (lhsState.masks[i] && rhsState.masks[i]) {
// And the mask.
masks.push_back(builder.create<arith::AndIOp>(loc, lhsState.masks[i],
rhsState.masks[i]));
Value lhsV = lhsState.masks[i];
Value rhsV = rhsState.masks[i];
if (!lhsV && !rhsV) {
masks.push_back(nullptr);
} else {
masks.push_back(lhsState.masks[i] ? lhsState.masks[i]
: rhsState.masks[i]);
uint32_t size = shapedType.getShape()[i];
auto structuredMaskToUnstructuredMask = [](MaskState state,
unsigned dim,
uint32_t size,
OpBuilder &builder,
Location loc) {
OpFoldResult ofr = state.isMask() ? state.dims[dim] : state.scalar;
if (auto intV = getIntAttr(ofr)) {
if (intV == size) {
// Full mask.
return Value();
}
}
auto targetTensorType =
RankedTensorType::get({size}, builder.getI32Type());
Value range =
builder
.create<triton::MakeRangeOp>(loc, targetTensorType, 0, size)
.getResult();
Value v = ofrToIndexValue(ofr, loc, builder);
v = builder
.create<arith::IndexCastUIOp>(loc, builder.getI32Type(), v)
.getResult();
v = builder.create<triton::SplatOp>(loc, targetTensorType, v)
.getResult();
return builder
.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult, range, v)
.getResult();
};
if (!lhsV) {
lhsV = structuredMaskToUnstructuredMask(lhsState, i, size, builder,
loc);
} else if (!rhsV) {
rhsV = structuredMaskToUnstructuredMask(rhsState, i, size, builder,
loc);
}
if (!lhsV) {
masks.push_back(rhsV);
continue;
} else if (!rhsV) {
masks.push_back(lhsV);
continue;
}
// And the mask.
masks.push_back(builder.create<arith::AndIOp>(loc, lhsV, rhsV));
}
}
// Only support one generic mask.
if (getGenericMasks().size() > 1) {
// Only support one unstructured mask.
if (getUnstructuredMasks().size() > 1) {
return failure();
}
}
Expand All @@ -408,6 +453,9 @@ LogicalResult MaskState::parseCmp(arith::CmpIOp cmpOp, const Location loc,
for (unsigned r = 0; r < shapedType.getRank(); r++) {
if (shapedType.getShape()[r] != 1) {
if (cmpOpDim != -1) {
// This will happen when the cmp has more than one dimension with size
// larger than 1.
// Like a < b while both a and b are tensors with shape 2x2.
cmpOpDim = -1;
Comment thread
bmyerz0 marked this conversation as resolved.
break;
}
Expand All @@ -419,11 +467,11 @@ LogicalResult MaskState::parseCmp(arith::CmpIOp cmpOp, const Location loc,
masks.push_back(nullptr);
}
// If cmpOpDim == -1, parseCmp must fail later.
// Here just setup generic masks when cmpOpDim != -1.
// Here just setup unstructured masks when cmpOpDim != -1.
if (cmpOpDim != -1) {
// Save cmpOp as generic mask for failure case, will recover it to nullptr
// later if success.
Value genericMask = cmpOp;
// Save cmpOp as unstructured mask for failure case, will recover it to
// nullptr later if success.
Value unstructuredMask = cmpOp;
if (shapedType.getRank() > 1) {
// If cmpOp is not 1D, collapse it to 1D.
auto flatType = RankedTensorType::get({shapedType.getShape()[cmpOpDim]},
Expand All @@ -433,10 +481,10 @@ LogicalResult MaskState::parseCmp(arith::CmpIOp cmpOp, const Location loc,
SmallVector<ReassociationIndices> reassociation =
*maybeReassociationMap;
// Set masks.
genericMask = builder.create<tensor::CollapseShapeOp>(
unstructuredMask = builder.create<tensor::CollapseShapeOp>(
loc, flatType, cmpOp, reassociation);
}
masks[cmpOpDim] = genericMask;
masks[cmpOpDim] = unstructuredMask;
}
} else {
cmpOpDim = 0;
Expand Down Expand Up @@ -746,23 +794,24 @@ LogicalResult MaskState::parseExpandDims(triton::ExpandDimsOp expandDimsOp,
// Recover dims to allow other dim to be processed.
dims.clear();
dims.push_back(builder.getIndexAttr(srcType.getShape()[0]));
// Save src as generic mask.
// Save src as unstructured mask.
masks[1 - axis] = src;
} else {
// save nullptr when parse success.
masks[1 - axis] = nullptr;
}
} else {
if (failed(result)) {
auto genericMasks = getGenericMasks();
if (genericMasks.empty()) {
auto unstructuredMasks = getUnstructuredMasks();
if (unstructuredMasks.empty()) {
return failure();
}
if (genericMasks.size() > 1) {
if (unstructuredMasks.size() > 1) {
return failure();
}
auto [dim, mask] = genericMasks.front();
// Recover dims for generic mask dim to allow other dim to be processed.
auto [dim, mask] = unstructuredMasks.front();
// Recover dims for unstructured mask dim to allow other dim to be
// processed.
dims[dim] = builder.getIndexAttr(srcType.getShape()[dim]);
}
masks.insert(masks.begin() + axis, nullptr);
Expand All @@ -777,7 +826,7 @@ LogicalResult MaskState::parseExpandDims(triton::ExpandDimsOp expandDimsOp,
}

// Return all non-nullptr masks along with their dimensions.
SmallVector<std::pair<unsigned, Value>> MaskState::getGenericMasks() {
SmallVector<std::pair<unsigned, Value>> MaskState::getUnstructuredMasks() {
SmallVector<std::pair<unsigned, Value>> result;

for (auto [i, m] : llvm::enumerate(masks)) {
Comment thread
bmyerz0 marked this conversation as resolved.
Expand Down
112 changes: 61 additions & 51 deletions lib/AnalysisStructured/PtrAnalysis.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,70 +39,81 @@

using namespace mlir;

// Try to apply generic mask on the ptr.
static Value applyGenericMask(Operation *op, Value ptr,
triton::MaskState &mstate, Location loc,
OpBuilder builder) {
SmallVector<std::pair<unsigned, Value>> masks = mstate.getGenericMasks();
// Try to apply unstructured mask on the ptr.
static Value applyUnstructuredMask(Operation *op, Value ptr,
Comment thread
bmyerz0 marked this conversation as resolved.
triton::MaskState &mstate, Location loc,
OpBuilder builder) {
SmallVector<std::pair<unsigned, Value>> masks = mstate.getUnstructuredMasks();
if (masks.empty()) {
return ptr;
}
if (masks.size() > 1) {
op->emitRemark("MaskAnalysis failed for more than one generic masks");
op->emitRemark("MaskAnalysis failed for more than one unstructured masks");
return nullptr;
}

auto [dim, genericMask] = masks[0];
auto [dim, unstructuredMask] = masks[0];
if (auto scatterPtr =
ptr.getDefiningOp<tts::MakeGatherScatterTensorPtrOp>()) {
if (dim != scatterPtr.getGatherScatterDim()) {
op->emitRemark("MaskAnalysis failed for generic mask dim not equal "
op->emitRemark("MaskAnalysis failed for unstructured mask dim not equal "
"gather scatter dim");
return nullptr;
}

ptr = builder
.create<tts::MakeGatherScatterTensorPtrOp>(
loc, scatterPtr.getBase(),
scatterPtr.getGatherScatterOffset(), genericMask,
scatterPtr.getGatherScatterOffset(), unstructuredMask,
scatterPtr.getGatherScatterDim(), scatterPtr.getSizes(),
scatterPtr.getMixedStrides(), scatterPtr.getMixedOffsets())
.getResult();

} else if (auto tptr = ptr.getDefiningOp<tts::MakeTensorPtrOp>()) {
OpFoldResult offsetFold = tptr.getMixedOffsets()[dim];
Value offset = dyn_cast<Value>(offsetFold);
if (!offset) {
offset = builder
.create<arith::ConstantOp>(
loc, cast<TypedAttr>(cast<Attribute>(offsetFold)))
.getResult();
}
// Cast to integer for splat and makerange.
if (isa<IndexType>(offset.getType())) {
offset =
builder.create<arith::IndexCastOp>(loc, builder.getI32Type(), offset)
.getResult();
} else if (offset.getType().isInteger(64)) {
offset =
builder.create<arith::TruncIOp>(loc, builder.getI32Type(), offset)
.getResult();
}
auto offsetRowType =
RankedTensorType::get({tptr.getSizes()[dim]}, offset.getType());
Value scatterOffset =
} else if (auto structuredPtr = ptr.getDefiningOp<tts::MakeTensorPtrOp>()) {
auto ofrToI32Value = [&](OpFoldResult ofr) {
Value v = dyn_cast<Value>(ofr);
if (!v) {
v = builder
.create<arith::ConstantOp>(
loc, cast<TypedAttr>(cast<Attribute>(ofr)))
.getResult();
}
if (isa<IndexType>(v.getType())) {
v = builder.create<arith::IndexCastOp>(loc, builder.getI32Type(), v)
.getResult();
} else if (v.getType().isInteger(64)) {
v = builder.create<arith::TruncIOp>(loc, builder.getI32Type(), v)
.getResult();
}

return v;
};
OpFoldResult offsetFold = structuredPtr.getMixedOffsets()[dim];
Value offset = ofrToI32Value(offsetFold);
auto offsetRowType = RankedTensorType::get({structuredPtr.getSizes()[dim]},
offset.getType());
OpFoldResult strideFold = structuredPtr.getMixedStrides()[dim];
Value stride = ofrToI32Value(strideFold);
// Divide stride since offset of tts::MakeTensorPtrOp already include the
// stride, but gatherScatterOffset of tts::MakeGatherScatterTensorPtrOp
// should not include stride.
offset = builder.create<arith::DivUIOp>(loc, offset, stride);

Value gatherScatterOffset =
builder.create<tensor::SplatOp>(loc, offsetRowType, offset).getResult();
Value range = builder
.create<triton::MakeRangeOp>(loc, offsetRowType, 0,
tptr.getSizes()[dim])
.create<triton::MakeRangeOp>(
loc, offsetRowType, 0, structuredPtr.getSizes()[dim])
.getResult();
scatterOffset = builder.create<arith::AddIOp>(loc, scatterOffset, range);
ptr =
builder
.create<tts::MakeGatherScatterTensorPtrOp>(
loc, tptr.getBase(), scatterOffset, genericMask, dim,
tptr.getSizes(), tptr.getMixedStrides(), tptr.getMixedOffsets())
.getResult();
gatherScatterOffset =
builder.create<arith::AddIOp>(loc, gatherScatterOffset, range);
ptr = builder
.create<tts::MakeGatherScatterTensorPtrOp>(
loc, structuredPtr.getBase(), gatherScatterOffset,
unstructuredMask, dim, structuredPtr.getSizes(),
structuredPtr.getMixedStrides(),
structuredPtr.getMixedOffsets())
.getResult();
} else {
return nullptr;
}
Expand Down Expand Up @@ -358,13 +369,13 @@ LogicalResult PtrState::addState(const PtrState &lhsState,
}

if (lhsStride == rhsStride) {
// For case like lhs_offset * stride + rhs_offset * stride, it is same as
// (lhs_offset + rhs_offset) * stride.
// We can just
// add the offsets and reuse the stride like this:
// For case like lhs_offset * stride + rhs_offset * stride, it is
// same as (lhs_offset + rhs_offset) * stride. We can just add the
// offsets and reuse the stride like this:
// offsets[i] = lhsOffset + rhsOffset
// strides[i] = lhsStride
// Expand structured offset since unstructured offset has tensor type.
// Expand structured offset since unstructured offset has tensor
// type.
if (!lhsState.dimIsStructured(i)) {
rhsOffset = expandOFRIndex(rhsOffset, lhsOffset, loc, builder);
} else {
Expand All @@ -380,10 +391,9 @@ LogicalResult PtrState::addState(const PtrState &lhsState,
// equal to 1 earlier for case both offsets and strides not equal.
assert(lhsOffset == rhsOffset &&
"If strides are not equal, offsets must be equal");
// For case like offset * lhs_stride + offset * rhs_stride, it is same as
// offset * (lhs_stride + rhs_stride).
// We can just
// add the strides and reuse the offset like this:
// For case like offset * lhs_stride + offset * rhs_stride, it is
// same as offset * (lhs_stride + rhs_stride). We can just add the
// strides and reuse the offset like this:
// offsets[i] = lhsOffset
// strides[i] = lhsStride + rhsStride

Expand Down Expand Up @@ -1676,7 +1686,7 @@ LogicalResult PtrAnalysis::rewriteLoadOp(triton::LoadOp op,
op->emitRemark("MaskAnalysis failed");
return failure();
}
ptr = applyGenericMask(op, ptr, mstate, loc, builder);
ptr = applyUnstructuredMask(op, ptr, mstate, loc, builder);
if (!ptr) {
return failure();
}
Expand Down Expand Up @@ -1815,7 +1825,7 @@ LogicalResult PtrAnalysis::rewriteStoreOp(triton::StoreOp op,
op->emitRemark("MaskAnalysis failed");
return failure();
}
ptr = applyGenericMask(op, ptr, mstate, loc, builder);
ptr = applyUnstructuredMask(op, ptr, mstate, loc, builder);
if (!ptr) {
return failure();
}
Expand Down
16 changes: 8 additions & 8 deletions lib/Conversion/StructuredToMemref/StructuredToMemref.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -886,8 +886,8 @@ struct LoadConverter : public OpConversionPattern<tts::LoadOp> {
// If the gather mask dimension is a constant, we can use it directly.
unsigned gatherMaskDimValue = gatherMaskDimIndex.value();
if (gatherMaskDimValue == 0 && ptr.getGatherScatterMask()) {
// For generic mask case, do the full loop and use the generic mask to
// guard the store.
// For unstructured mask case, loop over all elements and use the
// unstructured mask to guard the store.
gatherMaskDimValue = offsetSize;
}
offsetSize = std::min(offsetSize, gatherMaskDimValue);
Expand Down Expand Up @@ -915,11 +915,11 @@ struct LoadConverter : public OpConversionPattern<tts::LoadOp> {

Value inductionVar = loop.getInductionVar();

if (Value genericMask = ptr.getGatherScatterMask()) {
if (Value unstructuredMask = ptr.getGatherScatterMask()) {
// If the gather scatter mask is present, we need to use it to guard the
// load.
auto maskValue = rewriter.create<tensor::ExtractOp>(
loc, genericMask, ValueRange{inductionVar});
loc, unstructuredMask, ValueRange{inductionVar});
auto ifOp = rewriter.create<scf::IfOp>(loc, maskValue);
rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front());
}
Expand Down Expand Up @@ -1048,8 +1048,8 @@ struct StoreConverter : public OpConversionPattern<tts::StoreOp> {
// If the gather mask dimension is a constant, we can use it directly.
unsigned gatherMaskDimValue = gatherMaskDimIndex.value();
if (gatherMaskDimValue == 0 && ptr.getGatherScatterMask()) {
// For generic mask case, do the full loop and use the generic mask to
// guard the store.
// For unstructured mask case, loop over all elements and use the
// unstructured mask to guard the store.
gatherMaskDimValue = offsetSize;
}
offsetSize = std::min(offsetSize, gatherMaskDimValue);
Expand All @@ -1072,11 +1072,11 @@ struct StoreConverter : public OpConversionPattern<tts::StoreOp> {

Value inductionVar = loop.getInductionVar();

if (Value genericMask = ptr.getGatherScatterMask()) {
if (Value unstructuredMask = ptr.getGatherScatterMask()) {
// If the gather scatter mask is present, we need to use it to guard the
// store.
auto maskValue = rewriter.create<tensor::ExtractOp>(
loc, genericMask, ValueRange{inductionVar});
loc, unstructuredMask, ValueRange{inductionVar});
auto ifOp = rewriter.create<scf::IfOp>(loc, maskValue);
rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front());
}
Expand Down
Loading