From 287a9f492756f2377d477897c65856c13e43066d Mon Sep 17 00:00:00 2001 From: Xiang Li Date: Fri, 22 Aug 2025 17:33:14 +0000 Subject: [PATCH 1/4] Add support for generic masks in MaskAnalysis - Enhanced `MaskAnalysis` to handle generic masks. MaskState will save generic mask for a dimension if that dimension failed MaskAnalysis. - Updated `TritonStructuredDialect` to include `gather_scatter_mask` in `MakeGatherScatterTensorPtrOp`. - Modified `PtrAnalysis` to apply generic masks during pointer analysis. - Extended `StructuredToMemref` conversion to support generic masks in load/store operations. - Added verification logic for `MakeGatherScatterTensorPtrOp` to ensure compatibility with masks. --- include/triton-shared/Analysis/MaskAnalysis.h | 16 + .../IR/TritonStructuredDialect.td | 21 +- lib/Analysis/MaskAnalysis.cpp | 166 +++++++++- lib/AnalysisStructured/PtrAnalysis.cpp | 82 +++++ .../StructuredToMemref/StructuredToMemref.cpp | 33 +- .../IR/TritonStructuredOps.cpp | 147 ++++++++- python/examples/test_generic_mask.py | 284 ++++++++++++++++++ .../generic_mask_2d_kernel.mlir | 95 ++++++ ...ic_mask_2d_non_continuous_load_kernel.mlir | 84 ++++++ ...c_mask_2d_non_continuous_store_kernel.mlir | 88 ++++++ .../generic_mask_3d_kernel.mlir | 128 ++++++++ ...ic_mask_3d_non_continuous_load_kernel.mlir | 121 ++++++++ ...c_mask_3d_non_continuous_store_kernel.mlir | 119 ++++++++ 13 files changed, 1363 insertions(+), 21 deletions(-) create mode 100644 python/examples/test_generic_mask.py create mode 100644 test/Conversion/TritonToStructured/generic_mask_2d_kernel.mlir create mode 100644 test/Conversion/TritonToStructured/generic_mask_2d_non_continuous_load_kernel.mlir create mode 100644 test/Conversion/TritonToStructured/generic_mask_2d_non_continuous_store_kernel.mlir create mode 100644 test/Conversion/TritonToStructured/generic_mask_3d_kernel.mlir create mode 100644 test/Conversion/TritonToStructured/generic_mask_3d_non_continuous_load_kernel.mlir create mode 100644 test/Conversion/TritonToStructured/generic_mask_3d_non_continuous_store_kernel.mlir diff --git a/include/triton-shared/Analysis/MaskAnalysis.h b/include/triton-shared/Analysis/MaskAnalysis.h index 3dd1c6e93..4faa1b660 100644 --- a/include/triton-shared/Analysis/MaskAnalysis.h +++ b/include/triton-shared/Analysis/MaskAnalysis.h @@ -44,10 +44,24 @@ namespace triton { // // Example of creating 2D mask: // mask = (rows[:, None] < M) & (cols[None, :] < N) +// +// Bool tensor mask could be saved into masks in case that dimension failed +// MaskAnalysis. These is to allow case where only one dimension failed while +// others passed. A MakeGatherScatterTensorPtrOp operation could be generated +// for the failed dimension. Only 3 patterns are supported for this. +// 1. offsets[:, None] < n where the offsets is 1d tensor. +// It will in pattern of expandDims -> broadcast -> cmp +// 2. mask[:, None] where mask is 1d bool tensor. +// 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 +// Only one tensor mask is allowed. If multiple dimensions have failed +// MaskAnalysis, then MaskAnalysis will still fail on the current operation. struct MaskState { OpFoldResult start; OpFoldResult end; SmallVector dims; + SmallVector masks; OpFoldResult scalar; const bool useUnsafeMask; @@ -55,6 +69,8 @@ struct MaskState { MaskState(bool useUnsafeMask = false) : useUnsafeMask(useUnsafeMask) {} + SmallVector> getGenericMasks(); + int64_t getRank() const { return dims.size(); } bool isEmpty() const { return getRank() == 0 && !scalar && !start && !end; } diff --git a/include/triton-shared/Dialect/TritonStructured/IR/TritonStructuredDialect.td b/include/triton-shared/Dialect/TritonStructured/IR/TritonStructuredDialect.td index ae8065ced..17c3ce9a8 100644 --- a/include/triton-shared/Dialect/TritonStructured/IR/TritonStructuredDialect.td +++ b/include/triton-shared/Dialect/TritonStructured/IR/TritonStructuredDialect.td @@ -159,6 +159,7 @@ def TTS_MakeGatherScatterTensorPtrOp // strides: The strides of the parent tensor, which means how much to increase the pointer // by when moving by 1 element in a specific axis. // offsets: Offset of the block along each dimension from base. + // gather_scatter_mask: Optional bool mask for mask which failed MaskAnalysis. // result: A tensor of pointers. let arguments = (ins TT_Ptr:$base, @@ -168,7 +169,8 @@ def TTS_MakeGatherScatterTensorPtrOp Variadic:$strides, Variadic:$offsets, DenseI64ArrayAttr:$static_strides, - DenseI64ArrayAttr:$static_offsets); + DenseI64ArrayAttr:$static_offsets, + Optional:$gather_scatter_mask); let results = (outs TT_PtrLike:$result); @@ -176,11 +178,12 @@ def TTS_MakeGatherScatterTensorPtrOp $base `to` `sizes` `` `:` $sizes `gather_scatter_dim` `` `:` $gather_scatter_dim `gather_scatter_offset` `` `:` $gather_scatter_offset + (`gather_scatter_mask` `` `:` $gather_scatter_mask^)? `` `,` `strides` `` `:` custom($strides, $static_strides) `` `,` `offsets` `` `:` custom($offsets, $static_offsets) - attr-dict `:` type($gather_scatter_offset) type($base) `to` type($result) + attr-dict `:` type($gather_scatter_offset) type($gather_scatter_mask) type($base) `to` type($result) }]; @@ -193,6 +196,15 @@ def TTS_MakeGatherScatterTensorPtrOp "ArrayRef":$sizes, "ArrayRef":$strides, "ArrayRef":$offsets)>, + + OpBuilder<(ins + "Value":$base, + "Value":$gather_scatter_offset, + "Value":$gather_scatter_mask, + "int":$gather_scatter_dim, + "ArrayRef":$sizes, + "ArrayRef":$strides, + "ArrayRef":$offsets)>, ]; let extraClassDeclaration = [{ @@ -213,9 +225,8 @@ def TTS_MakeGatherScatterTensorPtrOp } }]; - // TODO - //let hasVerifier = 1; - //let hasCanonicalizer = 1; + let hasVerifier = 1; + let hasCanonicalizer = 0; } def TTS_GetStructuredStateOp : TTS_Op<"get_structured_state", [AttrSizedResultSegments, Pure]> { diff --git a/lib/Analysis/MaskAnalysis.cpp b/lib/Analysis/MaskAnalysis.cpp index df40b436e..0f95dcc08 100644 --- a/lib/Analysis/MaskAnalysis.cpp +++ b/lib/Analysis/MaskAnalysis.cpp @@ -316,6 +316,11 @@ void MaskState::dump() const { llvm::dbgs() << "dims: "; for (auto dim : dims) llvm::dbgs() << "\t" << dim << "\n"; + if (!masks.empty()) { + llvm::dbgs() << "masks: "; + for (auto mask : masks) + llvm::dbgs() << "\t" << mask << "\n"; + } llvm::dbgs() << "\n"; } @@ -337,14 +342,51 @@ LogicalResult MaskState::parseAdd(arith::AddIOp addOp, const Location loc, LogicalResult MaskState::parseAnd(arith::AndIOp andOp, const Location loc, OpBuilder &builder) { assert(this->isEmpty()); - + bool isBoolOp = false; + unsigned rank = 1; + if (auto shapedType = dyn_cast(andOp.getType())) { + isBoolOp = shapedType.getElementType().isInteger(1); + rank = shapedType.getRank(); + } MaskState lhsState; - if (failed(lhsState.parse(andOp.getLhs(), loc, builder))) + LogicalResult lResult = lhsState.parse(andOp.getLhs(), loc, builder); + if (failed(lResult) && !isBoolOp) { return failure(); + } MaskState rhsState; - if (failed(rhsState.parse(andOp.getRhs(), loc, builder))) + LogicalResult rResult = rhsState.parse(andOp.getRhs(), loc, builder); + if (failed(rResult) && !isBoolOp) { return failure(); + } + + if (isBoolOp) { + if (lhsState.masks.size() != rank) { + return failure(); + } + + if (lhsState.masks.size() != rhsState.masks.size()) { + return failure(); + } + + // merge the masks. + if (lhsState.masks.size() == rhsState.masks.size()) { + 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(loc, lhsState.masks[i], + rhsState.masks[i])); + } else { + masks.push_back(lhsState.masks[i] ? lhsState.masks[i] + : rhsState.masks[i]); + } + } + // Only support one generic mask. + if (getGenericMasks().size() > 1) { + return failure(); + } + } + } if (!lhsState.isMask() || !rhsState.isMask()) { return this->minStateScalar(lhsState, rhsState, loc, builder); @@ -361,7 +403,45 @@ LogicalResult MaskState::parseExtSI(arith::ExtSIOp op, const Location loc, LogicalResult MaskState::parseCmp(arith::CmpIOp cmpOp, const Location loc, OpBuilder &builder) { assert(this->isEmpty()); - + int cmpOpDim = -1; + if (auto shapedType = dyn_cast(cmpOp.getType())) { + for (unsigned r = 0; r < shapedType.getRank(); r++) { + if (shapedType.getShape()[r] != 1) { + if (cmpOpDim != -1) { + cmpOpDim = -1; + break; + } + cmpOpDim = r; + } + } + masks.clear(); + for (unsigned r = 0; r < shapedType.getRank(); r++) { + masks.push_back(nullptr); + } + // If cmpOpDim == -1, parseCmp must fail later. + // Here just setup generic 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; + if (shapedType.getRank() > 1) { + // If cmpOp is not 1D, collapse it to 1D. + auto flatType = RankedTensorType::get({shapedType.getShape()[cmpOpDim]}, + shapedType.getElementType()); + auto maybeReassociationMap = + getReassociationIndicesForReshape(shapedType, flatType); + SmallVector reassociation = + *maybeReassociationMap; + // Set masks. + genericMask = builder.create( + loc, flatType, cmpOp, reassociation); + } + masks[cmpOpDim] = genericMask; + } + } else { + cmpOpDim = 0; + masks.push_back(cmpOp); + } if (cmpOp.getPredicate() != arith::CmpIPredicate::slt && cmpOp.getPredicate() != arith::CmpIPredicate::ult && cmpOp.getPredicate() != arith::CmpIPredicate::sge) { @@ -449,7 +529,10 @@ LogicalResult MaskState::parseCmp(arith::CmpIOp cmpOp, const Location loc, else this->dims.push_back(lhsState.dims[i]); } - + if (cmpOpDim != -1) { + // Clear masks when success. + masks[cmpOpDim] = nullptr; + } return success(); } @@ -619,7 +702,15 @@ LogicalResult MaskState::parseSplat(triton::SplatOp splatOp, const Location loc, for (auto s : dstShape) this->dims.push_back(builder.getIndexAttr(s)); - + bool isBool = src.getType().isInteger(1); + if (isBool) { + // If src is a 1D boolean tensor and parse success. + // Create masks. + masks.clear(); + for (unsigned i = 0; i < dstShape.size(); i++) { + masks.push_back(nullptr); + } + } return success(); } @@ -628,12 +719,56 @@ LogicalResult MaskState::parseExpandDims(triton::ExpandDimsOp expandDimsOp, OpBuilder &builder) { assert(this->isEmpty()); - if (failed(this->parse(expandDimsOp.getSrc(), loc, builder))) - return failure(); - auto dstShape = cast(expandDimsOp.getResult().getType()).getShape(); auto axis = expandDimsOp.getAxis(); + Value src = expandDimsOp.getSrc(); + auto srcType = cast(src.getType()); + bool isBoolOp = srcType.getElementType().isInteger(1); + LogicalResult result = parse(src, loc, builder); + if (failed(result)) { + if (isBoolOp) { + if (srcType.getRank() > 1 && masks.size() != srcType.getRank()) { + return failure(); + } + } else { + return failure(); + } + } + + if (isBoolOp) { + // Save mask for 1D boolean tensor + if (srcType.getRank() == 1) { + assert(dstShape.size() == 2); + masks.resize(dstShape.size()); + masks[axis] = nullptr; + if (failed(result)) { + // Recover dims to allow other dim to be processed. + dims.clear(); + dims.push_back(builder.getIndexAttr(srcType.getShape()[0])); + // Save src as generic 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()) { + return failure(); + } + if (genericMasks.size() > 1) { + return failure(); + } + auto [dim, mask] = genericMasks.front(); + // Recover dims for generic mask dim to allow other dim to be processed. + dims[dim] = builder.getIndexAttr(srcType.getShape()[dim]); + } + masks.insert(masks.begin() + axis, nullptr); + } + } + assert(dstShape[axis] == 1 && "expect changed dimension to be 1 in expand_dims"); this->dims.insert(this->dims.begin() + axis, builder.getIndexAttr(1)); @@ -641,5 +776,18 @@ LogicalResult MaskState::parseExpandDims(triton::ExpandDimsOp expandDimsOp, return success(); } +// Return all non-nullptr masks along with their dimensions. +SmallVector> MaskState::getGenericMasks() { + SmallVector> result; + + for (auto [i, m] : llvm::enumerate(masks)) { + if (m) { + result.push_back({i, m}); + } + } + + return result; +} + } // namespace triton } // namespace mlir diff --git a/lib/AnalysisStructured/PtrAnalysis.cpp b/lib/AnalysisStructured/PtrAnalysis.cpp index 671c73590..cafc18f83 100644 --- a/lib/AnalysisStructured/PtrAnalysis.cpp +++ b/lib/AnalysisStructured/PtrAnalysis.cpp @@ -37,6 +37,80 @@ #define DEBUG_TYPE "triton-ptr-analysis" +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> masks = mstate.getGenericMasks(); + if (masks.empty()) { + return ptr; + } + if (masks.size() > 1) { + op->emitRemark("MaskAnalysis failed for more than one generic masks"); + return nullptr; + } + + auto [dim, genericMask] = masks[0]; + if (auto scatterPtr = + ptr.getDefiningOp()) { + if (dim != scatterPtr.getGatherScatterDim()) { + op->emitRemark("MaskAnalysis failed for generic mask dim not equal " + "gather scatter dim"); + return nullptr; + } + + ptr = builder + .create( + loc, scatterPtr.getBase(), + scatterPtr.getGatherScatterOffset(), genericMask, + scatterPtr.getGatherScatterDim(), scatterPtr.getSizes(), + scatterPtr.getMixedStrides(), scatterPtr.getMixedOffsets()) + .getResult(); + + } else if (auto tptr = ptr.getDefiningOp()) { + OpFoldResult offsetFold = tptr.getMixedOffsets()[dim]; + Value offset = dyn_cast(offsetFold); + if (!offset) { + offset = builder + .create( + loc, cast(cast(offsetFold))) + .getResult(); + } + // Cast to integer for splat and makerange. + if (isa(offset.getType())) { + offset = + builder.create(loc, builder.getI32Type(), offset) + .getResult(); + } else if (offset.getType().isInteger(64)) { + offset = + builder.create(loc, builder.getI32Type(), offset) + .getResult(); + } + auto offsetRowType = + RankedTensorType::get({tptr.getSizes()[dim]}, offset.getType()); + Value scatterOffset = + builder.create(loc, offsetRowType, offset).getResult(); + Value range = builder + .create(loc, offsetRowType, 0, + tptr.getSizes()[dim]) + .getResult(); + scatterOffset = builder.create(loc, scatterOffset, range); + ptr = + builder + .create( + loc, tptr.getBase(), scatterOffset, genericMask, dim, + tptr.getSizes(), tptr.getMixedStrides(), tptr.getMixedOffsets()) + .getResult(); + } else { + return nullptr; + } + // Clear the mask size for gather/scatter dim. + mstate.dims[dim] = OpFoldResult(builder.getI32IntegerAttr(0)); + return ptr; +} + namespace mlir { namespace tts { @@ -1602,6 +1676,10 @@ LogicalResult PtrAnalysis::rewriteLoadOp(triton::LoadOp op, op->emitRemark("MaskAnalysis failed"); return failure(); } + ptr = applyGenericMask(op, ptr, mstate, loc, builder); + if (!ptr) { + return failure(); + } dims = mstate.dims; } @@ -1737,6 +1815,10 @@ LogicalResult PtrAnalysis::rewriteStoreOp(triton::StoreOp op, op->emitRemark("MaskAnalysis failed"); return failure(); } + ptr = applyGenericMask(op, ptr, mstate, loc, builder); + if (!ptr) { + return failure(); + } dims = mstate.dims; } diff --git a/lib/Conversion/StructuredToMemref/StructuredToMemref.cpp b/lib/Conversion/StructuredToMemref/StructuredToMemref.cpp index 4ee575a88..afd7e67ad 100644 --- a/lib/Conversion/StructuredToMemref/StructuredToMemref.cpp +++ b/lib/Conversion/StructuredToMemref/StructuredToMemref.cpp @@ -885,6 +885,11 @@ struct LoadConverter : public OpConversionPattern { if (auto gatherMaskDimIndex = getIntAttr(gatherMaskDim)) { // 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. + gatherMaskDimValue = offsetSize; + } offsetSize = std::min(offsetSize, gatherMaskDimValue); upperBound = rewriter.create(loc, offsetSize).getResult(); } else { @@ -908,8 +913,18 @@ struct LoadConverter : public OpConversionPattern { // Build loop body. rewriter.setInsertionPointToStart(loop.getBody()); - // Load the offsetElt first. Value inductionVar = loop.getInductionVar(); + + if (Value genericMask = ptr.getGatherScatterMask()) { + // If the gather scatter mask is present, we need to use it to guard the + // load. + auto maskValue = rewriter.create( + loc, genericMask, ValueRange{inductionVar}); + auto ifOp = rewriter.create(loc, maskValue); + rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front()); + } + + // Load the offsetElt first. auto gatherOffsetElt = rewriter.create( loc, gatherOffset, ValueRange{inductionVar}); @@ -1032,6 +1047,11 @@ struct StoreConverter : public OpConversionPattern { if (auto gatherMaskDimIndex = getIntAttr(gatherMaskDim)) { // 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. + gatherMaskDimValue = offsetSize; + } offsetSize = std::min(offsetSize, gatherMaskDimValue); upperBound = rewriter.create(loc, offsetSize).getResult(); } else { @@ -1050,9 +1070,18 @@ struct StoreConverter : public OpConversionPattern { // Build loop body. rewriter.setInsertionPointToStart(loop.getBody()); - // Load the offsetElt first. Value inductionVar = loop.getInductionVar(); + if (Value genericMask = ptr.getGatherScatterMask()) { + // If the gather scatter mask is present, we need to use it to guard the + // store. + auto maskValue = rewriter.create( + loc, genericMask, ValueRange{inductionVar}); + auto ifOp = rewriter.create(loc, maskValue); + rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front()); + } + + // Load the offsetElt first. auto gatherOffsetElt = rewriter.create( loc, gatherOffset, ValueRange{inductionVar}); diff --git a/lib/Dialect/TritonStructured/IR/TritonStructuredOps.cpp b/lib/Dialect/TritonStructured/IR/TritonStructuredOps.cpp index 129db34a4..df7de9962 100644 --- a/lib/Dialect/TritonStructured/IR/TritonStructuredOps.cpp +++ b/lib/Dialect/TritonStructured/IR/TritonStructuredOps.cpp @@ -132,10 +132,11 @@ void MakeTensorPtrOp::build(OpBuilder &b, OperationState &state, Value base, } void MakeGatherScatterTensorPtrOp::build(OpBuilder &b, OperationState &state, - Value base, Value gatherScatterOffset, - int gatherScatterDim, ArrayRef sizes, - ArrayRef strides, - ArrayRef offsets) { + Value base, Value gatherScatterOffset, + int gatherScatterDim, + ArrayRef sizes, + ArrayRef strides, + ArrayRef offsets) { SmallVector staticStrides, staticOffsets; SmallVector dynamicStrides, dynamicOffsets; for (auto [i, offset] : llvm::enumerate(offsets)) { @@ -156,7 +157,143 @@ void MakeGatherScatterTensorPtrOp::build(OpBuilder &b, OperationState &state, build(b, state, resType, base, gatherScatterOffset, b.getI32IntegerAttr(gatherScatterDim), b.getDenseI64ArrayAttr(sizes), dynamicStrides, dynamicOffsets, b.getDenseI64ArrayAttr(staticStrides), - b.getDenseI64ArrayAttr(staticOffsets)); + b.getDenseI64ArrayAttr(staticOffsets), Value()); +} + +void MakeGatherScatterTensorPtrOp::build( + OpBuilder &b, OperationState &state, Value base, Value gatherScatterOffset, + Value gatherScatterMask, int gatherScatterDim, ArrayRef sizes, + ArrayRef strides, ArrayRef offsets) { + SmallVector staticStrides, staticOffsets; + SmallVector dynamicStrides, dynamicOffsets; + for (auto [i, offset] : llvm::enumerate(offsets)) { + if (i != gatherScatterDim) + dispatchIndexOpFoldResult(offset, dynamicOffsets, staticOffsets); + else + staticOffsets.push_back(0); + } + dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides); + + Type resType; + auto basePtr = cast(base.getType()); + auto elemType = basePtr.getPointeeType(); + + if (gatherScatterOffset.getType().isIntOrIndex()) { + assert(sizes.size() == 1 && sizes[0] == 1 && + "gatherScatterOffset should be a scalar for 1D gather/scatter"); + resType = triton::PointerType::get(elemType, basePtr.getAddressSpace()); + + } else { + resType = triton::PointerType::get(RankedTensorType::get(sizes, elemType), + basePtr.getAddressSpace()); + } + + build(b, state, resType, base, gatherScatterOffset, + b.getI32IntegerAttr(gatherScatterDim), b.getDenseI64ArrayAttr(sizes), + dynamicStrides, dynamicOffsets, b.getDenseI64ArrayAttr(staticStrides), + b.getDenseI64ArrayAttr(staticOffsets), gatherScatterMask); +} + +LogicalResult MakeGatherScatterTensorPtrOp::verify() { + // Verify that the gatherScatterDim is within the valid range. + if (getGatherScatterDim() < 0 || getGatherScatterDim() >= getSizes().size()) { + return emitError("gatherScatterDim is out of bounds"); + } + + // Verify that the sizes, strides, and offsets have compatible dimensions. + if (getMixedSizes().size() != getMixedStrides().size() || + getMixedSizes().size() != getMixedOffsets().size()) { + return emitError( + "sizes, strides, and offsets must have the same number of dimensions"); + } + + Type offsetType = getGatherScatterOffset().getType(); + int64_t offsetSize = 0; + Type offsetEltType = offsetType; + // Verify that the gatherScatterOffset is a 1D tensor. + auto rankedTensorType = dyn_cast(offsetType); + if (!rankedTensorType) { + return emitError("gatherScatterOffset must be a 1D tensor"); + } + if (rankedTensorType.getRank() != 1) { + return emitError("gatherScatterOffset must be a 1D tensor"); + } + offsetSize = rankedTensorType.getShape()[0]; + offsetEltType = rankedTensorType.getElementType(); + + if (!offsetEltType.isIntOrIndex()) { + return emitError("gatherScatterOffset must be a 1D tensor of " + "int or index type"); + } + + // Verify that the gatherScatterMask, if provided, is a 1D tensor. + if (getGatherScatterMask()) { + Type maskType = getGatherScatterMask().getType(); + Type maskEltType = maskType; + auto rankedTensorType = dyn_cast(maskType); + if (!rankedTensorType) { + return emitError("gatherScatterMask must be a 1D tensor"); + } + if (rankedTensorType.getRank() != 1) { + return emitError("gatherScatterMask must be a 1D tensor of boolean type"); + } + // Verify that the gatherScatterMask has the same size as the + // gatherScatterOffset. + if (rankedTensorType.getShape()[0] != offsetSize) { + return emitError( + "gatherScatterMask must have the same size as gatherScatterOffset"); + } + maskEltType = rankedTensorType.getElementType(); + if (!maskEltType.isInteger(1)) { + return emitError("gatherScatterMask must be a 1D tensor of boolean type"); + } + } + + // Verify that when gatherScatterMask is provided, all the user of + // MakeGatherScatterTensorPtrOp must have mask with size of 0. + if (getGatherScatterMask()) { + for (auto user : (*this)->getUsers()) { + if (auto loadOp = dyn_cast(user)) { + if (loadOp.hasMask()) { + OpFoldResult MaskedSize = + loadOp.getMixedMaskDims()[getGatherScatterDim()]; + auto intAttr = + dyn_cast_if_present(dyn_cast(MaskedSize)); + if (!intAttr || intAttr.getInt() != 0) { + return emitError("tts.load user of tts.make_gather_scatter_tptr " + "with gather_scatter_mask must have " + "mask size of 0 for gather_scatter_dim"); + } + } else { + return emitError("tts.load user of tts.make_gather_scatter_tptr with " + "gather_scatter_mask must have " + "mask provided"); + } + } else if (auto storeOp = dyn_cast(user)) { + if (storeOp.hasMask()) { + OpFoldResult MaskedSize = + storeOp.getMixedMaskDims()[getGatherScatterDim()]; + auto intAttr = + dyn_cast_if_present(dyn_cast(MaskedSize)); + if (!intAttr || intAttr.getInt() != 0) { + return emitError("tts.store user of tts.make_gather_scatter_tptr " + "with gather_scatter_mask must have " + "mask size of 0 for gather_scatter_dim"); + } + } else { + return emitError( + "tts.store user of tts.make_gather_scatter_tptr with " + "gather_scatter_mask must have " + "mask provided"); + } + } else { + return emitError("tts.make_gather_scatter_tptr can only be used in " + "tts.load or tts.store operations"); + } + } + } + + return success(); } void LoadOp::build(OpBuilder &b, OperationState &state, Value ptr, diff --git a/python/examples/test_generic_mask.py b/python/examples/test_generic_mask.py new file mode 100644 index 000000000..35ef65c29 --- /dev/null +++ b/python/examples/test_generic_mask.py @@ -0,0 +1,284 @@ +import torch + +import triton +import triton.language as tl + +from triton.backends.triton_shared.driver import CPUDriver + +@triton.jit +def generic_mask_2d_kernel(in_ptr, out_ptr, mask_m_ptr, mask_n_ptr, m, n, M:tl.constexpr, N:tl.constexpr): + offs_m = tl.arange(0, M) + offs_n = tl.arange(0, N) + + mask_m = tl.load(mask_m_ptr + offs_m, mask=offs_m < m, other=0) != 0 + mask_n = tl.load(mask_n_ptr + offs_n, mask=offs_n < n, other=0) != 0 + + in_ptrs = in_ptr + offs_m[:, None] * N + offs_n[None, :] + v = tl.load(in_ptrs, mask=mask_m[:, None] and offs_n[None, :] < n, other=-2) + out_ptrs = out_ptr + offs_m[:, None] * N + offs_n[None, :] + tl.store(out_ptrs, v, mask=offs_m[:, None] < m and mask_n[None, :]) + +def test_generic_mask_2d(device): + if device == 'cpu': + triton.runtime.driver.set_active(CPUDriver()) + m = 6 + n = 16 + M = triton.next_power_of_2(m) + N = triton.next_power_of_2(n) + input = torch.arange(2, 2 + (m * n), device=device, dtype=torch.float32).reshape(m, n) + output = torch.full_like(input, -1) + mask_m = torch.tensor([1, 0, 1, 0, 1, 0], device=device, dtype=torch.bool) + + mask_n = torch.tensor([0, 1, 0, 1, 0, 1, 0, 1, + 1, 0, 1, 0, 1, 0, 1, 0], device=device, dtype=torch.bool) + + generic_mask_2d_kernel[1, 1, 1](input, output, mask_m.to(torch.int8), mask_n.to(torch.int8), m, n, M, N) + + v = torch.full_like(input, -2) + v[mask_m,:] = input[mask_m,:] + + expected_output = torch.full_like(output, -1) + expected_output[:,mask_n] = v[:,mask_n] + + torch.testing.assert_close(output, expected_output) + + +@triton.jit +def generic_mask_3d_kernel(in_ptr, out_ptr, mask_m_ptr, mask_n_ptr, b, m, n, stride_b, stride_m, stride_n, B: tl.constexpr, M:tl.constexpr, N:tl.constexpr): + offs_m = tl.arange(0, M) + offs_n = tl.arange(0, N) + offs_b = tl.arange(0, B) + + mask_m = tl.load(mask_m_ptr + offs_m, mask=offs_m < m, other=0) != 0 + mask_n = tl.load(mask_n_ptr + offs_n, mask=offs_n < n, other=0) != 0 + + in_ptrs = in_ptr + offs_b[:, None, None] * stride_b + offs_m[None, :, None] * stride_m + offs_n[None, None, :] * stride_n + v = tl.load(in_ptrs, mask= offs_b[:, None, None] < b and mask_m[None, :, None] and offs_n[None, None, :] < n, other=-2) + out_ptrs = out_ptr + offs_b[:, None, None] * stride_b + offs_m[None, :, None] * stride_m + offs_n[None, None, :] * stride_n + tl.store(out_ptrs, v, mask= offs_b[:, None, None] < b and offs_m[None, :, None] < m and mask_n[None, None, :]) + +def test_generic_mask_3d(device): + if device == 'cpu': + triton.runtime.driver.set_active(CPUDriver()) + b = 4 + m = 6 + n = 16 + B = triton.next_power_of_2(b) + M = triton.next_power_of_2(m) + N = triton.next_power_of_2(n) + input = torch.arange(2, 2 + (b * m * n), device=device, dtype=torch.float32).reshape(b, m, n) + output = torch.full_like(input, -1) + mask_m = torch.tensor([1, 0, 1, 0, 1, 0], device=device, dtype=torch.bool) + + mask_n = torch.tensor([0, 1, 0, 1, 0, 1, 0, 1, + 1, 0, 1, 0, 1, 0, 1, 0], device=device, dtype=torch.bool) + + stride_b = input.stride(0) + stride_m = input.stride(1) + stride_n = input.stride(2) + generic_mask_3d_kernel[1, 1, 1](input, output, mask_m.to(torch.int8), mask_n.to(torch.int8), b, m, n, + stride_b, stride_m, stride_n, + B, M, N) + + v = torch.full_like(input, -2) + v[:, mask_m,:] = input[:, mask_m,:] + + expected_output = torch.full_like(output, -1) + expected_output[:, :, mask_n] = v[:, :, mask_n] + + torch.testing.assert_close(output, expected_output) + +# non-continuous ld/st and (offs_n < n)[:, None] pattern. + +@triton.jit +def generic_mask_2d_non_continuous_load_kernel(in_ptr, out_ptr, index_m_ptr, im, I_M:tl.constexpr, N:tl.constexpr, m, n, stride_m, stride_n): + offs_m = tl.arange(0, I_M) + offs_n = tl.arange(0, N) + + index_m = tl.load(index_m_ptr + offs_m, mask=offs_m < im, other=0) + mask_m_i = index_m < m and offs_m < im + + + in_ptrs = in_ptr + index_m[:, None] * stride_m + offs_n[None, :] * stride_n + v = tl.load(in_ptrs, mask=mask_m_i[:, None] and offs_n[None, :] < n, other=-2) + + out_ptrs = out_ptr + offs_m[:, None] * stride_m + offs_n[None, :] * stride_n + tl.store(out_ptrs, v, mask=offs_m[:, None] < im and offs_n[None, :] < n) + +def test_generic_mask_2d_non_continuous_load(device): + if device == 'cpu': + triton.runtime.driver.set_active(CPUDriver()) + m = 6 + n = 16 + M = triton.next_power_of_2(m) + N = triton.next_power_of_2(n) + input = torch.arange(2, 2 + (m * n), device=device, dtype=torch.float32).reshape(m, n) + index_m = torch.tensor([1, 3, 7], device=device, dtype=torch.int32) + + index_n = torch.tensor([10, 20, 15], device=device, dtype=torch.int32) + + stride_m = input.stride(0) + stride_n = input.stride(1) + I_M = triton.next_power_of_2(len(index_m)) + I_N = triton.next_power_of_2(len(index_n)) + + output = torch.full((len(index_m), n), -1, device=device, dtype=torch.float32) + generic_mask_2d_non_continuous_load_kernel[1, 1, 1](input, output, index_m, len(index_m), I_M, N, m, n, stride_m, stride_n) + + expected_output = torch.full((len(index_m), n), -2, device=device, dtype=torch.float32) + mask_m = index_m < m + index_m = index_m[mask_m] + expected_output[:len(index_m),:] = input[index_m,:] + + torch.testing.assert_close(output, expected_output) + + +@triton.jit +def generic_mask_2d_non_continuous_store_kernel(in_ptr, out_ptr, index_n_ptr, i_n, I_N:tl.constexpr, M:tl.constexpr, m, n, stride_m, stride_n): + offs_m = tl.arange(0, M) + offs_n = tl.arange(0, I_N) + + index_n = tl.load(index_n_ptr + offs_n, mask=offs_n < n, other=0) + mask_n_i = index_n < n and offs_n < i_n + + + in_ptrs = in_ptr + offs_m[:, None] * stride_m + offs_n[None, :] * stride_n + v = tl.load(in_ptrs, mask=offs_m[:, None] < m and offs_n[None, :] < i_n, other=-2) + + out_ptrs = out_ptr + offs_m[:, None] * stride_m + index_n[None, :] * stride_n + tl.store(out_ptrs, v, mask=offs_m[:, None] < m and mask_n_i[None, :]) + +def test_generic_mask_2d_non_continuous_store(device): + if device == 'cpu': + triton.runtime.driver.set_active(CPUDriver()) + m = 6 + n = 16 + M = triton.next_power_of_2(m) + N = triton.next_power_of_2(n) + input = torch.arange(2, 2 + (m * n), device=device, dtype=torch.float32).reshape(m, n) + + index_n = torch.tensor([10, 20, 15, 1, 3, 5, 7, 0, 2, 4, 6, 8], device=device, dtype=torch.int32) + + stride_m = input.stride(0) + stride_n = input.stride(1) + + I_N = triton.next_power_of_2(len(index_n)) + + output = torch.full((m, n), -1, device=device, dtype=torch.float32) + generic_mask_2d_non_continuous_store_kernel[1, 1, 1](input, output, index_n, len(index_n), I_N, M, m, n, stride_m, stride_n) + + expected_output = torch.full((m, n), -1, device=device, dtype=torch.float32) + mask_n = index_n < n + + v = input[:,:len(index_n)] + v = v[:,mask_n] + index_n = index_n[mask_n] + expected_output[:,index_n] = v + + torch.testing.assert_close(output, expected_output) + + +@triton.jit +def generic_mask_3d_non_continuous_load_kernel(in_ptr, out_ptr, index_m_ptr, im, I_M:tl.constexpr, N:tl.constexpr, B:tl.constexpr, b, m, n, stride_b, stride_m, stride_n, o_stride_b, o_stride_m, o_stride_n): + offs_m = tl.arange(0, I_M) + offs_n = tl.arange(0, N) + offs_b = tl.arange(0, B) + + index_m = tl.load(index_m_ptr + offs_m, mask=offs_m < im, other=0) + mask_m_i = index_m < m and offs_m < im + + + in_ptrs = in_ptr + offs_b[:, None, None] * stride_b + index_m[None, :, None] * stride_m + offs_n[None, None, :] * stride_n + v = tl.load(in_ptrs, mask=offs_b[:, None, None] < b and mask_m_i[None, :, None] and offs_n[None, None, :] < n, other=-2) + + out_ptrs = out_ptr + offs_b[:, None, None] * o_stride_b + offs_m[None, :, None] * o_stride_m + offs_n[None, None, :] * o_stride_n + tl.store(out_ptrs, v, mask=offs_b[:, None, None] < b and offs_m[None, :, None] < im and offs_n[None, None, :] < n) + +def test_generic_mask_3d_non_continuous_load(device): + if device == 'cpu': + triton.runtime.driver.set_active(CPUDriver()) + b = 4 + m = 6 + n = 16 + B = triton.next_power_of_2(b) + M = triton.next_power_of_2(m) + N = triton.next_power_of_2(n) + input = torch.arange(2, 2 + (b * m * n), device=device, dtype=torch.float32).reshape(b, m, n) + index_m = torch.tensor([1, 3, 7], device=device, dtype=torch.int32) + + index_n = torch.tensor([10, 20, 15], device=device, dtype=torch.int32) + + stride_b = input.stride(0) + stride_m = input.stride(1) + stride_n = input.stride(2) + + I_M = triton.next_power_of_2(len(index_m)) + I_N = triton.next_power_of_2(len(index_n)) + + output = torch.full((b, len(index_m), n), -1, device=device, dtype=torch.float32) + + o_stride_b = output.stride(0) + o_stride_m = output.stride(1) + o_stride_n = output.stride(2) + + generic_mask_3d_non_continuous_load_kernel[1, 1, 1](input, output, index_m, len(index_m), I_M, N, B, b, m, n, + stride_b, stride_m, stride_n, + o_stride_b, o_stride_m, o_stride_n) + + expected_output = torch.full((b, len(index_m), n), -2, device=device, dtype=torch.float32) + mask_m = index_m < m + index_m = index_m[mask_m] + expected_output[:,:len(index_m),:] = input[:,index_m,:] + + print(output) + print(expected_output) + torch.testing.assert_close(output, expected_output) + + +@triton.jit +def generic_mask_3d_non_continuous_store_kernel(in_ptr, out_ptr, index_n_ptr, i_n, I_N:tl.constexpr, M:tl.constexpr, B:tl.constexpr, b, m, n, stride_b, stride_m, stride_n): + offs_m = tl.arange(0, M) + offs_n = tl.arange(0, I_N) + offs_b = tl.arange(0, B) + + index_n = tl.load(index_n_ptr + offs_n, mask=offs_n < n, other=0) + mask_n_i = index_n < n and offs_n < i_n + + + in_ptrs = in_ptr + offs_b[:, None, None] * stride_b + offs_m[None, :, None] * stride_m + offs_n[None, None, :] * stride_n + v = tl.load(in_ptrs, mask=offs_b[:, None, None] < b and offs_m[None, :, None] < m and offs_n[None, None, :] < i_n, other=-2) + + out_ptrs = out_ptr + offs_b[:, None, None] * stride_b + offs_m[None, :, None] * stride_m + index_n[None, None, :] * stride_n + tl.store(out_ptrs, v, mask=offs_b[:, None, None] < b and offs_m[None, :, None] < m and mask_n_i[None, None, :]) + +def test_generic_mask_3d_non_continuous_store(device): + if device == 'cpu': + triton.runtime.driver.set_active(CPUDriver()) + b = 4 + m = 6 + n = 16 + B = triton.next_power_of_2(b) + M = triton.next_power_of_2(m) + N = triton.next_power_of_2(n) + input = torch.arange(2, 2 + (b * m * n), device=device, dtype=torch.float32).reshape(b, m, n) + + index_n = torch.tensor([10, 20, 15, 1, 3, 5, 7, 0, 2, 4, 6, 8], device=device, dtype=torch.int32) + + stride_b = input.stride(0) + stride_m = input.stride(1) + stride_n = input.stride(2) + + I_N = triton.next_power_of_2(len(index_n)) + + output = torch.full((b, m, n), -1, device=device, dtype=torch.float32) + generic_mask_3d_non_continuous_store_kernel[1, 1, 1](input, output, index_n, len(index_n), I_N, M, B, b, m, n, stride_b, stride_m, stride_n) + + expected_output = torch.full((b, m, n), -1, device=device, dtype=torch.float32) + mask_n = index_n < n + + v = input[:,:,:len(index_n)] + v = v[:,:,mask_n] + index_n = index_n[mask_n] + expected_output[:,:,index_n] = v + + torch.testing.assert_close(output, expected_output) \ No newline at end of file diff --git a/test/Conversion/TritonToStructured/generic_mask_2d_kernel.mlir b/test/Conversion/TritonToStructured/generic_mask_2d_kernel.mlir new file mode 100644 index 000000000..e635736b5 --- /dev/null +++ b/test/Conversion/TritonToStructured/generic_mask_2d_kernel.mlir @@ -0,0 +1,95 @@ +// RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s + +// Make sure make_gather_scatter_tptr with generic mask generate correctly. + +// CHECK-LABEL: tt.func public @generic_mask_2d_kernel( +// CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_1:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_2:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_3:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_4:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32, +// CHECK-SAME: %[[VAL_5:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { +// CHECK: %[[VAL_6:.*]] = arith.constant -2.000000e+00 : f32 +// CHECK: %[[VAL_7:.*]] = arith.constant 16 : index +// CHECK: %[[VAL_8:.*]] = arith.constant 0 : i8 +// CHECK: %[[VAL_9:.*]] = arith.constant 0 : index +// CHECK: %[[VAL_10:.*]] = arith.constant 8 : index +// CHECK: %[[VAL_11:.*]] = arith.constant dense<0> : tensor<16xi32> +// CHECK: %[[VAL_12:.*]] = arith.constant dense<0> : tensor<8xi32> +// CHECK: %[[VAL_13:.*]] = tts.make_tptr %[[VAL_2]] to sizes: [8], strides: [1], offsets: [0], shape: [0], order: [] : to tensor<8x!tt.ptr> +// CHECK: %[[VAL_14:.*]] = arith.index_cast %[[VAL_4]] : i32 to index +// CHECK: %[[VAL_15:.*]] = arith.minsi %[[VAL_14]], %[[VAL_10]] : index +// CHECK: %[[VAL_16:.*]] = arith.maxsi %[[VAL_15]], %[[VAL_9]] : index +// CHECK: %[[VAL_17:.*]] = "tts.load"(%[[VAL_13]], %[[VAL_16]], %[[VAL_8]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (tensor<8x!tt.ptr>, index, i8) -> tensor<8xi8> +// CHECK: %[[VAL_18:.*]] = arith.extsi %[[VAL_17]] : tensor<8xi8> to tensor<8xi32> +// CHECK: %[[VAL_19:.*]] = arith.cmpi ne, %[[VAL_18]], %[[VAL_12]] : tensor<8xi32> +// CHECK: %[[VAL_20:.*]] = tts.make_tptr %[[VAL_3]] to sizes: [16], strides: [1], offsets: [0], shape: [0], order: [] : to tensor<16x!tt.ptr> +// CHECK: %[[VAL_21:.*]] = arith.index_cast %[[VAL_5]] : i32 to index +// CHECK: %[[VAL_22:.*]] = arith.minsi %[[VAL_21]], %[[VAL_7]] : index +// CHECK: %[[VAL_23:.*]] = arith.maxsi %[[VAL_22]], %[[VAL_9]] : index +// CHECK: %[[VAL_24:.*]] = "tts.load"(%[[VAL_20]], %[[VAL_23]], %[[VAL_8]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (tensor<16x!tt.ptr>, index, i8) -> tensor<16xi8> +// CHECK: %[[VAL_25:.*]] = arith.extsi %[[VAL_24]] : tensor<16xi8> to tensor<16xi32> +// CHECK: %[[VAL_26:.*]] = arith.cmpi ne, %[[VAL_25]], %[[VAL_11]] : tensor<16xi32> +// CHECK: %[[VAL_27:.*]] = arith.minsi %[[VAL_23]], %[[VAL_7]] : index +// CHECK: %[[VAL_28:.*]] = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> +// CHECK: %[[VAL_29:.*]] = tts.make_gather_scatter_tptr %[[VAL_0]] to sizes: [8, 16] gather_scatter_dim: 0 gather_scatter_offset: %[[VAL_28]] gather_scatter_mask: %[[VAL_19]], strides: {{\[}}%[[VAL_7]], 1], offsets: [0, 0] : tensor<8xi32> tensor<8xi1> to !tt.ptr> +// CHECK: %[[VAL_30:.*]] = "tts.load"(%[[VAL_29]], %[[VAL_27]], %[[VAL_6]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (!tt.ptr>, index, f32) -> tensor<8x16xf32> +// CHECK: %[[VAL_31:.*]] = arith.minsi %[[VAL_16]], %[[VAL_10]] : index +// CHECK: %[[VAL_32:.*]] = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> +// CHECK: %[[VAL_33:.*]] = tts.make_gather_scatter_tptr %[[VAL_1]] to sizes: [8, 16] gather_scatter_dim: 1 gather_scatter_offset: %[[VAL_32]] gather_scatter_mask: %[[VAL_26]], strides: {{\[}}%[[VAL_7]], 1], offsets: [0, 0] : tensor<16xi32> tensor<16xi1> to !tt.ptr> +// CHECK: "tts.store"(%[[VAL_33]], %[[VAL_30]], %[[VAL_31]]) <{static_mask_dims = array}> : (!tt.ptr>, tensor<8x16xf32>, index) -> () +// CHECK: tt.return + +module { + tt.func public @generic_mask_2d_kernel(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: !tt.ptr {tt.divisibility = 16 : i32}, %arg3: !tt.ptr {tt.divisibility = 16 : i32}, %arg4: i32, %arg5: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %cst = arith.constant dense<-2.000000e+00> : tensor<8x16xf32> + %cst_0 = arith.constant dense<0> : tensor<16xi8> + %cst_1 = arith.constant dense<0> : tensor<8xi8> + %cst_2 = arith.constant dense<16> : tensor<8x1xi32> + %cst_3 = arith.constant dense<0> : tensor<16xi32> + %cst_4 = arith.constant dense<0> : tensor<8xi32> + %0 = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> + %1 = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> + %2 = tt.splat %arg4 : i32 -> tensor<8xi32> + %3 = arith.cmpi slt, %0, %2 : tensor<8xi32> + %4 = tt.splat %arg2 : !tt.ptr -> tensor<8x!tt.ptr> + %5 = tt.addptr %4, %0 : tensor<8x!tt.ptr>, tensor<8xi32> + %6 = tt.load %5, %3, %cst_1 : tensor<8x!tt.ptr> + %7 = arith.extsi %6 : tensor<8xi8> to tensor<8xi32> + %8 = arith.cmpi ne, %7, %cst_4 : tensor<8xi32> + %9 = tt.splat %arg5 : i32 -> tensor<16xi32> + %10 = arith.cmpi slt, %1, %9 : tensor<16xi32> + %11 = tt.splat %arg3 : !tt.ptr -> tensor<16x!tt.ptr> + %12 = tt.addptr %11, %1 : tensor<16x!tt.ptr>, tensor<16xi32> + %13 = tt.load %12, %10, %cst_0 : tensor<16x!tt.ptr> + %14 = arith.extsi %13 : tensor<16xi8> to tensor<16xi32> + %15 = arith.cmpi ne, %14, %cst_3 : tensor<16xi32> + %16 = tt.expand_dims %0 {axis = 1 : i32} : tensor<8xi32> -> tensor<8x1xi32> + %17 = arith.muli %16, %cst_2 : tensor<8x1xi32> + %18 = tt.splat %arg0 : !tt.ptr -> tensor<8x1x!tt.ptr> + %19 = tt.addptr %18, %17 : tensor<8x1x!tt.ptr>, tensor<8x1xi32> + %20 = tt.expand_dims %1 {axis = 0 : i32} : tensor<16xi32> -> tensor<1x16xi32> + %21 = tt.broadcast %19 : tensor<8x1x!tt.ptr> -> tensor<8x16x!tt.ptr> + %22 = tt.broadcast %20 : tensor<1x16xi32> -> tensor<8x16xi32> + %23 = tt.addptr %21, %22 : tensor<8x16x!tt.ptr>, tensor<8x16xi32> + %24 = tt.expand_dims %8 {axis = 1 : i32} : tensor<8xi1> -> tensor<8x1xi1> + %25 = tt.splat %arg5 : i32 -> tensor<1x16xi32> + %26 = arith.cmpi slt, %20, %25 : tensor<1x16xi32> + %27 = tt.broadcast %24 : tensor<8x1xi1> -> tensor<8x16xi1> + %28 = tt.broadcast %26 : tensor<1x16xi1> -> tensor<8x16xi1> + %29 = arith.andi %27, %28 : tensor<8x16xi1> + %30 = tt.load %23, %29, %cst : tensor<8x16x!tt.ptr> + %31 = tt.splat %arg1 : !tt.ptr -> tensor<8x1x!tt.ptr> + %32 = tt.addptr %31, %17 : tensor<8x1x!tt.ptr>, tensor<8x1xi32> + %33 = tt.broadcast %32 : tensor<8x1x!tt.ptr> -> tensor<8x16x!tt.ptr> + %34 = tt.addptr %33, %22 : tensor<8x16x!tt.ptr>, tensor<8x16xi32> + %35 = tt.splat %arg4 : i32 -> tensor<8x1xi32> + %36 = arith.cmpi slt, %16, %35 : tensor<8x1xi32> + %37 = tt.expand_dims %15 {axis = 0 : i32} : tensor<16xi1> -> tensor<1x16xi1> + %38 = tt.broadcast %36 : tensor<8x1xi1> -> tensor<8x16xi1> + %39 = tt.broadcast %37 : tensor<1x16xi1> -> tensor<8x16xi1> + %40 = arith.andi %38, %39 : tensor<8x16xi1> + tt.store %34, %30, %40 : tensor<8x16x!tt.ptr> + tt.return + } +} \ No newline at end of file diff --git a/test/Conversion/TritonToStructured/generic_mask_2d_non_continuous_load_kernel.mlir b/test/Conversion/TritonToStructured/generic_mask_2d_non_continuous_load_kernel.mlir new file mode 100644 index 000000000..3eafac2b7 --- /dev/null +++ b/test/Conversion/TritonToStructured/generic_mask_2d_non_continuous_load_kernel.mlir @@ -0,0 +1,84 @@ +// RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s + +// Make sure make_gather_scatter_tptr with generic mask generate correctly. + +// CHECK-LABEL: tt.func public @generic_mask_2d_non_continuous_load_kernel( +// CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_1:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_2:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_3:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32, +// CHECK-SAME: %[[VAL_4:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32, +// CHECK-SAME: %[[VAL_5:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32 {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_6:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { +// CHECK: %[[VAL_7:.*]] = arith.constant -2.000000e+00 : f32 +// CHECK: %[[VAL_8:.*]] = arith.constant 16 : index +// CHECK: %[[VAL_9:.*]] = arith.constant 0 : i32 +// CHECK: %[[VAL_10:.*]] = arith.constant 0 : index +// CHECK: %[[VAL_11:.*]] = arith.constant 4 : index +// CHECK: %[[VAL_12:.*]] = tt.make_range {end = 4 : i32, start = 0 : i32} : tensor<4xi32> +// CHECK: %[[VAL_13:.*]] = tt.splat %[[VAL_3]] : i32 -> tensor<4xi32> +// CHECK: %[[VAL_14:.*]] = arith.cmpi slt, %[[VAL_12]], %[[VAL_13]] : tensor<4xi32> +// CHECK: %[[VAL_15:.*]] = tts.make_tptr %[[VAL_2]] to sizes: [4], strides: [1], offsets: [0], shape: [0], order: [] : to tensor<4x!tt.ptr> +// CHECK: %[[VAL_16:.*]] = arith.index_cast %[[VAL_3]] : i32 to index +// CHECK: %[[VAL_17:.*]] = arith.minsi %[[VAL_16]], %[[VAL_11]] : index +// CHECK: %[[VAL_18:.*]] = arith.maxsi %[[VAL_17]], %[[VAL_10]] : index +// CHECK: %[[VAL_19:.*]] = "tts.load"(%[[VAL_15]], %[[VAL_18]], %[[VAL_9]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (tensor<4x!tt.ptr>, index, i32) -> tensor<4xi32> +// CHECK: %[[VAL_20:.*]] = tt.splat %[[VAL_4]] : i32 -> tensor<4xi32> +// CHECK: %[[VAL_21:.*]] = arith.cmpi slt, %[[VAL_19]], %[[VAL_20]] : tensor<4xi32> +// CHECK: %[[VAL_22:.*]] = arith.andi %[[VAL_21]], %[[VAL_14]] : tensor<4xi1> +// CHECK: %[[VAL_23:.*]] = arith.index_cast %[[VAL_6]] : i32 to index +// CHECK: %[[VAL_24:.*]] = arith.index_cast %[[VAL_5]] : i32 to index +// CHECK: %[[VAL_25:.*]] = arith.minsi %[[VAL_24]], %[[VAL_8]] : index +// CHECK: %[[VAL_26:.*]] = arith.maxsi %[[VAL_25]], %[[VAL_10]] : index +// CHECK: %[[VAL_27:.*]] = arith.minsi %[[VAL_26]], %[[VAL_8]] : index +// CHECK: %[[VAL_28:.*]] = tts.make_gather_scatter_tptr %[[VAL_0]] to sizes: [4, 16] gather_scatter_dim: 0 gather_scatter_offset: %[[VAL_19]] gather_scatter_mask: %[[VAL_22]], strides: {{\[}}%[[VAL_23]], 1], offsets: [0, 0] : tensor<4xi32> tensor<4xi1> to !tt.ptr> +// CHECK: %[[VAL_29:.*]] = "tts.load"(%[[VAL_28]], %[[VAL_27]], %[[VAL_7]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (!tt.ptr>, index, f32) -> tensor<4x16xf32> +// CHECK: %[[VAL_30:.*]] = tts.make_tptr %[[VAL_1]] to sizes: [4, 16], strides: {{\[}}%[[VAL_23]], 1], offsets: [0, 0], shape: [0, 0], order: [] : to tensor<4x16x!tt.ptr> +// CHECK: %[[VAL_31:.*]] = arith.minsi %[[VAL_18]], %[[VAL_11]] : index +// CHECK: "tts.store"(%[[VAL_30]], %[[VAL_29]], %[[VAL_31]], %[[VAL_27]]) <{static_mask_dims = array}> : (tensor<4x16x!tt.ptr>, tensor<4x16xf32>, index, index) -> () +// CHECK: tt.return + +module { + tt.func public @generic_mask_2d_non_continuous_load_kernel(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: !tt.ptr {tt.divisibility = 16 : i32}, %arg3: i32, %arg4: i32, %arg5: i32 {tt.divisibility = 16 : i32}, %arg6: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %cst = arith.constant dense<-2.000000e+00> : tensor<4x16xf32> + %cst_0 = arith.constant dense<0> : tensor<4xi32> + %0 = tt.make_range {end = 4 : i32, start = 0 : i32} : tensor<4xi32> + %1 = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> + %2 = tt.splat %arg3 : i32 -> tensor<4xi32> + %3 = arith.cmpi slt, %0, %2 : tensor<4xi32> + %4 = tt.splat %arg2 : !tt.ptr -> tensor<4x!tt.ptr> + %5 = tt.addptr %4, %0 : tensor<4x!tt.ptr>, tensor<4xi32> + %6 = tt.load %5, %3, %cst_0 : tensor<4x!tt.ptr> + %7 = tt.splat %arg4 : i32 -> tensor<4xi32> + %8 = arith.cmpi slt, %6, %7 : tensor<4xi32> + %9 = arith.andi %8, %3 : tensor<4xi1> + %10 = tt.expand_dims %6 {axis = 1 : i32} : tensor<4xi32> -> tensor<4x1xi32> + %11 = tt.splat %arg6 : i32 -> tensor<4x1xi32> + %12 = arith.muli %10, %11 : tensor<4x1xi32> + %13 = tt.splat %arg0 : !tt.ptr -> tensor<4x1x!tt.ptr> + %14 = tt.addptr %13, %12 : tensor<4x1x!tt.ptr>, tensor<4x1xi32> + %15 = tt.expand_dims %1 {axis = 0 : i32} : tensor<16xi32> -> tensor<1x16xi32> + %16 = tt.broadcast %14 : tensor<4x1x!tt.ptr> -> tensor<4x16x!tt.ptr> + %17 = tt.broadcast %15 : tensor<1x16xi32> -> tensor<4x16xi32> + %18 = tt.addptr %16, %17 : tensor<4x16x!tt.ptr>, tensor<4x16xi32> + %19 = tt.expand_dims %9 {axis = 1 : i32} : tensor<4xi1> -> tensor<4x1xi1> + %20 = tt.splat %arg5 : i32 -> tensor<1x16xi32> + %21 = arith.cmpi slt, %15, %20 : tensor<1x16xi32> + %22 = tt.broadcast %19 : tensor<4x1xi1> -> tensor<4x16xi1> + %23 = tt.broadcast %21 : tensor<1x16xi1> -> tensor<4x16xi1> + %24 = arith.andi %22, %23 : tensor<4x16xi1> + %25 = tt.load %18, %24, %cst : tensor<4x16x!tt.ptr> + %26 = tt.expand_dims %0 {axis = 1 : i32} : tensor<4xi32> -> tensor<4x1xi32> + %27 = arith.muli %26, %11 : tensor<4x1xi32> + %28 = tt.splat %arg1 : !tt.ptr -> tensor<4x1x!tt.ptr> + %29 = tt.addptr %28, %27 : tensor<4x1x!tt.ptr>, tensor<4x1xi32> + %30 = tt.broadcast %29 : tensor<4x1x!tt.ptr> -> tensor<4x16x!tt.ptr> + %31 = tt.addptr %30, %17 : tensor<4x16x!tt.ptr>, tensor<4x16xi32> + %32 = tt.splat %arg3 : i32 -> tensor<4x1xi32> + %33 = arith.cmpi slt, %26, %32 : tensor<4x1xi32> + %34 = tt.broadcast %33 : tensor<4x1xi1> -> tensor<4x16xi1> + %35 = arith.andi %34, %23 : tensor<4x16xi1> + tt.store %31, %25, %35 : tensor<4x16x!tt.ptr> + tt.return + } +} diff --git a/test/Conversion/TritonToStructured/generic_mask_2d_non_continuous_store_kernel.mlir b/test/Conversion/TritonToStructured/generic_mask_2d_non_continuous_store_kernel.mlir new file mode 100644 index 000000000..a306e1668 --- /dev/null +++ b/test/Conversion/TritonToStructured/generic_mask_2d_non_continuous_store_kernel.mlir @@ -0,0 +1,88 @@ +// RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s + +// Make sure make_gather_scatter_tptr with generic mask generate correctly. + +// CHECK-LABEL: tt.func public @generic_mask_2d_non_continuous_store_kernel( +// CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_1:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_2:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_3:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32, +// CHECK-SAME: %[[VAL_4:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32, +// CHECK-SAME: %[[VAL_5:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32 {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_6:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { +// CHECK: %[[VAL_7:.*]] = arith.constant -2.000000e+00 : f32 +// CHECK: %[[VAL_8:.*]] = arith.constant 8 : index +// CHECK: %[[VAL_9:.*]] = arith.constant 0 : i32 +// CHECK: %[[VAL_10:.*]] = arith.constant 0 : index +// CHECK: %[[VAL_11:.*]] = arith.constant 16 : index +// CHECK: %[[VAL_12:.*]] = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> +// CHECK: %[[VAL_13:.*]] = tt.splat %[[VAL_5]] : i32 -> tensor<16xi32> +// CHECK: %[[VAL_14:.*]] = tts.make_tptr %[[VAL_2]] to sizes: [16], strides: [1], offsets: [0], shape: [0], order: [] : to tensor<16x!tt.ptr> +// CHECK: %[[VAL_15:.*]] = arith.index_cast %[[VAL_5]] : i32 to index +// CHECK: %[[VAL_16:.*]] = arith.minsi %[[VAL_15]], %[[VAL_11]] : index +// CHECK: %[[VAL_17:.*]] = arith.maxsi %[[VAL_16]], %[[VAL_10]] : index +// CHECK: %[[VAL_18:.*]] = "tts.load"(%[[VAL_14]], %[[VAL_17]], %[[VAL_9]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (tensor<16x!tt.ptr>, index, i32) -> tensor<16xi32> +// CHECK: %[[VAL_19:.*]] = arith.cmpi slt, %[[VAL_18]], %[[VAL_13]] : tensor<16xi32> +// CHECK: %[[VAL_20:.*]] = tt.splat %[[VAL_3]] : i32 -> tensor<16xi32> +// CHECK: %[[VAL_21:.*]] = arith.cmpi slt, %[[VAL_12]], %[[VAL_20]] : tensor<16xi32> +// CHECK: %[[VAL_22:.*]] = arith.andi %[[VAL_19]], %[[VAL_21]] : tensor<16xi1> +// CHECK: %[[VAL_23:.*]] = arith.index_cast %[[VAL_6]] : i32 to index +// CHECK: %[[VAL_24:.*]] = tts.make_tptr %[[VAL_0]] to sizes: [8, 16], strides: {{\[}}%[[VAL_23]], 1], offsets: [0, 0], shape: [0, 0], order: [] : to tensor<8x16x!tt.ptr> +// CHECK: %[[VAL_25:.*]] = arith.index_cast %[[VAL_4]] : i32 to index +// CHECK: %[[VAL_26:.*]] = arith.minsi %[[VAL_25]], %[[VAL_8]] : index +// CHECK: %[[VAL_27:.*]] = arith.maxsi %[[VAL_26]], %[[VAL_10]] : index +// CHECK: %[[VAL_28:.*]] = arith.index_cast %[[VAL_3]] : i32 to index +// CHECK: %[[VAL_29:.*]] = arith.minsi %[[VAL_28]], %[[VAL_11]] : index +// CHECK: %[[VAL_30:.*]] = arith.maxsi %[[VAL_29]], %[[VAL_10]] : index +// CHECK: %[[VAL_31:.*]] = arith.minsi %[[VAL_27]], %[[VAL_8]] : index +// CHECK: %[[VAL_32:.*]] = arith.minsi %[[VAL_30]], %[[VAL_11]] : index +// CHECK: %[[VAL_33:.*]] = "tts.load"(%[[VAL_24]], %[[VAL_31]], %[[VAL_32]], %[[VAL_7]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (tensor<8x16x!tt.ptr>, index, index, f32) -> tensor<8x16xf32> +// CHECK: %[[VAL_34:.*]] = tts.make_gather_scatter_tptr %[[VAL_1]] to sizes: [8, 16] gather_scatter_dim: 1 gather_scatter_offset: %[[VAL_18]] gather_scatter_mask: %[[VAL_22]], strides: {{\[}}%[[VAL_23]], 1], offsets: [0, 0] : tensor<16xi32> tensor<16xi1> to !tt.ptr> +// CHECK: "tts.store"(%[[VAL_34]], %[[VAL_33]], %[[VAL_31]]) <{static_mask_dims = array}> : (!tt.ptr>, tensor<8x16xf32>, index) -> () +// CHECK: tt.return + +module { + tt.func public @generic_mask_2d_non_continuous_store_kernel(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: !tt.ptr {tt.divisibility = 16 : i32}, %arg3: i32, %arg4: i32, %arg5: i32 {tt.divisibility = 16 : i32}, %arg6: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %cst = arith.constant dense<-2.000000e+00> : tensor<8x16xf32> + %cst_0 = arith.constant dense<0> : tensor<16xi32> + %0 = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> + %1 = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> + %2 = tt.splat %arg5 : i32 -> tensor<16xi32> + %3 = arith.cmpi slt, %1, %2 : tensor<16xi32> + %4 = tt.splat %arg2 : !tt.ptr -> tensor<16x!tt.ptr> + %5 = tt.addptr %4, %1 : tensor<16x!tt.ptr>, tensor<16xi32> + %6 = tt.load %5, %3, %cst_0 : tensor<16x!tt.ptr> + %7 = arith.cmpi slt, %6, %2 : tensor<16xi32> + %8 = tt.splat %arg3 : i32 -> tensor<16xi32> + %9 = arith.cmpi slt, %1, %8 : tensor<16xi32> + %10 = arith.andi %7, %9 : tensor<16xi1> + %11 = tt.expand_dims %0 {axis = 1 : i32} : tensor<8xi32> -> tensor<8x1xi32> + %12 = tt.splat %arg6 : i32 -> tensor<8x1xi32> + %13 = arith.muli %11, %12 : tensor<8x1xi32> + %14 = tt.splat %arg0 : !tt.ptr -> tensor<8x1x!tt.ptr> + %15 = tt.addptr %14, %13 : tensor<8x1x!tt.ptr>, tensor<8x1xi32> + %16 = tt.expand_dims %1 {axis = 0 : i32} : tensor<16xi32> -> tensor<1x16xi32> + %17 = tt.broadcast %15 : tensor<8x1x!tt.ptr> -> tensor<8x16x!tt.ptr> + %18 = tt.broadcast %16 : tensor<1x16xi32> -> tensor<8x16xi32> + %19 = tt.addptr %17, %18 : tensor<8x16x!tt.ptr>, tensor<8x16xi32> + %20 = tt.splat %arg4 : i32 -> tensor<8x1xi32> + %21 = arith.cmpi slt, %11, %20 : tensor<8x1xi32> + %22 = tt.splat %arg3 : i32 -> tensor<1x16xi32> + %23 = arith.cmpi slt, %16, %22 : tensor<1x16xi32> + %24 = tt.broadcast %21 : tensor<8x1xi1> -> tensor<8x16xi1> + %25 = tt.broadcast %23 : tensor<1x16xi1> -> tensor<8x16xi1> + %26 = arith.andi %24, %25 : tensor<8x16xi1> + %27 = tt.load %19, %26, %cst : tensor<8x16x!tt.ptr> + %28 = tt.splat %arg1 : !tt.ptr -> tensor<8x1x!tt.ptr> + %29 = tt.addptr %28, %13 : tensor<8x1x!tt.ptr>, tensor<8x1xi32> + %30 = tt.expand_dims %6 {axis = 0 : i32} : tensor<16xi32> -> tensor<1x16xi32> + %31 = tt.broadcast %29 : tensor<8x1x!tt.ptr> -> tensor<8x16x!tt.ptr> + %32 = tt.broadcast %30 : tensor<1x16xi32> -> tensor<8x16xi32> + %33 = tt.addptr %31, %32 : tensor<8x16x!tt.ptr>, tensor<8x16xi32> + %34 = tt.expand_dims %10 {axis = 0 : i32} : tensor<16xi1> -> tensor<1x16xi1> + %35 = tt.broadcast %34 : tensor<1x16xi1> -> tensor<8x16xi1> + %36 = arith.andi %24, %35 : tensor<8x16xi1> + tt.store %33, %27, %36 : tensor<8x16x!tt.ptr> + tt.return + } +} \ No newline at end of file diff --git a/test/Conversion/TritonToStructured/generic_mask_3d_kernel.mlir b/test/Conversion/TritonToStructured/generic_mask_3d_kernel.mlir new file mode 100644 index 000000000..ff8d80f4d --- /dev/null +++ b/test/Conversion/TritonToStructured/generic_mask_3d_kernel.mlir @@ -0,0 +1,128 @@ +// RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s + +// Make sure make_gather_scatter_tptr with generic mask generate correctly. + +// CHECK-LABEL: tt.func public @generic_mask_3d_kernel( +// CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_1:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_2:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_3:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_4:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32, +// CHECK-SAME: %[[VAL_5:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32, +// CHECK-SAME: %[[VAL_6:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32 {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_7:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32 {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_8:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { +// CHECK: %[[VAL_9:.*]] = arith.constant -2.000000e+00 : f32 +// CHECK: %[[VAL_10:.*]] = arith.constant 4 : index +// CHECK: %[[VAL_11:.*]] = arith.constant 16 : index +// CHECK: %[[VAL_12:.*]] = arith.constant 0 : i8 +// CHECK: %[[VAL_13:.*]] = arith.constant 0 : index +// CHECK: %[[VAL_14:.*]] = arith.constant 8 : index +// CHECK: %[[VAL_15:.*]] = arith.constant dense<0> : tensor<16xi32> +// CHECK: %[[VAL_16:.*]] = arith.constant dense<0> : tensor<8xi32> +// CHECK: %[[VAL_17:.*]] = tts.make_tptr %[[VAL_2]] to sizes: [8], strides: [1], offsets: [0], shape: [0], order: [] : to tensor<8x!tt.ptr> +// CHECK: %[[VAL_18:.*]] = arith.index_cast %[[VAL_5]] : i32 to index +// CHECK: %[[VAL_19:.*]] = arith.minsi %[[VAL_18]], %[[VAL_14]] : index +// CHECK: %[[VAL_20:.*]] = arith.maxsi %[[VAL_19]], %[[VAL_13]] : index +// CHECK: %[[VAL_21:.*]] = "tts.load"(%[[VAL_17]], %[[VAL_20]], %[[VAL_12]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (tensor<8x!tt.ptr>, index, i8) -> tensor<8xi8> +// CHECK: %[[VAL_22:.*]] = arith.extsi %[[VAL_21]] : tensor<8xi8> to tensor<8xi32> +// CHECK: %[[VAL_23:.*]] = arith.cmpi ne, %[[VAL_22]], %[[VAL_16]] : tensor<8xi32> +// CHECK: %[[VAL_24:.*]] = tts.make_tptr %[[VAL_3]] to sizes: [16], strides: [1], offsets: [0], shape: [0], order: [] : to tensor<16x!tt.ptr> +// CHECK: %[[VAL_25:.*]] = arith.index_cast %[[VAL_6]] : i32 to index +// CHECK: %[[VAL_26:.*]] = arith.minsi %[[VAL_25]], %[[VAL_11]] : index +// CHECK: %[[VAL_27:.*]] = arith.maxsi %[[VAL_26]], %[[VAL_13]] : index +// CHECK: %[[VAL_28:.*]] = "tts.load"(%[[VAL_24]], %[[VAL_27]], %[[VAL_12]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (tensor<16x!tt.ptr>, index, i8) -> tensor<16xi8> +// CHECK: %[[VAL_29:.*]] = arith.extsi %[[VAL_28]] : tensor<16xi8> to tensor<16xi32> +// CHECK: %[[VAL_30:.*]] = arith.cmpi ne, %[[VAL_29]], %[[VAL_15]] : tensor<16xi32> +// CHECK: %[[VAL_31:.*]] = arith.index_cast %[[VAL_7]] : i32 to index +// CHECK: %[[VAL_32:.*]] = arith.index_cast %[[VAL_8]] : i32 to index +// CHECK: %[[VAL_33:.*]] = arith.index_cast %[[VAL_4]] : i32 to index +// CHECK: %[[VAL_34:.*]] = arith.minsi %[[VAL_33]], %[[VAL_10]] : index +// CHECK: %[[VAL_35:.*]] = arith.maxsi %[[VAL_34]], %[[VAL_13]] : index +// CHECK: %[[VAL_36:.*]] = arith.minsi %[[VAL_27]], %[[VAL_11]] : index +// CHECK: %[[VAL_37:.*]] = arith.minsi %[[VAL_35]], %[[VAL_10]] : index +// CHECK: %[[VAL_38:.*]] = arith.minsi %[[VAL_36]], %[[VAL_11]] : index +// CHECK: %[[VAL_39:.*]] = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> +// CHECK: %[[VAL_40:.*]] = tts.make_gather_scatter_tptr %[[VAL_0]] to sizes: [4, 8, 16] gather_scatter_dim: 1 gather_scatter_offset: %[[VAL_39]] gather_scatter_mask: %[[VAL_23]], strides: {{\[}}%[[VAL_31]], %[[VAL_32]], 1], offsets: [0, 0, 0] : tensor<8xi32> tensor<8xi1> to !tt.ptr> +// CHECK: %[[VAL_41:.*]] = "tts.load"(%[[VAL_40]], %[[VAL_37]], %[[VAL_38]], %[[VAL_9]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (!tt.ptr>, index, index, f32) -> tensor<4x8x16xf32> +// CHECK: %[[VAL_42:.*]] = arith.minsi %[[VAL_20]], %[[VAL_14]] : index +// CHECK: %[[VAL_43:.*]] = arith.minsi %[[VAL_42]], %[[VAL_14]] : index +// CHECK: %[[VAL_44:.*]] = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> +// CHECK: %[[VAL_45:.*]] = tts.make_gather_scatter_tptr %[[VAL_1]] to sizes: [4, 8, 16] gather_scatter_dim: 2 gather_scatter_offset: %[[VAL_44]] gather_scatter_mask: %[[VAL_30]], strides: {{\[}}%[[VAL_31]], %[[VAL_32]], 1], offsets: [0, 0, 0] : tensor<16xi32> tensor<16xi1> to !tt.ptr> +// CHECK: "tts.store"(%[[VAL_45]], %[[VAL_41]], %[[VAL_37]], %[[VAL_43]]) <{static_mask_dims = array}> : (!tt.ptr>, tensor<4x8x16xf32>, index, index) -> () +// CHECK: tt.return + +module { + tt.func public @generic_mask_3d_kernel(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: !tt.ptr {tt.divisibility = 16 : i32}, %arg3: !tt.ptr {tt.divisibility = 16 : i32}, %arg4: i32, %arg5: i32, %arg6: i32 {tt.divisibility = 16 : i32}, %arg7: i32 {tt.divisibility = 16 : i32}, %arg8: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %cst = arith.constant dense<-2.000000e+00> : tensor<4x8x16xf32> + %cst_0 = arith.constant dense<0> : tensor<16xi8> + %cst_1 = arith.constant dense<0> : tensor<8xi8> + %cst_2 = arith.constant dense<0> : tensor<16xi32> + %cst_3 = arith.constant dense<0> : tensor<8xi32> + %0 = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> + %1 = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> + %2 = tt.make_range {end = 4 : i32, start = 0 : i32} : tensor<4xi32> + %3 = tt.splat %arg5 : i32 -> tensor<8xi32> + %4 = arith.cmpi slt, %0, %3 : tensor<8xi32> + %5 = tt.splat %arg2 : !tt.ptr -> tensor<8x!tt.ptr> + %6 = tt.addptr %5, %0 : tensor<8x!tt.ptr>, tensor<8xi32> + %7 = tt.load %6, %4, %cst_1 : tensor<8x!tt.ptr> + %8 = arith.extsi %7 : tensor<8xi8> to tensor<8xi32> + %9 = arith.cmpi ne, %8, %cst_3 : tensor<8xi32> + %10 = tt.splat %arg6 : i32 -> tensor<16xi32> + %11 = arith.cmpi slt, %1, %10 : tensor<16xi32> + %12 = tt.splat %arg3 : !tt.ptr -> tensor<16x!tt.ptr> + %13 = tt.addptr %12, %1 : tensor<16x!tt.ptr>, tensor<16xi32> + %14 = tt.load %13, %11, %cst_0 : tensor<16x!tt.ptr> + %15 = arith.extsi %14 : tensor<16xi8> to tensor<16xi32> + %16 = arith.cmpi ne, %15, %cst_2 : tensor<16xi32> + %17 = tt.expand_dims %2 {axis = 1 : i32} : tensor<4xi32> -> tensor<4x1xi32> + %18 = tt.expand_dims %17 {axis = 2 : i32} : tensor<4x1xi32> -> tensor<4x1x1xi32> + %19 = tt.splat %arg7 : i32 -> tensor<4x1x1xi32> + %20 = arith.muli %18, %19 : tensor<4x1x1xi32> + %21 = tt.splat %arg0 : !tt.ptr -> tensor<4x1x1x!tt.ptr> + %22 = tt.addptr %21, %20 : tensor<4x1x1x!tt.ptr>, tensor<4x1x1xi32> + %23 = tt.expand_dims %0 {axis = 0 : i32} : tensor<8xi32> -> tensor<1x8xi32> + %24 = tt.expand_dims %23 {axis = 2 : i32} : tensor<1x8xi32> -> tensor<1x8x1xi32> + %25 = tt.splat %arg8 : i32 -> tensor<1x8x1xi32> + %26 = arith.muli %24, %25 : tensor<1x8x1xi32> + %27 = tt.broadcast %22 : tensor<4x1x1x!tt.ptr> -> tensor<4x8x1x!tt.ptr> + %28 = tt.broadcast %26 : tensor<1x8x1xi32> -> tensor<4x8x1xi32> + %29 = tt.addptr %27, %28 : tensor<4x8x1x!tt.ptr>, tensor<4x8x1xi32> + %30 = tt.expand_dims %1 {axis = 0 : i32} : tensor<16xi32> -> tensor<1x16xi32> + %31 = tt.expand_dims %30 {axis = 1 : i32} : tensor<1x16xi32> -> tensor<1x1x16xi32> + %32 = tt.broadcast %29 : tensor<4x8x1x!tt.ptr> -> tensor<4x8x16x!tt.ptr> + %33 = tt.broadcast %31 : tensor<1x1x16xi32> -> tensor<4x8x16xi32> + %34 = tt.addptr %32, %33 : tensor<4x8x16x!tt.ptr>, tensor<4x8x16xi32> + %35 = tt.splat %arg4 : i32 -> tensor<4x1x1xi32> + %36 = arith.cmpi slt, %18, %35 : tensor<4x1x1xi32> + %37 = tt.expand_dims %9 {axis = 0 : i32} : tensor<8xi1> -> tensor<1x8xi1> + %38 = tt.expand_dims %37 {axis = 2 : i32} : tensor<1x8xi1> -> tensor<1x8x1xi1> + %39 = tt.splat %arg6 : i32 -> tensor<1x1x16xi32> + %40 = arith.cmpi slt, %31, %39 : tensor<1x1x16xi32> + %41 = tt.broadcast %38 : tensor<1x8x1xi1> -> tensor<1x8x16xi1> + %42 = tt.broadcast %40 : tensor<1x1x16xi1> -> tensor<1x8x16xi1> + %43 = arith.andi %41, %42 : tensor<1x8x16xi1> + %44 = tt.broadcast %36 : tensor<4x1x1xi1> -> tensor<4x8x16xi1> + %45 = tt.broadcast %43 : tensor<1x8x16xi1> -> tensor<4x8x16xi1> + %46 = arith.andi %44, %45 : tensor<4x8x16xi1> + %47 = tt.load %34, %46, %cst : tensor<4x8x16x!tt.ptr> + %48 = tt.splat %arg1 : !tt.ptr -> tensor<4x1x1x!tt.ptr> + %49 = tt.addptr %48, %20 : tensor<4x1x1x!tt.ptr>, tensor<4x1x1xi32> + %50 = tt.broadcast %49 : tensor<4x1x1x!tt.ptr> -> tensor<4x8x1x!tt.ptr> + %51 = tt.addptr %50, %28 : tensor<4x8x1x!tt.ptr>, tensor<4x8x1xi32> + %52 = tt.broadcast %51 : tensor<4x8x1x!tt.ptr> -> tensor<4x8x16x!tt.ptr> + %53 = tt.addptr %52, %33 : tensor<4x8x16x!tt.ptr>, tensor<4x8x16xi32> + %54 = tt.splat %arg5 : i32 -> tensor<1x8x1xi32> + %55 = arith.cmpi slt, %24, %54 : tensor<1x8x1xi32> + %56 = tt.expand_dims %16 {axis = 0 : i32} : tensor<16xi1> -> tensor<1x16xi1> + %57 = tt.expand_dims %56 {axis = 1 : i32} : tensor<1x16xi1> -> tensor<1x1x16xi1> + %58 = tt.broadcast %55 : tensor<1x8x1xi1> -> tensor<1x8x16xi1> + %59 = tt.broadcast %57 : tensor<1x1x16xi1> -> tensor<1x8x16xi1> + %60 = arith.andi %58, %59 : tensor<1x8x16xi1> + %61 = tt.broadcast %60 : tensor<1x8x16xi1> -> tensor<4x8x16xi1> + %62 = arith.andi %44, %61 : tensor<4x8x16xi1> + tt.store %53, %47, %62 : tensor<4x8x16x!tt.ptr> + tt.return + } +} \ No newline at end of file diff --git a/test/Conversion/TritonToStructured/generic_mask_3d_non_continuous_load_kernel.mlir b/test/Conversion/TritonToStructured/generic_mask_3d_non_continuous_load_kernel.mlir new file mode 100644 index 000000000..b2c6e8839 --- /dev/null +++ b/test/Conversion/TritonToStructured/generic_mask_3d_non_continuous_load_kernel.mlir @@ -0,0 +1,121 @@ +// RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s + +// Make sure make_gather_scatter_tptr with generic mask generate correctly. + +// CHECK-LABEL: tt.func public @generic_mask_3d_non_continuous_load_kernel( +// CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_1:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_2:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_3:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32, +// CHECK-SAME: %[[VAL_4:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32, +// CHECK-SAME: %[[VAL_5:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32, +// CHECK-SAME: %[[VAL_6:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32 {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_7:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32 {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_8:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32 {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_9:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32 {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_10:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { +// CHECK: %[[VAL_11:.*]] = arith.constant -2.000000e+00 : f32 +// CHECK: %[[VAL_12:.*]] = arith.constant 16 : index +// CHECK: %[[VAL_13:.*]] = arith.constant 0 : i32 +// CHECK: %[[VAL_14:.*]] = arith.constant 0 : index +// CHECK: %[[VAL_15:.*]] = arith.constant 4 : index +// CHECK: %[[VAL_16:.*]] = tt.make_range {end = 4 : i32, start = 0 : i32} : tensor<4xi32> +// CHECK: %[[VAL_17:.*]] = tt.splat %[[VAL_3]] : i32 -> tensor<4xi32> +// CHECK: %[[VAL_18:.*]] = arith.cmpi slt, %[[VAL_16]], %[[VAL_17]] : tensor<4xi32> +// CHECK: %[[VAL_19:.*]] = tts.make_tptr %[[VAL_2]] to sizes: [4], strides: [1], offsets: [0], shape: [0], order: [] : to tensor<4x!tt.ptr> +// CHECK: %[[VAL_20:.*]] = arith.index_cast %[[VAL_3]] : i32 to index +// CHECK: %[[VAL_21:.*]] = arith.minsi %[[VAL_20]], %[[VAL_15]] : index +// CHECK: %[[VAL_22:.*]] = arith.maxsi %[[VAL_21]], %[[VAL_14]] : index +// CHECK: %[[VAL_23:.*]] = "tts.load"(%[[VAL_19]], %[[VAL_22]], %[[VAL_13]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (tensor<4x!tt.ptr>, index, i32) -> tensor<4xi32> +// CHECK: %[[VAL_24:.*]] = tt.splat %[[VAL_5]] : i32 -> tensor<4xi32> +// CHECK: %[[VAL_25:.*]] = arith.cmpi slt, %[[VAL_23]], %[[VAL_24]] : tensor<4xi32> +// CHECK: %[[VAL_26:.*]] = arith.andi %[[VAL_25]], %[[VAL_18]] : tensor<4xi1> +// CHECK: %[[VAL_27:.*]] = arith.index_cast %[[VAL_7]] : i32 to index +// CHECK: %[[VAL_28:.*]] = arith.index_cast %[[VAL_8]] : i32 to index +// CHECK: %[[VAL_29:.*]] = arith.index_cast %[[VAL_4]] : i32 to index +// CHECK: %[[VAL_30:.*]] = arith.minsi %[[VAL_29]], %[[VAL_15]] : index +// CHECK: %[[VAL_31:.*]] = arith.maxsi %[[VAL_30]], %[[VAL_14]] : index +// CHECK: %[[VAL_32:.*]] = arith.index_cast %[[VAL_6]] : i32 to index +// CHECK: %[[VAL_33:.*]] = arith.minsi %[[VAL_32]], %[[VAL_12]] : index +// CHECK: %[[VAL_34:.*]] = arith.maxsi %[[VAL_33]], %[[VAL_14]] : index +// CHECK: %[[VAL_35:.*]] = arith.minsi %[[VAL_34]], %[[VAL_12]] : index +// CHECK: %[[VAL_36:.*]] = arith.minsi %[[VAL_31]], %[[VAL_15]] : index +// CHECK: %[[VAL_37:.*]] = arith.minsi %[[VAL_35]], %[[VAL_12]] : index +// CHECK: %[[VAL_38:.*]] = tts.make_gather_scatter_tptr %[[VAL_0]] to sizes: [4, 4, 16] gather_scatter_dim: 1 gather_scatter_offset: %[[VAL_23]] gather_scatter_mask: %[[VAL_26]], strides: {{\[}}%[[VAL_27]], %[[VAL_28]], 1], offsets: [0, 0, 0] : tensor<4xi32> tensor<4xi1> to !tt.ptr> +// CHECK: %[[VAL_39:.*]] = "tts.load"(%[[VAL_38]], %[[VAL_36]], %[[VAL_37]], %[[VAL_11]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (!tt.ptr>, index, index, f32) -> tensor<4x4x16xf32> +// CHECK: %[[VAL_40:.*]] = arith.index_cast %[[VAL_9]] : i32 to index +// CHECK: %[[VAL_41:.*]] = arith.index_cast %[[VAL_10]] : i32 to index +// CHECK: %[[VAL_42:.*]] = tts.make_tptr %[[VAL_1]] to sizes: [4, 4, 16], strides: {{\[}}%[[VAL_40]], %[[VAL_41]], 1], offsets: [0, 0, 0], shape: [0, 0, 0], order: [] : to tensor<4x4x16x!tt.ptr> +// CHECK: %[[VAL_43:.*]] = arith.minsi %[[VAL_22]], %[[VAL_15]] : index +// CHECK: %[[VAL_44:.*]] = arith.minsi %[[VAL_43]], %[[VAL_15]] : index +// CHECK: "tts.store"(%[[VAL_42]], %[[VAL_39]], %[[VAL_36]], %[[VAL_44]], %[[VAL_37]]) <{static_mask_dims = array}> : (tensor<4x4x16x!tt.ptr>, tensor<4x4x16xf32>, index, index, index) -> () +// CHECK: tt.return + +module { + tt.func public @generic_mask_3d_non_continuous_load_kernel(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: !tt.ptr {tt.divisibility = 16 : i32}, %arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32 {tt.divisibility = 16 : i32}, %arg7: i32 {tt.divisibility = 16 : i32}, %arg8: i32 {tt.divisibility = 16 : i32}, %arg9: i32 {tt.divisibility = 16 : i32}, %arg10: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %cst = arith.constant dense<-2.000000e+00> : tensor<4x4x16xf32> + %cst_0 = arith.constant dense<0> : tensor<4xi32> + %0 = tt.make_range {end = 4 : i32, start = 0 : i32} : tensor<4xi32> + %1 = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> + %2 = tt.splat %arg3 : i32 -> tensor<4xi32> + %3 = arith.cmpi slt, %0, %2 : tensor<4xi32> + %4 = tt.splat %arg2 : !tt.ptr -> tensor<4x!tt.ptr> + %5 = tt.addptr %4, %0 : tensor<4x!tt.ptr>, tensor<4xi32> + %6 = tt.load %5, %3, %cst_0 : tensor<4x!tt.ptr> + %7 = tt.splat %arg5 : i32 -> tensor<4xi32> + %8 = arith.cmpi slt, %6, %7 : tensor<4xi32> + %9 = arith.andi %8, %3 : tensor<4xi1> + %10 = tt.expand_dims %0 {axis = 1 : i32} : tensor<4xi32> -> tensor<4x1xi32> + %11 = tt.expand_dims %10 {axis = 2 : i32} : tensor<4x1xi32> -> tensor<4x1x1xi32> + %12 = tt.splat %arg7 : i32 -> tensor<4x1x1xi32> + %13 = arith.muli %11, %12 : tensor<4x1x1xi32> + %14 = tt.splat %arg0 : !tt.ptr -> tensor<4x1x1x!tt.ptr> + %15 = tt.addptr %14, %13 : tensor<4x1x1x!tt.ptr>, tensor<4x1x1xi32> + %16 = tt.expand_dims %6 {axis = 0 : i32} : tensor<4xi32> -> tensor<1x4xi32> + %17 = tt.expand_dims %16 {axis = 2 : i32} : tensor<1x4xi32> -> tensor<1x4x1xi32> + %18 = tt.splat %arg8 : i32 -> tensor<1x4x1xi32> + %19 = arith.muli %17, %18 : tensor<1x4x1xi32> + %20 = tt.broadcast %15 : tensor<4x1x1x!tt.ptr> -> tensor<4x4x1x!tt.ptr> + %21 = tt.broadcast %19 : tensor<1x4x1xi32> -> tensor<4x4x1xi32> + %22 = tt.addptr %20, %21 : tensor<4x4x1x!tt.ptr>, tensor<4x4x1xi32> + %23 = tt.expand_dims %1 {axis = 0 : i32} : tensor<16xi32> -> tensor<1x16xi32> + %24 = tt.expand_dims %23 {axis = 1 : i32} : tensor<1x16xi32> -> tensor<1x1x16xi32> + %25 = tt.broadcast %22 : tensor<4x4x1x!tt.ptr> -> tensor<4x4x16x!tt.ptr> + %26 = tt.broadcast %24 : tensor<1x1x16xi32> -> tensor<4x4x16xi32> + %27 = tt.addptr %25, %26 : tensor<4x4x16x!tt.ptr>, tensor<4x4x16xi32> + %28 = tt.splat %arg4 : i32 -> tensor<4x1x1xi32> + %29 = arith.cmpi slt, %11, %28 : tensor<4x1x1xi32> + %30 = tt.expand_dims %9 {axis = 0 : i32} : tensor<4xi1> -> tensor<1x4xi1> + %31 = tt.expand_dims %30 {axis = 2 : i32} : tensor<1x4xi1> -> tensor<1x4x1xi1> + %32 = tt.splat %arg6 : i32 -> tensor<1x1x16xi32> + %33 = arith.cmpi slt, %24, %32 : tensor<1x1x16xi32> + %34 = tt.broadcast %31 : tensor<1x4x1xi1> -> tensor<1x4x16xi1> + %35 = tt.broadcast %33 : tensor<1x1x16xi1> -> tensor<1x4x16xi1> + %36 = arith.andi %34, %35 : tensor<1x4x16xi1> + %37 = tt.broadcast %29 : tensor<4x1x1xi1> -> tensor<4x4x16xi1> + %38 = tt.broadcast %36 : tensor<1x4x16xi1> -> tensor<4x4x16xi1> + %39 = arith.andi %37, %38 : tensor<4x4x16xi1> + %40 = tt.load %27, %39, %cst : tensor<4x4x16x!tt.ptr> + %41 = tt.splat %arg9 : i32 -> tensor<4x1x1xi32> + %42 = arith.muli %11, %41 : tensor<4x1x1xi32> + %43 = tt.splat %arg1 : !tt.ptr -> tensor<4x1x1x!tt.ptr> + %44 = tt.addptr %43, %42 : tensor<4x1x1x!tt.ptr>, tensor<4x1x1xi32> + %45 = tt.expand_dims %0 {axis = 0 : i32} : tensor<4xi32> -> tensor<1x4xi32> + %46 = tt.expand_dims %45 {axis = 2 : i32} : tensor<1x4xi32> -> tensor<1x4x1xi32> + %47 = tt.splat %arg10 : i32 -> tensor<1x4x1xi32> + %48 = arith.muli %46, %47 : tensor<1x4x1xi32> + %49 = tt.broadcast %44 : tensor<4x1x1x!tt.ptr> -> tensor<4x4x1x!tt.ptr> + %50 = tt.broadcast %48 : tensor<1x4x1xi32> -> tensor<4x4x1xi32> + %51 = tt.addptr %49, %50 : tensor<4x4x1x!tt.ptr>, tensor<4x4x1xi32> + %52 = tt.broadcast %51 : tensor<4x4x1x!tt.ptr> -> tensor<4x4x16x!tt.ptr> + %53 = tt.addptr %52, %26 : tensor<4x4x16x!tt.ptr>, tensor<4x4x16xi32> + %54 = tt.splat %arg3 : i32 -> tensor<1x4x1xi32> + %55 = arith.cmpi slt, %46, %54 : tensor<1x4x1xi32> + %56 = tt.broadcast %55 : tensor<1x4x1xi1> -> tensor<1x4x16xi1> + %57 = arith.andi %56, %35 : tensor<1x4x16xi1> + %58 = tt.broadcast %57 : tensor<1x4x16xi1> -> tensor<4x4x16xi1> + %59 = arith.andi %37, %58 : tensor<4x4x16xi1> + tt.store %53, %40, %59 : tensor<4x4x16x!tt.ptr> + tt.return + } +} \ No newline at end of file diff --git a/test/Conversion/TritonToStructured/generic_mask_3d_non_continuous_store_kernel.mlir b/test/Conversion/TritonToStructured/generic_mask_3d_non_continuous_store_kernel.mlir new file mode 100644 index 000000000..681d1c8e9 --- /dev/null +++ b/test/Conversion/TritonToStructured/generic_mask_3d_non_continuous_store_kernel.mlir @@ -0,0 +1,119 @@ +// RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s + +// Make sure make_gather_scatter_tptr with generic mask generate correctly. + +// CHECK-LABEL: tt.func public @generic_mask_3d_non_continuous_store_kernel( +// CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_1:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_2:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_3:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32, +// CHECK-SAME: %[[VAL_4:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32, +// CHECK-SAME: %[[VAL_5:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32, +// CHECK-SAME: %[[VAL_6:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32 {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_7:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32 {tt.divisibility = 16 : i32}, +// CHECK-SAME: %[[VAL_8:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { +// CHECK: %[[VAL_9:.*]] = arith.constant -2.000000e+00 : f32 +// CHECK: %[[VAL_10:.*]] = arith.constant 8 : index +// CHECK: %[[VAL_11:.*]] = arith.constant 4 : index +// CHECK: %[[VAL_12:.*]] = arith.constant 0 : i32 +// CHECK: %[[VAL_13:.*]] = arith.constant 0 : index +// CHECK: %[[VAL_14:.*]] = arith.constant 16 : index +// CHECK: %[[VAL_15:.*]] = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> +// CHECK: %[[VAL_16:.*]] = tt.splat %[[VAL_6]] : i32 -> tensor<16xi32> +// CHECK: %[[VAL_17:.*]] = tts.make_tptr %[[VAL_2]] to sizes: [16], strides: [1], offsets: [0], shape: [0], order: [] : to tensor<16x!tt.ptr> +// CHECK: %[[VAL_18:.*]] = arith.index_cast %[[VAL_6]] : i32 to index +// CHECK: %[[VAL_19:.*]] = arith.minsi %[[VAL_18]], %[[VAL_14]] : index +// CHECK: %[[VAL_20:.*]] = arith.maxsi %[[VAL_19]], %[[VAL_13]] : index +// CHECK: %[[VAL_21:.*]] = "tts.load"(%[[VAL_17]], %[[VAL_20]], %[[VAL_12]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (tensor<16x!tt.ptr>, index, i32) -> tensor<16xi32> +// CHECK: %[[VAL_22:.*]] = arith.cmpi slt, %[[VAL_21]], %[[VAL_16]] : tensor<16xi32> +// CHECK: %[[VAL_23:.*]] = tt.splat %[[VAL_3]] : i32 -> tensor<16xi32> +// CHECK: %[[VAL_24:.*]] = arith.cmpi slt, %[[VAL_15]], %[[VAL_23]] : tensor<16xi32> +// CHECK: %[[VAL_25:.*]] = arith.andi %[[VAL_22]], %[[VAL_24]] : tensor<16xi1> +// CHECK: %[[VAL_26:.*]] = arith.index_cast %[[VAL_7]] : i32 to index +// CHECK: %[[VAL_27:.*]] = arith.index_cast %[[VAL_8]] : i32 to index +// CHECK: %[[VAL_28:.*]] = tts.make_tptr %[[VAL_0]] to sizes: [4, 8, 16], strides: {{\[}}%[[VAL_26]], %[[VAL_27]], 1], offsets: [0, 0, 0], shape: [0, 0, 0], order: [] : to tensor<4x8x16x!tt.ptr> +// CHECK: %[[VAL_29:.*]] = arith.index_cast %[[VAL_4]] : i32 to index +// CHECK: %[[VAL_30:.*]] = arith.minsi %[[VAL_29]], %[[VAL_11]] : index +// CHECK: %[[VAL_31:.*]] = arith.maxsi %[[VAL_30]], %[[VAL_13]] : index +// CHECK: %[[VAL_32:.*]] = arith.index_cast %[[VAL_5]] : i32 to index +// CHECK: %[[VAL_33:.*]] = arith.minsi %[[VAL_32]], %[[VAL_10]] : index +// CHECK: %[[VAL_34:.*]] = arith.maxsi %[[VAL_33]], %[[VAL_13]] : index +// CHECK: %[[VAL_35:.*]] = arith.index_cast %[[VAL_3]] : i32 to index +// CHECK: %[[VAL_36:.*]] = arith.minsi %[[VAL_35]], %[[VAL_14]] : index +// CHECK: %[[VAL_37:.*]] = arith.maxsi %[[VAL_36]], %[[VAL_13]] : index +// CHECK: %[[VAL_38:.*]] = arith.minsi %[[VAL_34]], %[[VAL_10]] : index +// CHECK: %[[VAL_39:.*]] = arith.minsi %[[VAL_37]], %[[VAL_14]] : index +// CHECK: %[[VAL_40:.*]] = arith.minsi %[[VAL_31]], %[[VAL_11]] : index +// CHECK: %[[VAL_41:.*]] = arith.minsi %[[VAL_38]], %[[VAL_10]] : index +// CHECK: %[[VAL_42:.*]] = arith.minsi %[[VAL_39]], %[[VAL_14]] : index +// CHECK: %[[VAL_43:.*]] = "tts.load"(%[[VAL_28]], %[[VAL_40]], %[[VAL_41]], %[[VAL_42]], %[[VAL_9]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (tensor<4x8x16x!tt.ptr>, index, index, index, f32) -> tensor<4x8x16xf32> +// CHECK: %[[VAL_44:.*]] = tts.make_gather_scatter_tptr %[[VAL_1]] to sizes: [4, 8, 16] gather_scatter_dim: 2 gather_scatter_offset: %[[VAL_21]] gather_scatter_mask: %[[VAL_25]], strides: {{\[}}%[[VAL_26]], %[[VAL_27]], 1], offsets: [0, 0, 0] : tensor<16xi32> tensor<16xi1> to !tt.ptr> +// CHECK: "tts.store"(%[[VAL_44]], %[[VAL_43]], %[[VAL_40]], %[[VAL_41]]) <{static_mask_dims = array}> : (!tt.ptr>, tensor<4x8x16xf32>, index, index) -> () +// CHECK: tt.return + +module { + tt.func public @generic_mask_3d_non_continuous_store_kernel(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: !tt.ptr {tt.divisibility = 16 : i32}, %arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32 {tt.divisibility = 16 : i32}, %arg7: i32 {tt.divisibility = 16 : i32}, %arg8: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %cst = arith.constant dense<-2.000000e+00> : tensor<4x8x16xf32> + %cst_0 = arith.constant dense<0> : tensor<16xi32> + %0 = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> + %1 = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> + %2 = tt.make_range {end = 4 : i32, start = 0 : i32} : tensor<4xi32> + %3 = tt.splat %arg6 : i32 -> tensor<16xi32> + %4 = arith.cmpi slt, %1, %3 : tensor<16xi32> + %5 = tt.splat %arg2 : !tt.ptr -> tensor<16x!tt.ptr> + %6 = tt.addptr %5, %1 : tensor<16x!tt.ptr>, tensor<16xi32> + %7 = tt.load %6, %4, %cst_0 : tensor<16x!tt.ptr> + %8 = arith.cmpi slt, %7, %3 : tensor<16xi32> + %9 = tt.splat %arg3 : i32 -> tensor<16xi32> + %10 = arith.cmpi slt, %1, %9 : tensor<16xi32> + %11 = arith.andi %8, %10 : tensor<16xi1> + %12 = tt.expand_dims %2 {axis = 1 : i32} : tensor<4xi32> -> tensor<4x1xi32> + %13 = tt.expand_dims %12 {axis = 2 : i32} : tensor<4x1xi32> -> tensor<4x1x1xi32> + %14 = tt.splat %arg7 : i32 -> tensor<4x1x1xi32> + %15 = arith.muli %13, %14 : tensor<4x1x1xi32> + %16 = tt.splat %arg0 : !tt.ptr -> tensor<4x1x1x!tt.ptr> + %17 = tt.addptr %16, %15 : tensor<4x1x1x!tt.ptr>, tensor<4x1x1xi32> + %18 = tt.expand_dims %0 {axis = 0 : i32} : tensor<8xi32> -> tensor<1x8xi32> + %19 = tt.expand_dims %18 {axis = 2 : i32} : tensor<1x8xi32> -> tensor<1x8x1xi32> + %20 = tt.splat %arg8 : i32 -> tensor<1x8x1xi32> + %21 = arith.muli %19, %20 : tensor<1x8x1xi32> + %22 = tt.broadcast %17 : tensor<4x1x1x!tt.ptr> -> tensor<4x8x1x!tt.ptr> + %23 = tt.broadcast %21 : tensor<1x8x1xi32> -> tensor<4x8x1xi32> + %24 = tt.addptr %22, %23 : tensor<4x8x1x!tt.ptr>, tensor<4x8x1xi32> + %25 = tt.expand_dims %1 {axis = 0 : i32} : tensor<16xi32> -> tensor<1x16xi32> + %26 = tt.expand_dims %25 {axis = 1 : i32} : tensor<1x16xi32> -> tensor<1x1x16xi32> + %27 = tt.broadcast %24 : tensor<4x8x1x!tt.ptr> -> tensor<4x8x16x!tt.ptr> + %28 = tt.broadcast %26 : tensor<1x1x16xi32> -> tensor<4x8x16xi32> + %29 = tt.addptr %27, %28 : tensor<4x8x16x!tt.ptr>, tensor<4x8x16xi32> + %30 = tt.splat %arg4 : i32 -> tensor<4x1x1xi32> + %31 = arith.cmpi slt, %13, %30 : tensor<4x1x1xi32> + %32 = tt.splat %arg5 : i32 -> tensor<1x8x1xi32> + %33 = arith.cmpi slt, %19, %32 : tensor<1x8x1xi32> + %34 = tt.splat %arg3 : i32 -> tensor<1x1x16xi32> + %35 = arith.cmpi slt, %26, %34 : tensor<1x1x16xi32> + %36 = tt.broadcast %33 : tensor<1x8x1xi1> -> tensor<1x8x16xi1> + %37 = tt.broadcast %35 : tensor<1x1x16xi1> -> tensor<1x8x16xi1> + %38 = arith.andi %36, %37 : tensor<1x8x16xi1> + %39 = tt.broadcast %31 : tensor<4x1x1xi1> -> tensor<4x8x16xi1> + %40 = tt.broadcast %38 : tensor<1x8x16xi1> -> tensor<4x8x16xi1> + %41 = arith.andi %39, %40 : tensor<4x8x16xi1> + %42 = tt.load %29, %41, %cst : tensor<4x8x16x!tt.ptr> + %43 = tt.splat %arg1 : !tt.ptr -> tensor<4x1x1x!tt.ptr> + %44 = tt.addptr %43, %15 : tensor<4x1x1x!tt.ptr>, tensor<4x1x1xi32> + %45 = tt.broadcast %44 : tensor<4x1x1x!tt.ptr> -> tensor<4x8x1x!tt.ptr> + %46 = tt.addptr %45, %23 : tensor<4x8x1x!tt.ptr>, tensor<4x8x1xi32> + %47 = tt.expand_dims %7 {axis = 0 : i32} : tensor<16xi32> -> tensor<1x16xi32> + %48 = tt.expand_dims %47 {axis = 1 : i32} : tensor<1x16xi32> -> tensor<1x1x16xi32> + %49 = tt.broadcast %46 : tensor<4x8x1x!tt.ptr> -> tensor<4x8x16x!tt.ptr> + %50 = tt.broadcast %48 : tensor<1x1x16xi32> -> tensor<4x8x16xi32> + %51 = tt.addptr %49, %50 : tensor<4x8x16x!tt.ptr>, tensor<4x8x16xi32> + %52 = tt.expand_dims %11 {axis = 0 : i32} : tensor<16xi1> -> tensor<1x16xi1> + %53 = tt.expand_dims %52 {axis = 1 : i32} : tensor<1x16xi1> -> tensor<1x1x16xi1> + %54 = tt.broadcast %53 : tensor<1x1x16xi1> -> tensor<1x8x16xi1> + %55 = arith.andi %36, %54 : tensor<1x8x16xi1> + %56 = tt.broadcast %55 : tensor<1x8x16xi1> -> tensor<4x8x16xi1> + %57 = arith.andi %39, %56 : tensor<4x8x16xi1> + tt.store %51, %42, %57 : tensor<4x8x16x!tt.ptr> + tt.return + } +} From 45ddfc06c44d4b0346461d1beb1c37a8d011221d Mon Sep 17 00:00:00 2001 From: Xiang Li Date: Fri, 29 Aug 2025 15:11:35 +0000 Subject: [PATCH 2/4] Update per comment. Also fix the offset divide stride issue when convert tts::MakeTensorPtrOp to tts::MakeGatherScatterTensorPtrOp. --- include/triton-shared/Analysis/MaskAnalysis.h | 5 +- lib/Analysis/MaskAnalysis.cpp | 91 +++++++++--- lib/AnalysisStructured/PtrAnalysis.cpp | 112 +++++++------- .../StructuredToMemref/StructuredToMemref.cpp | 16 +- ...eric_mask.py => test_unstructured_mask.py} | 139 +++++++++++++++--- ....mlir => unstructured_mask_2d_kernel.mlir} | 2 +- ...d_mask_2d_non_continuous_load_kernel.mlir} | 7 +- ..._mask_2d_non_continuous_store_kernel.mlir} | 2 +- ....mlir => unstructured_mask_3d_kernel.mlir} | 2 +- ...d_mask_3d_non_continuous_load_kernel.mlir} | 10 +- ..._mask_3d_non_continuous_store_kernel.mlir} | 3 +- 11 files changed, 270 insertions(+), 119 deletions(-) rename python/examples/{test_generic_mask.py => test_unstructured_mask.py} (59%) rename test/Conversion/TritonToStructured/{generic_mask_2d_kernel.mlir => unstructured_mask_2d_kernel.mlir} (98%) rename test/Conversion/TritonToStructured/{generic_mask_2d_non_continuous_load_kernel.mlir => unstructured_mask_2d_non_continuous_load_kernel.mlir} (91%) rename test/Conversion/TritonToStructured/{generic_mask_2d_non_continuous_store_kernel.mlir => unstructured_mask_2d_non_continuous_store_kernel.mlir} (98%) rename test/Conversion/TritonToStructured/{generic_mask_3d_kernel.mlir => unstructured_mask_3d_kernel.mlir} (99%) rename test/Conversion/TritonToStructured/{generic_mask_3d_non_continuous_load_kernel.mlir => unstructured_mask_3d_non_continuous_load_kernel.mlir} (91%) rename test/Conversion/TritonToStructured/{generic_mask_3d_non_continuous_store_kernel.mlir => unstructured_mask_3d_non_continuous_store_kernel.mlir} (99%) diff --git a/include/triton-shared/Analysis/MaskAnalysis.h b/include/triton-shared/Analysis/MaskAnalysis.h index 4faa1b660..e8dc43d17 100644 --- a/include/triton-shared/Analysis/MaskAnalysis.h +++ b/include/triton-shared/Analysis/MaskAnalysis.h @@ -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 +// 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 { @@ -69,7 +72,7 @@ struct MaskState { MaskState(bool useUnsafeMask = false) : useUnsafeMask(useUnsafeMask) {} - SmallVector> getGenericMasks(); + SmallVector> getUnstructuredMasks(); int64_t getRank() const { return dims.size(); } diff --git a/lib/Analysis/MaskAnalysis.cpp b/lib/Analysis/MaskAnalysis.cpp index 0f95dcc08..386a1cad3 100644 --- a/lib/Analysis/MaskAnalysis.cpp +++ b/lib/Analysis/MaskAnalysis.cpp @@ -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(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(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(loc, targetTensorType, 0, size) + .getResult(); + Value v = ofrToIndexValue(ofr, loc, builder); + v = builder + .create(loc, builder.getI32Type(), v) + .getResult(); + v = builder.create(loc, targetTensorType, v) + .getResult(); + return builder + .create(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(loc, lhsV, rhsV)); } } - // Only support one generic mask. - if (getGenericMasks().size() > 1) { + // Only support one unstructured mask. + if (getUnstructuredMasks().size() > 1) { return failure(); } } @@ -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; break; } @@ -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]}, @@ -433,10 +481,10 @@ LogicalResult MaskState::parseCmp(arith::CmpIOp cmpOp, const Location loc, SmallVector reassociation = *maybeReassociationMap; // Set masks. - genericMask = builder.create( + unstructuredMask = builder.create( loc, flatType, cmpOp, reassociation); } - masks[cmpOpDim] = genericMask; + masks[cmpOpDim] = unstructuredMask; } } else { cmpOpDim = 0; @@ -746,7 +794,7 @@ 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. @@ -754,15 +802,16 @@ LogicalResult MaskState::parseExpandDims(triton::ExpandDimsOp expandDimsOp, } } 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); @@ -777,7 +826,7 @@ LogicalResult MaskState::parseExpandDims(triton::ExpandDimsOp expandDimsOp, } // Return all non-nullptr masks along with their dimensions. -SmallVector> MaskState::getGenericMasks() { +SmallVector> MaskState::getUnstructuredMasks() { SmallVector> result; for (auto [i, m] : llvm::enumerate(masks)) { diff --git a/lib/AnalysisStructured/PtrAnalysis.cpp b/lib/AnalysisStructured/PtrAnalysis.cpp index cafc18f83..00b5718cd 100644 --- a/lib/AnalysisStructured/PtrAnalysis.cpp +++ b/lib/AnalysisStructured/PtrAnalysis.cpp @@ -39,24 +39,24 @@ 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> masks = mstate.getGenericMasks(); +// Try to apply unstructured mask on the ptr. +static Value applyUnstructuredMask(Operation *op, Value ptr, + triton::MaskState &mstate, Location loc, + OpBuilder builder) { + SmallVector> 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()) { 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; } @@ -64,45 +64,56 @@ static Value applyGenericMask(Operation *op, Value ptr, ptr = builder .create( loc, scatterPtr.getBase(), - scatterPtr.getGatherScatterOffset(), genericMask, + scatterPtr.getGatherScatterOffset(), unstructuredMask, scatterPtr.getGatherScatterDim(), scatterPtr.getSizes(), scatterPtr.getMixedStrides(), scatterPtr.getMixedOffsets()) .getResult(); - } else if (auto tptr = ptr.getDefiningOp()) { - OpFoldResult offsetFold = tptr.getMixedOffsets()[dim]; - Value offset = dyn_cast(offsetFold); - if (!offset) { - offset = builder - .create( - loc, cast(cast(offsetFold))) - .getResult(); - } - // Cast to integer for splat and makerange. - if (isa(offset.getType())) { - offset = - builder.create(loc, builder.getI32Type(), offset) - .getResult(); - } else if (offset.getType().isInteger(64)) { - offset = - builder.create(loc, builder.getI32Type(), offset) - .getResult(); - } - auto offsetRowType = - RankedTensorType::get({tptr.getSizes()[dim]}, offset.getType()); - Value scatterOffset = + } else if (auto structuredPtr = ptr.getDefiningOp()) { + auto ofrToI32Value = [&](OpFoldResult ofr) { + Value v = dyn_cast(ofr); + if (!v) { + v = builder + .create( + loc, cast(cast(ofr))) + .getResult(); + } + if (isa(v.getType())) { + v = builder.create(loc, builder.getI32Type(), v) + .getResult(); + } else if (v.getType().isInteger(64)) { + v = builder.create(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(loc, offset, stride); + + Value gatherScatterOffset = builder.create(loc, offsetRowType, offset).getResult(); Value range = builder - .create(loc, offsetRowType, 0, - tptr.getSizes()[dim]) + .create( + loc, offsetRowType, 0, structuredPtr.getSizes()[dim]) .getResult(); - scatterOffset = builder.create(loc, scatterOffset, range); - ptr = - builder - .create( - loc, tptr.getBase(), scatterOffset, genericMask, dim, - tptr.getSizes(), tptr.getMixedStrides(), tptr.getMixedOffsets()) - .getResult(); + gatherScatterOffset = + builder.create(loc, gatherScatterOffset, range); + ptr = builder + .create( + loc, structuredPtr.getBase(), gatherScatterOffset, + unstructuredMask, dim, structuredPtr.getSizes(), + structuredPtr.getMixedStrides(), + structuredPtr.getMixedOffsets()) + .getResult(); } else { return nullptr; } @@ -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 { @@ -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 @@ -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(); } @@ -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(); } diff --git a/lib/Conversion/StructuredToMemref/StructuredToMemref.cpp b/lib/Conversion/StructuredToMemref/StructuredToMemref.cpp index afd7e67ad..ae201a884 100644 --- a/lib/Conversion/StructuredToMemref/StructuredToMemref.cpp +++ b/lib/Conversion/StructuredToMemref/StructuredToMemref.cpp @@ -886,8 +886,8 @@ struct LoadConverter : public OpConversionPattern { // 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); @@ -915,11 +915,11 @@ struct LoadConverter : public OpConversionPattern { 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( - loc, genericMask, ValueRange{inductionVar}); + loc, unstructuredMask, ValueRange{inductionVar}); auto ifOp = rewriter.create(loc, maskValue); rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front()); } @@ -1048,8 +1048,8 @@ struct StoreConverter : public OpConversionPattern { // 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); @@ -1072,11 +1072,11 @@ struct StoreConverter : public OpConversionPattern { 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( - loc, genericMask, ValueRange{inductionVar}); + loc, unstructuredMask, ValueRange{inductionVar}); auto ifOp = rewriter.create(loc, maskValue); rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front()); } diff --git a/python/examples/test_generic_mask.py b/python/examples/test_unstructured_mask.py similarity index 59% rename from python/examples/test_generic_mask.py rename to python/examples/test_unstructured_mask.py index 35ef65c29..2b37c2a36 100644 --- a/python/examples/test_generic_mask.py +++ b/python/examples/test_unstructured_mask.py @@ -4,9 +4,27 @@ import triton.language as tl from triton.backends.triton_shared.driver import CPUDriver +from triton.backends.triton_shared.compiler import _get_triton_shared_opt_path + +import os +from pathlib import Path +import subprocess +import tempfile + +def run_triton_to_structured(ttir_code): + with tempfile.TemporaryDirectory() as tmpdir: + src_path = os.path.join(tmpdir, "tt.mlir") + dst_path = os.path.join(tmpdir, "ttshared.mlir") + Path(src_path).write_text(ttir_code) + triton_shared_opt_path = _get_triton_shared_opt_path() + + subprocess_args = [triton_shared_opt_path, src_path, "--triton-to-structured", "--remove-dead-values", "-o", dst_path] + + subprocess.check_call(subprocess_args) + return Path(dst_path).read_text() @triton.jit -def generic_mask_2d_kernel(in_ptr, out_ptr, mask_m_ptr, mask_n_ptr, m, n, M:tl.constexpr, N:tl.constexpr): +def unstructured_mask_2d_kernel(in_ptr, out_ptr, mask_m_ptr, mask_n_ptr, m, n, M:tl.constexpr, N:tl.constexpr): offs_m = tl.arange(0, M) offs_n = tl.arange(0, N) @@ -14,11 +32,13 @@ def generic_mask_2d_kernel(in_ptr, out_ptr, mask_m_ptr, mask_n_ptr, m, n, M:tl.c mask_n = tl.load(mask_n_ptr + offs_n, mask=offs_n < n, other=0) != 0 in_ptrs = in_ptr + offs_m[:, None] * N + offs_n[None, :] + # dim 0 with unstructured mask. v = tl.load(in_ptrs, mask=mask_m[:, None] and offs_n[None, :] < n, other=-2) out_ptrs = out_ptr + offs_m[:, None] * N + offs_n[None, :] + # dim 1 with unstructured mask. tl.store(out_ptrs, v, mask=offs_m[:, None] < m and mask_n[None, :]) -def test_generic_mask_2d(device): +def test_unstructured_mask_2d(device): if device == 'cpu': triton.runtime.driver.set_active(CPUDriver()) m = 6 @@ -32,7 +52,9 @@ def test_generic_mask_2d(device): mask_n = torch.tensor([0, 1, 0, 1, 0, 1, 0, 1, 1, 0, 1, 0, 1, 0, 1, 0], device=device, dtype=torch.bool) - generic_mask_2d_kernel[1, 1, 1](input, output, mask_m.to(torch.int8), mask_n.to(torch.int8), m, n, M, N) + res = unstructured_mask_2d_kernel[1,](input, output, mask_m.to(torch.int8), mask_n.to(torch.int8), m, n, M, N) + tts_code = run_triton_to_structured(res.asm['ttir']) + assert tts_code.count('tts.make_gather_scatter_tptr') == 2, "Expected exactly two calls to tts.make_gather_scatter_tptr" v = torch.full_like(input, -2) v[mask_m,:] = input[mask_m,:] @@ -44,7 +66,7 @@ def test_generic_mask_2d(device): @triton.jit -def generic_mask_3d_kernel(in_ptr, out_ptr, mask_m_ptr, mask_n_ptr, b, m, n, stride_b, stride_m, stride_n, B: tl.constexpr, M:tl.constexpr, N:tl.constexpr): +def unstructured_mask_3d_kernel(in_ptr, out_ptr, mask_m_ptr, mask_n_ptr, b, m, n, stride_b, stride_m, stride_n, B: tl.constexpr, M:tl.constexpr, N:tl.constexpr): offs_m = tl.arange(0, M) offs_n = tl.arange(0, N) offs_b = tl.arange(0, B) @@ -53,11 +75,13 @@ def generic_mask_3d_kernel(in_ptr, out_ptr, mask_m_ptr, mask_n_ptr, b, m, n, str mask_n = tl.load(mask_n_ptr + offs_n, mask=offs_n < n, other=0) != 0 in_ptrs = in_ptr + offs_b[:, None, None] * stride_b + offs_m[None, :, None] * stride_m + offs_n[None, None, :] * stride_n + # dim 1 with unstructured mask. v = tl.load(in_ptrs, mask= offs_b[:, None, None] < b and mask_m[None, :, None] and offs_n[None, None, :] < n, other=-2) out_ptrs = out_ptr + offs_b[:, None, None] * stride_b + offs_m[None, :, None] * stride_m + offs_n[None, None, :] * stride_n + # dim 2 with unstructured mask. tl.store(out_ptrs, v, mask= offs_b[:, None, None] < b and offs_m[None, :, None] < m and mask_n[None, None, :]) -def test_generic_mask_3d(device): +def test_unstructured_mask_3d(device): if device == 'cpu': triton.runtime.driver.set_active(CPUDriver()) b = 4 @@ -76,9 +100,11 @@ def test_generic_mask_3d(device): stride_b = input.stride(0) stride_m = input.stride(1) stride_n = input.stride(2) - generic_mask_3d_kernel[1, 1, 1](input, output, mask_m.to(torch.int8), mask_n.to(torch.int8), b, m, n, + res = unstructured_mask_3d_kernel[1,](input, output, mask_m.to(torch.int8), mask_n.to(torch.int8), b, m, n, stride_b, stride_m, stride_n, B, M, N) + tts_code = run_triton_to_structured(res.asm['ttir']) + assert tts_code.count('tts.make_gather_scatter_tptr') == 2, "Expected exactly two calls to tts.make_gather_scatter_tptr" v = torch.full_like(input, -2) v[:, mask_m,:] = input[:, mask_m,:] @@ -91,7 +117,7 @@ def test_generic_mask_3d(device): # non-continuous ld/st and (offs_n < n)[:, None] pattern. @triton.jit -def generic_mask_2d_non_continuous_load_kernel(in_ptr, out_ptr, index_m_ptr, im, I_M:tl.constexpr, N:tl.constexpr, m, n, stride_m, stride_n): +def unstructured_mask_2d_non_continuous_load_kernel(in_ptr, out_ptr, index_m_ptr, im, I_M:tl.constexpr, N:tl.constexpr, m, n, stride_m, stride_n): offs_m = tl.arange(0, I_M) offs_n = tl.arange(0, N) @@ -99,13 +125,15 @@ def generic_mask_2d_non_continuous_load_kernel(in_ptr, out_ptr, index_m_ptr, im, mask_m_i = index_m < m and offs_m < im + # dim 0 with unstructured offset. in_ptrs = in_ptr + index_m[:, None] * stride_m + offs_n[None, :] * stride_n + # dim 0 with unstructured mask. v = tl.load(in_ptrs, mask=mask_m_i[:, None] and offs_n[None, :] < n, other=-2) out_ptrs = out_ptr + offs_m[:, None] * stride_m + offs_n[None, :] * stride_n tl.store(out_ptrs, v, mask=offs_m[:, None] < im and offs_n[None, :] < n) -def test_generic_mask_2d_non_continuous_load(device): +def test_unstructured_mask_2d_non_continuous_load(device): if device == 'cpu': triton.runtime.driver.set_active(CPUDriver()) m = 6 @@ -123,7 +151,9 @@ def test_generic_mask_2d_non_continuous_load(device): I_N = triton.next_power_of_2(len(index_n)) output = torch.full((len(index_m), n), -1, device=device, dtype=torch.float32) - generic_mask_2d_non_continuous_load_kernel[1, 1, 1](input, output, index_m, len(index_m), I_M, N, m, n, stride_m, stride_n) + res = unstructured_mask_2d_non_continuous_load_kernel[1,](input, output, index_m, len(index_m), I_M, N, m, n, stride_m, stride_n) + tts_code = run_triton_to_structured(res.asm['ttir']) + assert tts_code.count('tts.make_gather_scatter_tptr') == 1, "Expected exactly one call to tts.make_gather_scatter_tptr" expected_output = torch.full((len(index_m), n), -2, device=device, dtype=torch.float32) mask_m = index_m < m @@ -134,7 +164,7 @@ def test_generic_mask_2d_non_continuous_load(device): @triton.jit -def generic_mask_2d_non_continuous_store_kernel(in_ptr, out_ptr, index_n_ptr, i_n, I_N:tl.constexpr, M:tl.constexpr, m, n, stride_m, stride_n): +def unstructured_mask_2d_non_continuous_store_kernel(in_ptr, out_ptr, index_n_ptr, i_n, I_N:tl.constexpr, M:tl.constexpr, m, n, stride_m, stride_n): offs_m = tl.arange(0, M) offs_n = tl.arange(0, I_N) @@ -144,11 +174,12 @@ def generic_mask_2d_non_continuous_store_kernel(in_ptr, out_ptr, index_n_ptr, i_ in_ptrs = in_ptr + offs_m[:, None] * stride_m + offs_n[None, :] * stride_n v = tl.load(in_ptrs, mask=offs_m[:, None] < m and offs_n[None, :] < i_n, other=-2) - + # dim 1 with unstructured offset. out_ptrs = out_ptr + offs_m[:, None] * stride_m + index_n[None, :] * stride_n + # dim 1 with unstructured mask. tl.store(out_ptrs, v, mask=offs_m[:, None] < m and mask_n_i[None, :]) -def test_generic_mask_2d_non_continuous_store(device): +def test_unstructured_mask_2d_non_continuous_store(device): if device == 'cpu': triton.runtime.driver.set_active(CPUDriver()) m = 6 @@ -165,7 +196,9 @@ def test_generic_mask_2d_non_continuous_store(device): I_N = triton.next_power_of_2(len(index_n)) output = torch.full((m, n), -1, device=device, dtype=torch.float32) - generic_mask_2d_non_continuous_store_kernel[1, 1, 1](input, output, index_n, len(index_n), I_N, M, m, n, stride_m, stride_n) + res = unstructured_mask_2d_non_continuous_store_kernel[1,](input, output, index_n, len(index_n), I_N, M, m, n, stride_m, stride_n) + tts_code = run_triton_to_structured(res.asm['ttir']) + assert tts_code.count('tts.make_gather_scatter_tptr') == 1, "Expected exactly one call to tts.make_gather_scatter_tptr" expected_output = torch.full((m, n), -1, device=device, dtype=torch.float32) mask_n = index_n < n @@ -179,7 +212,7 @@ def test_generic_mask_2d_non_continuous_store(device): @triton.jit -def generic_mask_3d_non_continuous_load_kernel(in_ptr, out_ptr, index_m_ptr, im, I_M:tl.constexpr, N:tl.constexpr, B:tl.constexpr, b, m, n, stride_b, stride_m, stride_n, o_stride_b, o_stride_m, o_stride_n): +def unstructured_mask_3d_non_continuous_load_kernel(in_ptr, out_ptr, index_m_ptr, im, I_M:tl.constexpr, N:tl.constexpr, B:tl.constexpr, b, m, n, stride_b, stride_m, stride_n, o_stride_b, o_stride_m, o_stride_n): offs_m = tl.arange(0, I_M) offs_n = tl.arange(0, N) offs_b = tl.arange(0, B) @@ -187,14 +220,15 @@ def generic_mask_3d_non_continuous_load_kernel(in_ptr, out_ptr, index_m_ptr, im, index_m = tl.load(index_m_ptr + offs_m, mask=offs_m < im, other=0) mask_m_i = index_m < m and offs_m < im - + # dim 1 with unstructured offset. in_ptrs = in_ptr + offs_b[:, None, None] * stride_b + index_m[None, :, None] * stride_m + offs_n[None, None, :] * stride_n + # dim 1 with unstructured mask. v = tl.load(in_ptrs, mask=offs_b[:, None, None] < b and mask_m_i[None, :, None] and offs_n[None, None, :] < n, other=-2) out_ptrs = out_ptr + offs_b[:, None, None] * o_stride_b + offs_m[None, :, None] * o_stride_m + offs_n[None, None, :] * o_stride_n tl.store(out_ptrs, v, mask=offs_b[:, None, None] < b and offs_m[None, :, None] < im and offs_n[None, None, :] < n) -def test_generic_mask_3d_non_continuous_load(device): +def test_unstructured_mask_3d_non_continuous_load(device): if device == 'cpu': triton.runtime.driver.set_active(CPUDriver()) b = 4 @@ -221,9 +255,11 @@ def test_generic_mask_3d_non_continuous_load(device): o_stride_m = output.stride(1) o_stride_n = output.stride(2) - generic_mask_3d_non_continuous_load_kernel[1, 1, 1](input, output, index_m, len(index_m), I_M, N, B, b, m, n, + res = unstructured_mask_3d_non_continuous_load_kernel[1,](input, output, index_m, len(index_m), I_M, N, B, b, m, n, stride_b, stride_m, stride_n, o_stride_b, o_stride_m, o_stride_n) + tts_code = run_triton_to_structured(res.asm['ttir']) + assert tts_code.count('tts.make_gather_scatter_tptr') == 1, "Expected exactly one call to tts.make_gather_scatter_tptr" expected_output = torch.full((b, len(index_m), n), -2, device=device, dtype=torch.float32) mask_m = index_m < m @@ -236,7 +272,7 @@ def test_generic_mask_3d_non_continuous_load(device): @triton.jit -def generic_mask_3d_non_continuous_store_kernel(in_ptr, out_ptr, index_n_ptr, i_n, I_N:tl.constexpr, M:tl.constexpr, B:tl.constexpr, b, m, n, stride_b, stride_m, stride_n): +def unstructured_mask_3d_non_continuous_store_kernel(in_ptr, out_ptr, index_n_ptr, i_n, I_N:tl.constexpr, M:tl.constexpr, B:tl.constexpr, b, m, n, stride_b, stride_m, stride_n): offs_m = tl.arange(0, M) offs_n = tl.arange(0, I_N) offs_b = tl.arange(0, B) @@ -248,10 +284,12 @@ def generic_mask_3d_non_continuous_store_kernel(in_ptr, out_ptr, index_n_ptr, i_ in_ptrs = in_ptr + offs_b[:, None, None] * stride_b + offs_m[None, :, None] * stride_m + offs_n[None, None, :] * stride_n v = tl.load(in_ptrs, mask=offs_b[:, None, None] < b and offs_m[None, :, None] < m and offs_n[None, None, :] < i_n, other=-2) + # dim 2 with unstructured offset. out_ptrs = out_ptr + offs_b[:, None, None] * stride_b + offs_m[None, :, None] * stride_m + index_n[None, None, :] * stride_n + # dim 2 with unstructured mask. tl.store(out_ptrs, v, mask=offs_b[:, None, None] < b and offs_m[None, :, None] < m and mask_n_i[None, None, :]) -def test_generic_mask_3d_non_continuous_store(device): +def test_unstructured_mask_3d_non_continuous_store(device): if device == 'cpu': triton.runtime.driver.set_active(CPUDriver()) b = 4 @@ -271,7 +309,9 @@ def test_generic_mask_3d_non_continuous_store(device): I_N = triton.next_power_of_2(len(index_n)) output = torch.full((b, m, n), -1, device=device, dtype=torch.float32) - generic_mask_3d_non_continuous_store_kernel[1, 1, 1](input, output, index_n, len(index_n), I_N, M, B, b, m, n, stride_b, stride_m, stride_n) + res = unstructured_mask_3d_non_continuous_store_kernel[1,](input, output, index_n, len(index_n), I_N, M, B, b, m, n, stride_b, stride_m, stride_n) + tts_code = run_triton_to_structured(res.asm['ttir']) + assert tts_code.count('tts.make_gather_scatter_tptr') == 1, "Expected exactly one call to tts.make_gather_scatter_tptr" expected_output = torch.full((b, m, n), -1, device=device, dtype=torch.float32) mask_n = index_n < n @@ -281,4 +321,61 @@ def test_generic_mask_3d_non_continuous_store(device): index_n = index_n[mask_n] expected_output[:,:,index_n] = v - torch.testing.assert_close(output, expected_output) \ No newline at end of file + torch.testing.assert_close(output, expected_output) + + +@triton.jit +def unstructured_mask_and_structured_mask_kernel(in_ptr, out_ptr, m, n, stride_pid, stride_m, stride_n, M:tl.constexpr, N:tl.constexpr): + pid = tl.program_id(0) + offs_m = tl.arange(0, M) + offs_n = tl.arange(0, N) + + mask_m = offs_m > pid + mask_n = offs_n > pid + + in_ptrs = in_ptr + pid * stride_pid + offs_m[:, None] * stride_m + offs_n[None, :] * stride_n + # dim 0 with unstructured mask. + v = tl.load(in_ptrs, mask=offs_m[:, None] < m and offs_n[None, :] < n and mask_m[:, None], + other=-2) + out_ptrs = out_ptr + pid * stride_pid + offs_m[:, None] * stride_m + offs_n[None, :] * stride_n + # dim 1 with unstructured mask. + tl.store(out_ptrs, v, mask=offs_m[:, None] < m and offs_n[None, :] < n and mask_n[None, :]) + + +def unstructured_mask_and_structured_mask_torch(input, n_pids, m, n, M, N): + output = torch.full_like(input, -1) + tmp = torch.full_like(input, -2) + for pid in range(n_pids): + offs_m = torch.arange(0, m) + offs_n = torch.arange(0, n) + + mask_m = (offs_m > pid) + mask_n = (offs_n > pid) + + tmp[pid, mask_m, :] = input[pid, mask_m, :] + output[pid, :, mask_n] = tmp[pid, :, mask_n] + + return output + +def test_unstructured_mask_and_structured_mask(device): + if device == 'cpu': + triton.runtime.driver.set_active(CPUDriver()) + n_pids = 3 + m = 6 + n = 16 + M = triton.next_power_of_2(m) + N = triton.next_power_of_2(n) + input = torch.arange(2, 2 + (n_pids * m * n), device=device, dtype=torch.float32).reshape(n_pids, m, n) + output = torch.full_like(input, -1) + + stride_pid = input.stride(0) + stride_m = input.stride(1) + stride_n = input.stride(2) + + res = unstructured_mask_and_structured_mask_kernel[n_pids,](input, output, m, n, stride_pid, stride_m, stride_n, M, N) + tts_code = run_triton_to_structured(res.asm['ttir']) + assert tts_code.count('tts.make_gather_scatter_tptr') == 2, "Expected exactly two calls to tts.make_gather_scatter_tptr" + + expected_output = unstructured_mask_and_structured_mask_torch(input, n_pids, m, n, M, N) + + torch.testing.assert_close(output, expected_output) diff --git a/test/Conversion/TritonToStructured/generic_mask_2d_kernel.mlir b/test/Conversion/TritonToStructured/unstructured_mask_2d_kernel.mlir similarity index 98% rename from test/Conversion/TritonToStructured/generic_mask_2d_kernel.mlir rename to test/Conversion/TritonToStructured/unstructured_mask_2d_kernel.mlir index e635736b5..990f997a3 100644 --- a/test/Conversion/TritonToStructured/generic_mask_2d_kernel.mlir +++ b/test/Conversion/TritonToStructured/unstructured_mask_2d_kernel.mlir @@ -1,6 +1,6 @@ // RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s -// Make sure make_gather_scatter_tptr with generic mask generate correctly. +// Make sure make_gather_scatter_tptr with unsturctured mask generate correctly from structured ptr with unstructured mask. // CHECK-LABEL: tt.func public @generic_mask_2d_kernel( // CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, diff --git a/test/Conversion/TritonToStructured/generic_mask_2d_non_continuous_load_kernel.mlir b/test/Conversion/TritonToStructured/unstructured_mask_2d_non_continuous_load_kernel.mlir similarity index 91% rename from test/Conversion/TritonToStructured/generic_mask_2d_non_continuous_load_kernel.mlir rename to test/Conversion/TritonToStructured/unstructured_mask_2d_non_continuous_load_kernel.mlir index 3eafac2b7..b943465b3 100644 --- a/test/Conversion/TritonToStructured/generic_mask_2d_non_continuous_load_kernel.mlir +++ b/test/Conversion/TritonToStructured/unstructured_mask_2d_non_continuous_load_kernel.mlir @@ -1,6 +1,6 @@ // RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s -// Make sure make_gather_scatter_tptr with generic mask generate correctly. +// Make sure make_gather_scatter_tptr with unsturctured mask generate correctly from row-structured ptr with unstructured mask. // CHECK-LABEL: tt.func public @generic_mask_2d_non_continuous_load_kernel( // CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, @@ -33,10 +33,7 @@ // CHECK: %[[VAL_27:.*]] = arith.minsi %[[VAL_26]], %[[VAL_8]] : index // CHECK: %[[VAL_28:.*]] = tts.make_gather_scatter_tptr %[[VAL_0]] to sizes: [4, 16] gather_scatter_dim: 0 gather_scatter_offset: %[[VAL_19]] gather_scatter_mask: %[[VAL_22]], strides: {{\[}}%[[VAL_23]], 1], offsets: [0, 0] : tensor<4xi32> tensor<4xi1> to !tt.ptr> // CHECK: %[[VAL_29:.*]] = "tts.load"(%[[VAL_28]], %[[VAL_27]], %[[VAL_7]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (!tt.ptr>, index, f32) -> tensor<4x16xf32> -// CHECK: %[[VAL_30:.*]] = tts.make_tptr %[[VAL_1]] to sizes: [4, 16], strides: {{\[}}%[[VAL_23]], 1], offsets: [0, 0], shape: [0, 0], order: [] : to tensor<4x16x!tt.ptr> -// CHECK: %[[VAL_31:.*]] = arith.minsi %[[VAL_18]], %[[VAL_11]] : index -// CHECK: "tts.store"(%[[VAL_30]], %[[VAL_29]], %[[VAL_31]], %[[VAL_27]]) <{static_mask_dims = array}> : (tensor<4x16x!tt.ptr>, tensor<4x16xf32>, index, index) -> () -// CHECK: tt.return + module { tt.func public @generic_mask_2d_non_continuous_load_kernel(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: !tt.ptr {tt.divisibility = 16 : i32}, %arg3: i32, %arg4: i32, %arg5: i32 {tt.divisibility = 16 : i32}, %arg6: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { diff --git a/test/Conversion/TritonToStructured/generic_mask_2d_non_continuous_store_kernel.mlir b/test/Conversion/TritonToStructured/unstructured_mask_2d_non_continuous_store_kernel.mlir similarity index 98% rename from test/Conversion/TritonToStructured/generic_mask_2d_non_continuous_store_kernel.mlir rename to test/Conversion/TritonToStructured/unstructured_mask_2d_non_continuous_store_kernel.mlir index a306e1668..afe131fff 100644 --- a/test/Conversion/TritonToStructured/generic_mask_2d_non_continuous_store_kernel.mlir +++ b/test/Conversion/TritonToStructured/unstructured_mask_2d_non_continuous_store_kernel.mlir @@ -1,6 +1,6 @@ // RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s -// Make sure make_gather_scatter_tptr with generic mask generate correctly. +// Make sure make_gather_scatter_tptr with unsturctured mask generate correctly from column-structured ptr with unstructured mask. // CHECK-LABEL: tt.func public @generic_mask_2d_non_continuous_store_kernel( // CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, diff --git a/test/Conversion/TritonToStructured/generic_mask_3d_kernel.mlir b/test/Conversion/TritonToStructured/unstructured_mask_3d_kernel.mlir similarity index 99% rename from test/Conversion/TritonToStructured/generic_mask_3d_kernel.mlir rename to test/Conversion/TritonToStructured/unstructured_mask_3d_kernel.mlir index ff8d80f4d..cb55c732d 100644 --- a/test/Conversion/TritonToStructured/generic_mask_3d_kernel.mlir +++ b/test/Conversion/TritonToStructured/unstructured_mask_3d_kernel.mlir @@ -1,6 +1,6 @@ // RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s -// Make sure make_gather_scatter_tptr with generic mask generate correctly. +// Make sure make_gather_scatter_tptr with generic mask generate correctly from structured ptr with unstructured mask. // CHECK-LABEL: tt.func public @generic_mask_3d_kernel( // CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, diff --git a/test/Conversion/TritonToStructured/generic_mask_3d_non_continuous_load_kernel.mlir b/test/Conversion/TritonToStructured/unstructured_mask_3d_non_continuous_load_kernel.mlir similarity index 91% rename from test/Conversion/TritonToStructured/generic_mask_3d_non_continuous_load_kernel.mlir rename to test/Conversion/TritonToStructured/unstructured_mask_3d_non_continuous_load_kernel.mlir index b2c6e8839..ff924202e 100644 --- a/test/Conversion/TritonToStructured/generic_mask_3d_non_continuous_load_kernel.mlir +++ b/test/Conversion/TritonToStructured/unstructured_mask_3d_non_continuous_load_kernel.mlir @@ -1,6 +1,6 @@ // RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s -// Make sure make_gather_scatter_tptr with generic mask generate correctly. +// Make sure make_gather_scatter_tptr with generic mask generate correctly from row-structured ptr with unstructured mask. // CHECK-LABEL: tt.func public @generic_mask_3d_non_continuous_load_kernel( // CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, @@ -43,13 +43,7 @@ // CHECK: %[[VAL_37:.*]] = arith.minsi %[[VAL_35]], %[[VAL_12]] : index // CHECK: %[[VAL_38:.*]] = tts.make_gather_scatter_tptr %[[VAL_0]] to sizes: [4, 4, 16] gather_scatter_dim: 1 gather_scatter_offset: %[[VAL_23]] gather_scatter_mask: %[[VAL_26]], strides: {{\[}}%[[VAL_27]], %[[VAL_28]], 1], offsets: [0, 0, 0] : tensor<4xi32> tensor<4xi1> to !tt.ptr> // CHECK: %[[VAL_39:.*]] = "tts.load"(%[[VAL_38]], %[[VAL_36]], %[[VAL_37]], %[[VAL_11]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (!tt.ptr>, index, index, f32) -> tensor<4x4x16xf32> -// CHECK: %[[VAL_40:.*]] = arith.index_cast %[[VAL_9]] : i32 to index -// CHECK: %[[VAL_41:.*]] = arith.index_cast %[[VAL_10]] : i32 to index -// CHECK: %[[VAL_42:.*]] = tts.make_tptr %[[VAL_1]] to sizes: [4, 4, 16], strides: {{\[}}%[[VAL_40]], %[[VAL_41]], 1], offsets: [0, 0, 0], shape: [0, 0, 0], order: [] : to tensor<4x4x16x!tt.ptr> -// CHECK: %[[VAL_43:.*]] = arith.minsi %[[VAL_22]], %[[VAL_15]] : index -// CHECK: %[[VAL_44:.*]] = arith.minsi %[[VAL_43]], %[[VAL_15]] : index -// CHECK: "tts.store"(%[[VAL_42]], %[[VAL_39]], %[[VAL_36]], %[[VAL_44]], %[[VAL_37]]) <{static_mask_dims = array}> : (tensor<4x4x16x!tt.ptr>, tensor<4x4x16xf32>, index, index, index) -> () -// CHECK: tt.return + module { tt.func public @generic_mask_3d_non_continuous_load_kernel(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: !tt.ptr {tt.divisibility = 16 : i32}, %arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32 {tt.divisibility = 16 : i32}, %arg7: i32 {tt.divisibility = 16 : i32}, %arg8: i32 {tt.divisibility = 16 : i32}, %arg9: i32 {tt.divisibility = 16 : i32}, %arg10: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { diff --git a/test/Conversion/TritonToStructured/generic_mask_3d_non_continuous_store_kernel.mlir b/test/Conversion/TritonToStructured/unstructured_mask_3d_non_continuous_store_kernel.mlir similarity index 99% rename from test/Conversion/TritonToStructured/generic_mask_3d_non_continuous_store_kernel.mlir rename to test/Conversion/TritonToStructured/unstructured_mask_3d_non_continuous_store_kernel.mlir index 681d1c8e9..b44ef8c81 100644 --- a/test/Conversion/TritonToStructured/generic_mask_3d_non_continuous_store_kernel.mlir +++ b/test/Conversion/TritonToStructured/unstructured_mask_3d_non_continuous_store_kernel.mlir @@ -1,6 +1,7 @@ // RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s -// Make sure make_gather_scatter_tptr with generic mask generate correctly. +// Make sure make_gather_scatter_tptr with generic mask generate correctly from column-structured ptr with unstructured mask. + // CHECK-LABEL: tt.func public @generic_mask_3d_non_continuous_store_kernel( // CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, From b06f5f3f75e2db55057199b215d22e4b2956646e Mon Sep 17 00:00:00 2001 From: Xiang Li Date: Fri, 29 Aug 2025 15:23:23 +0000 Subject: [PATCH 3/4] Update test. --- .../unstructured_mask_3d_kernel.mlir | 78 ++++++++++--------- 1 file changed, 41 insertions(+), 37 deletions(-) diff --git a/test/Conversion/TritonToStructured/unstructured_mask_3d_kernel.mlir b/test/Conversion/TritonToStructured/unstructured_mask_3d_kernel.mlir index cb55c732d..cb7462e3e 100644 --- a/test/Conversion/TritonToStructured/unstructured_mask_3d_kernel.mlir +++ b/test/Conversion/TritonToStructured/unstructured_mask_3d_kernel.mlir @@ -12,44 +12,48 @@ // CHECK-SAME: %[[VAL_6:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32 {tt.divisibility = 16 : i32}, // CHECK-SAME: %[[VAL_7:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32 {tt.divisibility = 16 : i32}, // CHECK-SAME: %[[VAL_8:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { -// CHECK: %[[VAL_9:.*]] = arith.constant -2.000000e+00 : f32 -// CHECK: %[[VAL_10:.*]] = arith.constant 4 : index -// CHECK: %[[VAL_11:.*]] = arith.constant 16 : index -// CHECK: %[[VAL_12:.*]] = arith.constant 0 : i8 -// CHECK: %[[VAL_13:.*]] = arith.constant 0 : index -// CHECK: %[[VAL_14:.*]] = arith.constant 8 : index -// CHECK: %[[VAL_15:.*]] = arith.constant dense<0> : tensor<16xi32> -// CHECK: %[[VAL_16:.*]] = arith.constant dense<0> : tensor<8xi32> -// CHECK: %[[VAL_17:.*]] = tts.make_tptr %[[VAL_2]] to sizes: [8], strides: [1], offsets: [0], shape: [0], order: [] : to tensor<8x!tt.ptr> -// CHECK: %[[VAL_18:.*]] = arith.index_cast %[[VAL_5]] : i32 to index -// CHECK: %[[VAL_19:.*]] = arith.minsi %[[VAL_18]], %[[VAL_14]] : index -// CHECK: %[[VAL_20:.*]] = arith.maxsi %[[VAL_19]], %[[VAL_13]] : index -// CHECK: %[[VAL_21:.*]] = "tts.load"(%[[VAL_17]], %[[VAL_20]], %[[VAL_12]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (tensor<8x!tt.ptr>, index, i8) -> tensor<8xi8> -// CHECK: %[[VAL_22:.*]] = arith.extsi %[[VAL_21]] : tensor<8xi8> to tensor<8xi32> -// CHECK: %[[VAL_23:.*]] = arith.cmpi ne, %[[VAL_22]], %[[VAL_16]] : tensor<8xi32> -// CHECK: %[[VAL_24:.*]] = tts.make_tptr %[[VAL_3]] to sizes: [16], strides: [1], offsets: [0], shape: [0], order: [] : to tensor<16x!tt.ptr> -// CHECK: %[[VAL_25:.*]] = arith.index_cast %[[VAL_6]] : i32 to index -// CHECK: %[[VAL_26:.*]] = arith.minsi %[[VAL_25]], %[[VAL_11]] : index -// CHECK: %[[VAL_27:.*]] = arith.maxsi %[[VAL_26]], %[[VAL_13]] : index -// CHECK: %[[VAL_28:.*]] = "tts.load"(%[[VAL_24]], %[[VAL_27]], %[[VAL_12]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (tensor<16x!tt.ptr>, index, i8) -> tensor<16xi8> -// CHECK: %[[VAL_29:.*]] = arith.extsi %[[VAL_28]] : tensor<16xi8> to tensor<16xi32> -// CHECK: %[[VAL_30:.*]] = arith.cmpi ne, %[[VAL_29]], %[[VAL_15]] : tensor<16xi32> -// CHECK: %[[VAL_31:.*]] = arith.index_cast %[[VAL_7]] : i32 to index -// CHECK: %[[VAL_32:.*]] = arith.index_cast %[[VAL_8]] : i32 to index -// CHECK: %[[VAL_33:.*]] = arith.index_cast %[[VAL_4]] : i32 to index -// CHECK: %[[VAL_34:.*]] = arith.minsi %[[VAL_33]], %[[VAL_10]] : index -// CHECK: %[[VAL_35:.*]] = arith.maxsi %[[VAL_34]], %[[VAL_13]] : index -// CHECK: %[[VAL_36:.*]] = arith.minsi %[[VAL_27]], %[[VAL_11]] : index -// CHECK: %[[VAL_37:.*]] = arith.minsi %[[VAL_35]], %[[VAL_10]] : index +// CHECK: %[[VAL_9:.*]] = arith.constant 0 : i32 +// CHECK: %[[VAL_10:.*]] = arith.constant -2.000000e+00 : f32 +// CHECK: %[[VAL_11:.*]] = arith.constant 4 : index +// CHECK: %[[VAL_12:.*]] = arith.constant 16 : index +// CHECK: %[[VAL_13:.*]] = arith.constant 0 : i8 +// CHECK: %[[VAL_14:.*]] = arith.constant 0 : index +// CHECK: %[[VAL_15:.*]] = arith.constant 8 : index +// CHECK: %[[VAL_16:.*]] = arith.constant dense<0> : tensor<16xi32> +// CHECK: %[[VAL_17:.*]] = arith.constant dense<0> : tensor<8xi32> +// CHECK: %[[VAL_18:.*]] = tts.make_tptr %[[VAL_2]] to sizes: [8], strides: [1], offsets: [0], shape: [0], order: [] : to tensor<8x!tt.ptr> +// CHECK: %[[VAL_19:.*]] = arith.index_cast %[[VAL_5]] : i32 to index +// CHECK: %[[VAL_20:.*]] = arith.minsi %[[VAL_19]], %[[VAL_15]] : index +// CHECK: %[[VAL_21:.*]] = arith.maxsi %[[VAL_20]], %[[VAL_14]] : index +// CHECK: %[[VAL_22:.*]] = "tts.load"(%[[VAL_18]], %[[VAL_21]], %[[VAL_13]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (tensor<8x!tt.ptr>, index, i8) -> tensor<8xi8> +// CHECK: %[[VAL_23:.*]] = arith.extsi %[[VAL_22]] : tensor<8xi8> to tensor<8xi32> +// CHECK: %[[VAL_24:.*]] = arith.cmpi ne, %[[VAL_23]], %[[VAL_17]] : tensor<8xi32> +// CHECK: %[[VAL_25:.*]] = tts.make_tptr %[[VAL_3]] to sizes: [16], strides: [1], offsets: [0], shape: [0], order: [] : to tensor<16x!tt.ptr> +// CHECK: %[[VAL_26:.*]] = arith.index_cast %[[VAL_6]] : i32 to index +// CHECK: %[[VAL_27:.*]] = arith.minsi %[[VAL_26]], %[[VAL_12]] : index +// CHECK: %[[VAL_28:.*]] = arith.maxsi %[[VAL_27]], %[[VAL_14]] : index +// CHECK: %[[VAL_29:.*]] = "tts.load"(%[[VAL_25]], %[[VAL_28]], %[[VAL_13]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (tensor<16x!tt.ptr>, index, i8) -> tensor<16xi8> +// CHECK: %[[VAL_30:.*]] = arith.extsi %[[VAL_29]] : tensor<16xi8> to tensor<16xi32> +// CHECK: %[[VAL_31:.*]] = arith.cmpi ne, %[[VAL_30]], %[[VAL_16]] : tensor<16xi32> +// CHECK: %[[VAL_32:.*]] = arith.index_cast %[[VAL_7]] : i32 to index +// CHECK: %[[VAL_33:.*]] = arith.index_cast %[[VAL_8]] : i32 to index +// CHECK: %[[VAL_34:.*]] = arith.index_cast %[[VAL_4]] : i32 to index +// CHECK: %[[VAL_35:.*]] = arith.minsi %[[VAL_34]], %[[VAL_11]] : index +// CHECK: %[[VAL_36:.*]] = arith.maxsi %[[VAL_35]], %[[VAL_14]] : index +// CHECK: %[[VAL_37:.*]] = arith.minsi %[[VAL_28]], %[[VAL_12]] : index // CHECK: %[[VAL_38:.*]] = arith.minsi %[[VAL_36]], %[[VAL_11]] : index -// CHECK: %[[VAL_39:.*]] = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> -// CHECK: %[[VAL_40:.*]] = tts.make_gather_scatter_tptr %[[VAL_0]] to sizes: [4, 8, 16] gather_scatter_dim: 1 gather_scatter_offset: %[[VAL_39]] gather_scatter_mask: %[[VAL_23]], strides: {{\[}}%[[VAL_31]], %[[VAL_32]], 1], offsets: [0, 0, 0] : tensor<8xi32> tensor<8xi1> to !tt.ptr> -// CHECK: %[[VAL_41:.*]] = "tts.load"(%[[VAL_40]], %[[VAL_37]], %[[VAL_38]], %[[VAL_9]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (!tt.ptr>, index, index, f32) -> tensor<4x8x16xf32> -// CHECK: %[[VAL_42:.*]] = arith.minsi %[[VAL_20]], %[[VAL_14]] : index -// CHECK: %[[VAL_43:.*]] = arith.minsi %[[VAL_42]], %[[VAL_14]] : index -// CHECK: %[[VAL_44:.*]] = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> -// CHECK: %[[VAL_45:.*]] = tts.make_gather_scatter_tptr %[[VAL_1]] to sizes: [4, 8, 16] gather_scatter_dim: 2 gather_scatter_offset: %[[VAL_44]] gather_scatter_mask: %[[VAL_30]], strides: {{\[}}%[[VAL_31]], %[[VAL_32]], 1], offsets: [0, 0, 0] : tensor<16xi32> tensor<16xi1> to !tt.ptr> -// CHECK: "tts.store"(%[[VAL_45]], %[[VAL_41]], %[[VAL_37]], %[[VAL_43]]) <{static_mask_dims = array}> : (!tt.ptr>, tensor<4x8x16xf32>, index, index) -> () +// CHECK: %[[VAL_39:.*]] = arith.minsi %[[VAL_37]], %[[VAL_12]] : index +// CHECK: %[[VAL_40:.*]] = arith.divui %[[VAL_9]], %[[VAL_8]] : i32 +// CHECK: %[[VAL_41:.*]] = tensor.splat %[[VAL_40]] : tensor<8xi32> +// CHECK: %[[VAL_42:.*]] = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> +// CHECK: %[[VAL_43:.*]] = arith.addi %[[VAL_41]], %[[VAL_42]] : tensor<8xi32> +// CHECK: %[[VAL_44:.*]] = tts.make_gather_scatter_tptr %[[VAL_0]] to sizes: [4, 8, 16] gather_scatter_dim: 1 gather_scatter_offset: %[[VAL_43]] gather_scatter_mask: %[[VAL_24]], strides: {{\[}}%[[VAL_32]], %[[VAL_33]], 1], offsets: [0, 0, 0] : tensor<8xi32> tensor<8xi1> to !tt.ptr> +// CHECK: %[[VAL_45:.*]] = "tts.load"(%[[VAL_44]], %[[VAL_38]], %[[VAL_39]], %[[VAL_10]]) <{operandSegmentSizes = array, static_mask_dims = array}> : (!tt.ptr>, index, index, f32) -> tensor<4x8x16xf32> +// CHECK: %[[VAL_46:.*]] = arith.minsi %[[VAL_21]], %[[VAL_15]] : index +// CHECK: %[[VAL_47:.*]] = arith.minsi %[[VAL_46]], %[[VAL_15]] : index +// CHECK: %[[VAL_48:.*]] = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> +// CHECK: %[[VAL_49:.*]] = tts.make_gather_scatter_tptr %[[VAL_1]] to sizes: [4, 8, 16] gather_scatter_dim: 2 gather_scatter_offset: %[[VAL_48]] gather_scatter_mask: %[[VAL_31]], strides: {{\[}}%[[VAL_32]], %[[VAL_33]], 1], offsets: [0, 0, 0] : tensor<16xi32> tensor<16xi1> to !tt.ptr> +// CHECK: "tts.store"(%[[VAL_49]], %[[VAL_45]], %[[VAL_38]], %[[VAL_47]]) <{static_mask_dims = array}> : (!tt.ptr>, tensor<4x8x16xf32>, index, index) -> () // CHECK: tt.return module { From 76eb4e03736542a535aaa2af186b4af7b446776d Mon Sep 17 00:00:00 2001 From: Xiang Li Date: Wed, 3 Sep 2025 02:18:17 +0000 Subject: [PATCH 4/4] Update per comment. --- lib/AnalysisStructured/PtrAnalysis.cpp | 21 ++++++++++--------- .../unstructured_mask_2d_kernel.mlir | 4 +++- ...ed_mask_2d_non_continuous_load_kernel.mlir | 3 ++- ...d_mask_2d_non_continuous_store_kernel.mlir | 3 ++- .../unstructured_mask_3d_kernel.mlir | 2 ++ ...ed_mask_3d_non_continuous_load_kernel.mlir | 1 + ...d_mask_3d_non_continuous_store_kernel.mlir | 1 + 7 files changed, 22 insertions(+), 13 deletions(-) diff --git a/lib/AnalysisStructured/PtrAnalysis.cpp b/lib/AnalysisStructured/PtrAnalysis.cpp index 00b5718cd..65534d3c8 100644 --- a/lib/AnalysisStructured/PtrAnalysis.cpp +++ b/lib/AnalysisStructured/PtrAnalysis.cpp @@ -53,22 +53,23 @@ static Value applyUnstructuredMask(Operation *op, Value ptr, } auto [dim, unstructuredMask] = masks[0]; - if (auto scatterPtr = + if (auto gatherScatterPtr = ptr.getDefiningOp()) { - if (dim != scatterPtr.getGatherScatterDim()) { + if (dim != gatherScatterPtr.getGatherScatterDim()) { op->emitRemark("MaskAnalysis failed for unstructured mask dim not equal " "gather scatter dim"); return nullptr; } - ptr = builder - .create( - loc, scatterPtr.getBase(), - scatterPtr.getGatherScatterOffset(), unstructuredMask, - scatterPtr.getGatherScatterDim(), scatterPtr.getSizes(), - scatterPtr.getMixedStrides(), scatterPtr.getMixedOffsets()) - .getResult(); - + ptr = + builder + .create( + loc, gatherScatterPtr.getBase(), + gatherScatterPtr.getGatherScatterOffset(), unstructuredMask, + gatherScatterPtr.getGatherScatterDim(), + gatherScatterPtr.getSizes(), gatherScatterPtr.getMixedStrides(), + gatherScatterPtr.getMixedOffsets()) + .getResult(); } else if (auto structuredPtr = ptr.getDefiningOp()) { auto ofrToI32Value = [&](OpFoldResult ofr) { Value v = dyn_cast(ofr); diff --git a/test/Conversion/TritonToStructured/unstructured_mask_2d_kernel.mlir b/test/Conversion/TritonToStructured/unstructured_mask_2d_kernel.mlir index 990f997a3..bc6dfa709 100644 --- a/test/Conversion/TritonToStructured/unstructured_mask_2d_kernel.mlir +++ b/test/Conversion/TritonToStructured/unstructured_mask_2d_kernel.mlir @@ -1,6 +1,8 @@ // RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s -// Make sure make_gather_scatter_tptr with unsturctured mask generate correctly from structured ptr with unstructured mask. +// Make sure make_gather_scatter_tptr with unstructuredmask generate correctly from structured ptr with unstructured mask. +// The load is structured ptr, with unstructured mask on dim 0. +// The store is structured ptr, with unstructured mask on dim 1. // CHECK-LABEL: tt.func public @generic_mask_2d_kernel( // CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, diff --git a/test/Conversion/TritonToStructured/unstructured_mask_2d_non_continuous_load_kernel.mlir b/test/Conversion/TritonToStructured/unstructured_mask_2d_non_continuous_load_kernel.mlir index b943465b3..0abd0c8f9 100644 --- a/test/Conversion/TritonToStructured/unstructured_mask_2d_non_continuous_load_kernel.mlir +++ b/test/Conversion/TritonToStructured/unstructured_mask_2d_non_continuous_load_kernel.mlir @@ -1,6 +1,7 @@ // RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s -// Make sure make_gather_scatter_tptr with unsturctured mask generate correctly from row-structured ptr with unstructured mask. +// Make sure make_gather_scatter_tptr with unstructuredmask generate correctly from row-structured ptr with unstructured mask. +// The load is unstructured ptr on dim 0 and unstructured mask on dim 0. // CHECK-LABEL: tt.func public @generic_mask_2d_non_continuous_load_kernel( // CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, diff --git a/test/Conversion/TritonToStructured/unstructured_mask_2d_non_continuous_store_kernel.mlir b/test/Conversion/TritonToStructured/unstructured_mask_2d_non_continuous_store_kernel.mlir index afe131fff..854c38a1b 100644 --- a/test/Conversion/TritonToStructured/unstructured_mask_2d_non_continuous_store_kernel.mlir +++ b/test/Conversion/TritonToStructured/unstructured_mask_2d_non_continuous_store_kernel.mlir @@ -1,6 +1,7 @@ // RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s -// Make sure make_gather_scatter_tptr with unsturctured mask generate correctly from column-structured ptr with unstructured mask. +// Make sure make_gather_scatter_tptr with unstructuredmask generate correctly from column-structured ptr with unstructured mask. +// The store is unstructured ptr on dim 1 and unstructured mask on dim 1. // CHECK-LABEL: tt.func public @generic_mask_2d_non_continuous_store_kernel( // CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, diff --git a/test/Conversion/TritonToStructured/unstructured_mask_3d_kernel.mlir b/test/Conversion/TritonToStructured/unstructured_mask_3d_kernel.mlir index cb7462e3e..64473d259 100644 --- a/test/Conversion/TritonToStructured/unstructured_mask_3d_kernel.mlir +++ b/test/Conversion/TritonToStructured/unstructured_mask_3d_kernel.mlir @@ -1,6 +1,8 @@ // RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s // Make sure make_gather_scatter_tptr with generic mask generate correctly from structured ptr with unstructured mask. +// The load is structured ptr, with unstructured mask on dim 1. +// The store is structured ptr, with unstructured mask on dim 2. // CHECK-LABEL: tt.func public @generic_mask_3d_kernel( // CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, diff --git a/test/Conversion/TritonToStructured/unstructured_mask_3d_non_continuous_load_kernel.mlir b/test/Conversion/TritonToStructured/unstructured_mask_3d_non_continuous_load_kernel.mlir index ff924202e..691399b90 100644 --- a/test/Conversion/TritonToStructured/unstructured_mask_3d_non_continuous_load_kernel.mlir +++ b/test/Conversion/TritonToStructured/unstructured_mask_3d_non_continuous_load_kernel.mlir @@ -1,6 +1,7 @@ // RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s // Make sure make_gather_scatter_tptr with generic mask generate correctly from row-structured ptr with unstructured mask. +// The load is unstructured ptr on dim 1 and unstructured mask on dim 1. // CHECK-LABEL: tt.func public @generic_mask_3d_non_continuous_load_kernel( // CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr {tt.divisibility = 16 : i32}, diff --git a/test/Conversion/TritonToStructured/unstructured_mask_3d_non_continuous_store_kernel.mlir b/test/Conversion/TritonToStructured/unstructured_mask_3d_non_continuous_store_kernel.mlir index b44ef8c81..8f99d00b6 100644 --- a/test/Conversion/TritonToStructured/unstructured_mask_3d_non_continuous_store_kernel.mlir +++ b/test/Conversion/TritonToStructured/unstructured_mask_3d_non_continuous_store_kernel.mlir @@ -1,6 +1,7 @@ // RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s // Make sure make_gather_scatter_tptr with generic mask generate correctly from column-structured ptr with unstructured mask. +// The store is unstructured ptr on dim 2 and unstructured mask on dim 2. // CHECK-LABEL: tt.func public @generic_mask_3d_non_continuous_store_kernel(