From 0eebeb9d2fc10a0df8ff6c35892bfa6613e4a110 Mon Sep 17 00:00:00 2001 From: chengwei Date: Wed, 10 Jun 2026 11:15:57 +0800 Subject: [PATCH] [BACKEND] Sync tsingmicro code --- include/triton-shared/Analysis/MaskAnalysis.h | 30 +- .../Analysis/OpFoldResultUtils.h | 2 + .../AnalysisStructured/PtrAnalysis.h | 23 +- .../IR/TritonStructuredDialect.td | 114 +++- lib/Analysis/CMakeLists.txt | 1 + lib/Analysis/MaskAnalysis.cpp | 209 +++++- lib/Analysis/OpFoldResultUtils.cpp | 20 + lib/AnalysisStructured/PtrAnalysis.cpp | 13 +- lib/AnalysisStructured/PtrAnalysisTS.cpp | 640 ++++++++++++++---- .../MemrefCopyToDMAFlagTree.cpp | 3 - .../NoBufferizeFlagTree.cpp | 1 - .../TritonToUnstructuredPass.cpp | 42 +- .../TritonStructured/IR/CMakeLists.txt | 1 + .../IR/TritonStructuredDialect.cpp | 109 +++ .../IR/TritonStructuredOps.cpp | 66 ++ lib/Utils/Utils.cpp | 8 +- 16 files changed, 1106 insertions(+), 176 deletions(-) diff --git a/include/triton-shared/Analysis/MaskAnalysis.h b/include/triton-shared/Analysis/MaskAnalysis.h index fefea7a0..c0412791 100644 --- a/include/triton-shared/Analysis/MaskAnalysis.h +++ b/include/triton-shared/Analysis/MaskAnalysis.h @@ -85,6 +85,23 @@ struct dimInfo { // // Example of creating 2D mask: // mask = (rows[:, None] < M) & (cols[None, :] < N) +// +// NOTE: Bool tensor mask could be saved into dims and record the dim in +// nonContinuousDim 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 +// 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 { OpFoldResult start; OpFoldResult end; @@ -93,13 +110,24 @@ struct MaskState { const bool useUnsafeMask; ///ASCEND SmallVector stateInfo; - + + SmallVector nonContinuousDim; void dump() const; MaskState(bool useUnsafeMask = false) : useUnsafeMask(useUnsafeMask) {} int64_t getRank() const { return dims.size(); } + bool hasNonContinuousDim() const { return !nonContinuousDim.empty(); } + + std::optional getNonContinuousDim() const { + if (hasNonContinuousDim()) { + assert(nonContinuousDim.size() == 1); + return nonContinuousDim[0]; + } + return std::nullopt; + } + bool isEmpty() const { return getRank() == 0 && !scalar && !start && !end; } bool isMask() const { return !start && !end && !scalar && dims.size() != 0; } diff --git a/include/triton-shared/Analysis/OpFoldResultUtils.h b/include/triton-shared/Analysis/OpFoldResultUtils.h index b1517224..84e1f341 100644 --- a/include/triton-shared/Analysis/OpFoldResultUtils.h +++ b/include/triton-shared/Analysis/OpFoldResultUtils.h @@ -24,6 +24,8 @@ Value materializeValue(OpBuilder &builder, Location loc, OpFoldResult ofr); // result of an operation too. std::optional getIntAttr(const OpFoldResult ofr); +std::optional getConstValue(const OpFoldResult ofr); + // Return if ofr contains a constant zero, either represented by an integer // attribute or a constant value. bool hasConstZero(const OpFoldResult ofr); diff --git a/include/triton-shared/AnalysisStructured/PtrAnalysis.h b/include/triton-shared/AnalysisStructured/PtrAnalysis.h index f9c9b60b..661d0439 100644 --- a/include/triton-shared/AnalysisStructured/PtrAnalysis.h +++ b/include/triton-shared/AnalysisStructured/PtrAnalysis.h @@ -38,6 +38,11 @@ const extern std::string ptrAnalysisAttr; // shape field means the same field as tt.make_tensor_ptr; when it describes a // non-block pointer, shape field indicates how address wraps around (i.e., // modulo); a constant 0 indicates no modulo for the dimension. +// Multi-dimension PtrState, which has one unstructured dimension, is supported +// for gather/scatter access. The unstructured dimension is marked by a tensor +// type offset. The tensor offset for the unstructured dimension must be +// expanded from a 1D tensor. The analysis will fail for multi-dimension +// unstructured offsets. struct PtrState { SmallVector offsets; @@ -57,6 +62,10 @@ struct PtrState { bool dimHasModulo(uint32_t dim) const; + bool dimIsStructured(uint32_t dim) const; + int32_t getNonStructuredDim() const; + bool isStructured() const; + bool isBlockPtr() const; void dump() const; @@ -73,8 +82,8 @@ struct PtrState { LogicalResult mulState(const PtrState &lhsState, const PtrState &rhsState, Operation *op, OpBuilder &builder); - tts::MakeTensorPtrOp createTTSMakeTensorPtrOp(OpBuilder &builder, - Location loc); + // All dim structured or only one dim unstructured + Operation *createTTSMakeTensorPtrOp(OpBuilder &builder, Location loc); }; class PtrAnalysis { @@ -162,6 +171,16 @@ class PtrAnalysis { LogicalResult visitOperandRem(arith::RemSIOp mulOp, PtrState &state, const Location loc, OpBuilder &builder); + // Operand is the result of arith.divsi (signed integer division). + // Main assumptions: + // Divisor must be a constant scalar + // Dividend must be 1D tensor or 2D tensor with one dimension being 1 + // Division is applied to offsets, strides, and scalar values + // Expected result: + // All offset and stride values are divided by the constant divisor + LogicalResult visitOperandDivSI(arith::DivSIOp divOp, PtrState &state, + const Location loc, OpBuilder &builder); + // Operand is the result of make_range. // Main assumptions: // start, end, and shape are all statically known diff --git a/include/triton-shared/Dialect/TritonStructured/IR/TritonStructuredDialect.td b/include/triton-shared/Dialect/TritonStructured/IR/TritonStructuredDialect.td index 8412c4d0..6bcbe3cb 100644 --- a/include/triton-shared/Dialect/TritonStructured/IR/TritonStructuredDialect.td +++ b/include/triton-shared/Dialect/TritonStructured/IR/TritonStructuredDialect.td @@ -18,12 +18,17 @@ def Triton_Structured_Dialect : Dialect { }]; let dependentDialects = [ - "triton::TritonDialect" + "triton::TritonDialect", + "tensor::TensorDialect" ]; let usePropertiesForAttributes = 1; + let hasCanonicalizer = 1; } +// FIXME: AnyTensor +def TT_IndexTensorLike : AnyTypeOf<[I1Tensor, I8Tensor, I16Tensor, I32Tensor, I64Tensor, IndexTensor]>; + // // Op Base // @@ -121,6 +126,113 @@ def TTS_MakeTensorPtrOp //let hasCanonicalizer = 1; } +def TTS_MakeGatherScatterTensorPtrOp + : TTS_Op<"make_gather_scatter_tptr", [AttrSizedOperandSegments, Pure]> { + // NOTE: Only support cases where the offset for each dimension is defined in a different operation. + // Not support case where the offset is a tensor load from other ptr which for multiple dimension. + // + // offset_m = tl.arange(0, M) + // offset_n = tl.arange(0, N) + // offset_k = tl.arange(0, K) + // ld_offsets = tl.load(a_ptr + offset_m[:,None]+offsets_n[None,:]) + // not_support = tl.load(b_ptr + ld_offsets) + // not_support2 = tl.load(b_ptr + ld_offsets * (offset_m[:,None]+offsets_n[None,:])) + // not_support3 = tl.load(b_ptr + (ld_offsets * (offset_m[:,None]+offsets_n[None,:]))[:, :, None] + offset_k[None,None,:]) + // + // # Support cases where one dimension is structured while the other is not. + // # For example, `offset_m[:, None] // K` is not structured, whereas `offset_n[None, :]` is structured in next line. + // supported = tl.load(b_ptr + offset_m[:, None] // K + offset_n[None, :]) + // # `ld_offsets_1d[:, None]` is not structured, whereas `offset_n[None, :]`. + // ld_offsets_1d = tl.load(a_ptr + offsets_m) + // supported2 = tl.load(b_ptr + ld_offsets_1d[:, None] + offsets_n[None, :]) + + let summary = "create an pointer that points to a tensor in memory for gather/scatter"; + let description = [{ + The `tts.make_gather_scatter_tptr` operation is similar to `tts.make_tptr`. + The key difference is that `make_gather_scatter_tptr` accesses the tensor non-continuously. + Currently, only one dimension is allowed to be non-continuous. + This dimension is saved in `gather_scatter_dim`, and the offset for that dimension is saved in `gather_scatter_offset`. + Each contiguous load will load from this offset. + Cases with more than one non-continuous dimension are not supported. + }]; + + // base: Base pointer used to contruct the tensor of pointers or pointer to tensor. + // gather_scatter_dim: The dimension for gather_scatter_offset. + // sizes: Size of the data being loaded or stored. + // 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. + // result: A tensor of pointers. + + let arguments = (ins TT_Ptr:$base, + TT_IndexTensorLike:$gather_scatter_offset, + I32Attr:$gather_scatter_dim, + DenseI64ArrayAttr:$sizes, + Variadic:$strides, + Variadic:$offsets, + DenseI64ArrayAttr:$static_strides, + DenseI64ArrayAttr:$static_offsets, + Optional:$gather_scatter_mask); + + let results = (outs TT_PtrLike:$result); + + let assemblyFormat = [{ + $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($gather_scatter_mask) type($base) `to` type($result) + }]; + + + let builders = [ + // Build with mixed static and dynamic entries. + OpBuilder<(ins + "Value":$base, + "Value":$gather_scatter_offset, + "int":$gather_scatter_dim, + "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 = [{ + /// Return a vector of all the static or dynamic fields + SmallVector getMixedSizes() { + Builder b(getContext()); + SmallVector dynSizes; // sizes are always static + return ::mlir::getMixedValues(getSizes(), dynSizes, b); + } + /// Return a vector of all the static or dynamic fields + SmallVector getMixedStrides() { + Builder b(getContext()); + return ::mlir::getMixedValues(getStaticStrides(), getStrides(), b); + } + SmallVector getMixedOffsets() { + Builder b(getContext()); + return ::mlir::getMixedValues(getStaticOffsets(), getOffsets(), b); + } + }]; + + let hasVerifier = 1; + let hasCanonicalizer = 0; +} + +// SameVariadicResultSize +// AttrSizedResultSegments def TTS_GetStructuredStateOp : TTS_Op<"get_structured_state", [AttrSizedResultSegments, Pure]> { let summary = "Placeholder for the structured pointer states computed during PtrAnalysis."; let description = "Used to pass the offsets and strides to scf.for op to simplify IR rewrites."; diff --git a/lib/Analysis/CMakeLists.txt b/lib/Analysis/CMakeLists.txt index 71a89f94..0165047c 100644 --- a/lib/Analysis/CMakeLists.txt +++ b/lib/Analysis/CMakeLists.txt @@ -11,4 +11,5 @@ add_triton_library(TritonSharedAnalysis LINK_LIBS PUBLIC MLIRAnalysis + Common ) diff --git a/lib/Analysis/MaskAnalysis.cpp b/lib/Analysis/MaskAnalysis.cpp index 4391c53a..e2635b5d 100644 --- a/lib/Analysis/MaskAnalysis.cpp +++ b/lib/Analysis/MaskAnalysis.cpp @@ -40,6 +40,23 @@ void dimInfo::dump() const { }); }; +static Value structuredMaskToUnstructuredMask(OpBuilder &builder, Location loc, + uint32_t size, uint32_t maskDim) { + + if (maskDim == size) { + // Full mask. + return Value(); + } + + auto i32TensorType = RankedTensorType::get({size}, builder.getI32Type()); + Value range = + builder.create(loc, i32TensorType, 0, size); + Value v = builder.create(loc, maskDim); + v = builder.create(loc, builder.getI32Type(), v); + v = builder.create(loc, i32TensorType, v).getResult(); + return builder.create(loc, arith::CmpIPredicate::ult, range, + v); +}; /////////////ascend LogicalResult MaskState::parse(Value operand, const Location loc, @@ -80,8 +97,16 @@ LogicalResult MaskState::parse(Value operand, const Location loc, return this->parseDivsi(op, loc, builder); else return failure(); - } + } else { + auto resultType = operand.getType(); + if (isa(resultType) && + cast(resultType).getRank() == 1) { + assert(isEmpty()); + dims.push_back(operand); + nonContinuousDim.push_back(0); + return success(); + } return failure(); } } @@ -238,13 +263,13 @@ LogicalResult MaskState::subStates(const MaskState &lhsState, OpBuilder &builder) { if (lhsState.scalar && rhsState.scalar) { InFlightDiagnostic diag = - emitError(loc) << "Unexpected case where both lhs and rhs are scalars"; + emitWarning(loc) << "Unexpected case where both lhs and rhs are scalars"; return failure(); } if (!lhsState.scalar && !rhsState.scalar) { InFlightDiagnostic diag = - emitError(loc) + emitWarning(loc) << "Unsupported scenario where neither lhs nor rhs is a scalar"; return failure(); } @@ -260,13 +285,13 @@ LogicalResult MaskState::addStates(const MaskState &lhsState, OpBuilder &builder) { if (lhsState.scalar && rhsState.scalar) { InFlightDiagnostic diag = - emitError(loc) << "Unexpected case where both lhs and rhs are scalars"; + emitWarning(loc) << "Unexpected case where both lhs and rhs are scalars"; return failure(); } if (!lhsState.scalar && !rhsState.scalar) { InFlightDiagnostic diag = - emitError(loc) + emitWarning(loc) << "Unsupported scenario where neither lhs nor rhs is a scalar"; return failure(); } @@ -296,7 +321,7 @@ LogicalResult MaskState::minStateScalar(const MaskState &lhsState, } } else { InFlightDiagnostic diag = - emitError(loc) << "Unexpected case where both lhs and rhs are not scalars"; + emitWarning(loc) << "Unexpected case where both lhs and rhs are not scalars"; return failure(); } return success(); @@ -307,7 +332,7 @@ LogicalResult MaskState::minStates(const MaskState &lhsState, OpBuilder &builder) { if (lhsState.getRank() != rhsState.getRank()) { InFlightDiagnostic diag = - emitError(loc) + emitWarning(loc) << "Unexpected case where lhs and rhs have different ranks"; return failure(); } @@ -375,6 +400,12 @@ LogicalResult MaskState::parseAdd(arith::AddIOp addOp, const Location loc, if (failed(rhsState.parse(addOp.getRhs(), loc, builder))) return failure(); + // WORKAROUND: avoid other non-structured cases. Since mask analysis is not + // consistent with noncontinuous dims now + if (lhsState.hasNonContinuousDim() || rhsState.hasNonContinuousDim()) { + return failure(); + } + return this->addStates(lhsState, rhsState, loc, builder); } @@ -390,6 +421,12 @@ LogicalResult MaskState::parseSub(arith::SubIOp subOp, const Location loc, if (failed(rhsState.parse(subOp.getRhs(), loc, builder))) return failure(); + // WORKAROUND: avoid other non-structured cases. Since mask analysis is not + // consistent with noncontinuous dims now + if (lhsState.hasNonContinuousDim() || rhsState.hasNonContinuousDim()) { + return failure(); + } + return this->subStates(lhsState, rhsState, loc, builder); } @@ -405,6 +442,101 @@ LogicalResult MaskState::parseAnd(arith::AndIOp andOp, const Location loc, if (failed(rhsState.parse(andOp.getRhs(), loc, builder))) return failure(); + if (lhsState.hasNonContinuousDim() || rhsState.hasNonContinuousDim()) { + + auto lhsNonContinuousDim = lhsState.getNonContinuousDim(); + auto rhsNonContinuousDim = rhsState.getNonContinuousDim(); + + // assert( + // !(lhsNonContinuousDim.has_value() && rhsNonContinuousDim.has_value()) && + // "Cannot have both non continuous dims"); + + // If both sides have non-continuous dims at the same dimension: + // directly AND the two boolean tensors element-wise. + if (lhsNonContinuousDim.has_value() && rhsNonContinuousDim.has_value()) { + if (lhsNonContinuousDim.value() != rhsNonContinuousDim.value()) { + InFlightDiagnostic diag = andOp->emitWarning( + "Cannot have non-continuous dims at different positions"); + return failure(); + } + + auto nonContinuousDim = lhsNonContinuousDim.value(); + auto tensorType = cast(andOp.getType()); + assert(tensorType.getElementTypeBitWidth() == 1 && + "Unexpected case where type is not i1"); + + for (uint32_t i = 0; i < lhsState.getRank(); i++) { + if (i == nonContinuousDim) { + auto lhsVal = dyn_cast(lhsState.dims[nonContinuousDim]); + auto rhsVal = dyn_cast(rhsState.dims[nonContinuousDim]); + assert(lhsVal && rhsVal && + "Expected boolean tensor Values for non-continuous dims"); + dims.push_back( + builder.create(loc, lhsVal, rhsVal).getResult()); + continue; + } + auto lhsDim = lhsState.dims[i]; + auto rhsDim = rhsState.dims[i]; + dims.push_back(minOFRs(lhsDim, rhsDim, loc, builder)); + } + this->nonContinuousDim = lhsState.nonContinuousDim; + return success(); + } + + auto [lhs, rhs] = + lhsNonContinuousDim.has_value() + ? std::tuple(&lhsState, &rhsState) + : std::tuple(&rhsState, &lhsState); + + auto nonContinuousDim = lhs->getNonContinuousDim().value(); + auto rhsDim = rhs->dims[nonContinuousDim]; + auto rhsDimInt = getIntAttr(rhsDim); + + auto tensorType = cast(andOp.getType()); + assert(tensorType.getElementTypeBitWidth() == 1 && + "Unexpected case where type is not i1"); + + for (uint32_t i = 0; i < lhsState.getRank(); i++) { + if (i == nonContinuousDim) { + Value rhsDimValue; + if (rhsDimInt.has_value()) { + rhsDimValue = structuredMaskToUnstructuredMask( + builder, loc, tensorType.getShape()[nonContinuousDim], + rhsDimInt.value()); + } else { + // dynamic rhsDim: create range < splat(rhsDim) directly + uint32_t dimSize = tensorType.getShape()[nonContinuousDim]; + auto i32TensorType = + RankedTensorType::get({dimSize}, builder.getI32Type()); + Value range = builder.create(loc, i32TensorType, 0, + dimSize); + Value rhsVal = ofrToIndexValue(rhsDim, loc, builder); + Value rhsValI32 = builder.create( + loc, builder.getI32Type(), rhsVal); + Value rhsValSplat = + builder.create(loc, i32TensorType, rhsValI32); + rhsDimValue = builder.create( + loc, arith::CmpIPredicate::ult, range, rhsValSplat); + } + + dims.push_back( + rhsDimValue + ? builder + .create( + loc, dyn_cast(lhs->dims[nonContinuousDim]), + rhsDimValue) + .getResult() + : lhs->dims[nonContinuousDim]); + continue; + } + + auto lhsDim = lhsState.dims[i]; + auto rhsDim = rhsState.dims[i]; + dims.push_back(minOFRs(lhsDim, rhsDim, loc, builder)); + } + this->nonContinuousDim = lhs->nonContinuousDim; + return success(); + } // TODO(FLIR): should be isMask() if(!lhsState.isMaskWithoutScalar() || !rhsState.isMaskWithoutScalar()) { return this->minStateScalar(lhsState, rhsState, loc, builder); @@ -425,7 +557,7 @@ LogicalResult MaskState::parseCmp(arith::CmpIOp cmpOp, const Location loc, if (cmpOp.getPredicate() != arith::CmpIPredicate::slt && cmpOp.getPredicate() != arith::CmpIPredicate::ult && cmpOp.getPredicate() != arith::CmpIPredicate::sge) { - InFlightDiagnostic diag = emitError(loc) << "Unsupported cmpi"; + InFlightDiagnostic diag = cmpOp->emitWarning("Unsupported cmpi"); return failure(); } @@ -442,19 +574,23 @@ LogicalResult MaskState::parseCmp(arith::CmpIOp cmpOp, const Location loc, // the comparison evaluates to true. if (cmpOp.getPredicate() == arith::CmpIPredicate::sge && !(rhsState.scalar && hasConstZero(rhsState.scalar))) { - InFlightDiagnostic diag = emitError(loc) - << "Unsupported cmpi with rhs not equal to 0"; + InFlightDiagnostic diag = cmpOp->emitWarning( + "Unsupported cmpi with rhs not equal to 0"); return failure(); } + // This cmpDim is used to determine which dimension contains the actual + // data range (i.e., the dimension with a size greater than 1) + // Because mask analysis assumes that only one dimension can have an actual + // range int32_t cmpDim = lhsState.scalar && rhsState.scalar ? 0 : -1; for (int32_t i = 0; i < lhsState.getRank(); i++) { auto dimIntAttr = getIntAttr(lhsState.dims[i]); if (!dimIntAttr || dimIntAttr.value() != 1) { if (cmpDim != -1) { - InFlightDiagnostic diag = emitError(loc) - << "Unsupported cmpi with more than one " - "dimension with size larger than 1"; + InFlightDiagnostic diag = cmpOp->emitWarning( + "Unsupported cmpi with more than one " + "dimension with size larger than 1"); return failure(); } cmpDim = i; @@ -464,7 +600,25 @@ LogicalResult MaskState::parseCmp(arith::CmpIOp cmpOp, const Location loc, "Unexpected case where no dimension has size larger than 1"); OpFoldResult newDim; - if (lhsState.scalar) { + if (lhsState.hasNonContinuousDim()) { + assert(rhsState.scalar && "Unexpected case where rhs is not a scalar"); + assert(cmpDim == lhsState.getNonContinuousDim()); + auto lhsValue = dyn_cast(lhsState.dims[cmpDim]); + auto tensorType = cast(lhsValue.getType()); + + auto rhsValue = builder.create( + loc, tensorType.getElementType(), + ofrToIndexValue(rhsState.scalar, loc, builder)); + auto rhsTensorValue = + builder.create(loc, tensorType, rhsValue); + + newDim = builder + .create(loc, cmpOp.getPredicate(), lhsValue, + rhsTensorValue) + ->getResult(0); + + nonContinuousDim = lhsState.nonContinuousDim; + } else if (lhsState.scalar) { assert(rhsState.scalar && "Unexpected case where rhs is not a scalar"); // If both lhs and rhs are scalars, we can't just derive the dimension of // the mask as the minimum value: lhs/rhs could be 0 and then we don't @@ -492,6 +646,7 @@ LogicalResult MaskState::parseCmp(arith::CmpIOp cmpOp, const Location loc, // The correct formula is to optionally move `newDim` back to `start` using // max(newEnd, start). auto newEnd = minOFRs(lhsState.end, rhsState.scalar, loc, builder); + // Avoid minus dim. ([5, 10] < 3: 3 - 5 = -2�� newEnd = maxOFRs(newEnd, lhsState.start, loc, builder); newDim = subOFRs(newEnd, lhsState.start, loc, builder); } else { @@ -585,9 +740,9 @@ LogicalResult MaskState::parseMakeRange(triton::MakeRangeOp rangeOp, if (stride != 1) { InFlightDiagnostic diag = - emitError(loc) - << "stride must be 1 for make_range whose result is used " - "as load or store masks"; + rangeOp->emitWarning( + "stride must be 1 for make_range whose result is used " + "as load or store masks"); return failure(); } @@ -638,8 +793,8 @@ LogicalResult MaskState::parseSplat(triton::SplatOp splatOp, const Location loc, if (!isa(src.getType())) { InFlightDiagnostic diag = - emitError(loc) - << "splat source must be an integer scalar for load/store masks"; + splatOp->emitWarning( + "splat source must be an integer scalar for load/store masks"); return failure(); } @@ -667,6 +822,10 @@ LogicalResult MaskState::parseExpandDims(triton::ExpandDimsOp expandDimsOp, "expect changed dimension to be 1 in expand_dims"); this->dims.insert(this->dims.begin() + axis, builder.getIndexAttr(1)); + if (hasNonContinuousDim()) { + assert(dstShape.size() == 2); + this->nonContinuousDim[0] = 1 - axis; + } return success(); } ////////ASCEND @@ -698,7 +857,7 @@ LogicalResult MaskState::parseRemsi(arith::RemSIOp remsiOp, assert(rhsIntAttr.has_value()); staticShape = rhsIntAttr.value(); }else{ - remsiOp->emitError("MaskAnalysis: Static compilation cannot determine the value of this parameter"); + remsiOp->emitWarning("MaskAnalysis: Static compilation cannot determine the value of this parameter"); return failure(); } @@ -710,11 +869,11 @@ LogicalResult MaskState::parseRemsi(arith::RemSIOp remsiOp, if(info.isRealDim){ auto staticDim = getIntAttr(rhsState.scalar); if(!staticDim.has_value() || (staticDim.value() % staticShape != 0 && staticShape % staticDim.value() != 0)){ - remsiOp->emitError("MaskAnalysis: The shape of the mask is not divisible by the shape of the block"); + remsiOp->emitWarning("MaskAnalysis: The shape of the mask is not divisible by the shape of the block"); return failure(); } // if(getIntAttr(info.div).has_value() && getIntAttr(info.div).value() != 0){ - // remsiOp->emitError("MaskAnalysis: do not support remsi after div"); + // remsiOp->emitWarning("MaskAnalysis: do not support remsi after div"); // return failure(); // } info.shape = builder.getIndexAttr(staticShape); @@ -753,7 +912,7 @@ LogicalResult MaskState::parseDivsi(arith::DivSIOp divsiOp, assert(rhsIntAttr.has_value()); staticDiv = rhsIntAttr.value(); }else{ - divsiOp->emitError("MaskAnalysis: Static compilation cannot determine the value of this parameter"); + divsiOp->emitWarning("MaskAnalysis: Static compilation cannot determine the value of this parameter"); return failure(); } @@ -769,11 +928,11 @@ LogicalResult MaskState::parseDivsi(arith::DivSIOp divsiOp, if(info.isRealDim){ auto staticDim = getIntAttr(rhsState.scalar); if(!staticDim.has_value() || (staticDim.value() % staticDiv != 0 && staticDiv % staticDim.value() != 0)){ - divsiOp->emitError("MaskAnalysis: The shape of the mask is not divisible by the shape of the block"); + divsiOp->emitWarning("MaskAnalysis: The shape of the mask is not divisible by the shape of the block"); return failure(); } if(getIntAttr(info.shape).has_value() && getIntAttr(info.shape).value() != 0){ - divsiOp->emitError("MaskAnalysis: do not support div after remsi"); + divsiOp->emitWarning("MaskAnalysis: do not support div after remsi"); return failure(); } info.div = builder.getIndexAttr(staticDiv); diff --git a/lib/Analysis/OpFoldResultUtils.cpp b/lib/Analysis/OpFoldResultUtils.cpp index 95ad6f3b..b5e666cc 100644 --- a/lib/Analysis/OpFoldResultUtils.cpp +++ b/lib/Analysis/OpFoldResultUtils.cpp @@ -50,6 +50,26 @@ bool hasConstZero(const OpFoldResult ofr) { return false; } +std::optional getConstValue(const OpFoldResult ofr) { + auto intAttr = getIntAttr(ofr); + if (intAttr.has_value()) { + return intAttr.value(); + } + + auto val = dyn_cast(ofr); + assert(val); + auto constOp = val.getDefiningOp(); + if (!constOp) + return std::nullopt; + + intAttr = getIntAttr(constOp.getValue()); + if (intAttr.has_value()) { + return intAttr.value(); + } else { + return std::nullopt; + } +} + Value ofrToIndexValue(const OpFoldResult ofr, const Location loc, OpBuilder &b) { if (Value val = dyn_cast(ofr)) { diff --git a/lib/AnalysisStructured/PtrAnalysis.cpp b/lib/AnalysisStructured/PtrAnalysis.cpp index 1955b466..031fd82e 100644 --- a/lib/AnalysisStructured/PtrAnalysis.cpp +++ b/lib/AnalysisStructured/PtrAnalysis.cpp @@ -9,6 +9,7 @@ #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Operation.h" #include "mlir/IR/Value.h" #include "mlir/IR/Visitors.h" #include "mlir/Support/LLVM.h" @@ -258,7 +259,7 @@ LogicalResult PtrState::mulState(const PtrState &lhsState, return success(); } -tts::MakeTensorPtrOp PtrState::createTTSMakeTensorPtrOp(OpBuilder &builder, +mlir::Operation* PtrState::createTTSMakeTensorPtrOp(OpBuilder &builder, Location loc) { SmallVector staticSizes; for (size_t i = 0; i < getRank(); i++) { @@ -772,7 +773,7 @@ LogicalResult PtrAnalysis::rewriteAddptrOp(triton::AddPtrOp op) { if (isa(op.getPtr().getType())) { auto maketptrOp = state.createTTSMakeTensorPtrOp(builder, op.getLoc()); - ptrMap.map(op.getResult(), maketptrOp.getResult()); + ptrMap.map(op.getResult(), maketptrOp->getResult(0)); } else { // record the ptr as we have visited and built up the state for this scalar // pointer, which may be used by rewriteForOp later. @@ -810,7 +811,7 @@ LogicalResult PtrAnalysis::rewriteBitcastOp(triton::BitcastOp op) { knownPtrs[op.getResult()] = state; auto maketptrOp = state.createTTSMakeTensorPtrOp(builder, op.getLoc()); - ptrMap.map(op.getResult(), maketptrOp.getResult()); + ptrMap.map(op.getResult(), maketptrOp->getResult(0)); return success(); } @@ -825,7 +826,7 @@ LogicalResult PtrAnalysis::rewriteMakeTensorPtrOp(triton::MakeTensorPtrOp op) { auto maketptrOp = state.createTTSMakeTensorPtrOp(builder, op.getLoc()); knownPtrs[op.getResult()] = state; - ptrMap.map(op.getResult(), maketptrOp.getResult()); + ptrMap.map(op.getResult(), maketptrOp->getResult(0)); return success(); } @@ -867,7 +868,7 @@ LogicalResult PtrAnalysis::rewriteAdvanceOp(triton::AdvanceOp op) { auto newOp = state.createTTSMakeTensorPtrOp(builder, loc); knownPtrs[op.getResult()] = state; - ptrMap.map(op.getResult(), newOp.getResult()); + ptrMap.map(op.getResult(), newOp->getResult(0)); return success(); } @@ -1032,7 +1033,7 @@ LogicalResult PtrAnalysis::rewriteForOp(scf::ForOp op) { if (state->getRank() != 0) { OpBuilder builder(op.getRegion()); auto maketptrOp = state->createTTSMakeTensorPtrOp(builder, op.getLoc()); - ptrMap.map(arg, maketptrOp.getResult()); + ptrMap.map(arg, maketptrOp->getResult(0)); } } } diff --git a/lib/AnalysisStructured/PtrAnalysisTS.cpp b/lib/AnalysisStructured/PtrAnalysisTS.cpp index 313a8067..f4a919f7 100644 --- a/lib/AnalysisStructured/PtrAnalysisTS.cpp +++ b/lib/AnalysisStructured/PtrAnalysisTS.cpp @@ -101,6 +101,89 @@ bool PtrState::dimHasModulo(uint32_t dim) const { return intAttr.value() != 0; } +static Value applyUnstructuredMask(Operation *op, Value ptr, + triton::MaskState &mstate, Location loc, + OpBuilder builder) { + // structured mask + if (!mstate.hasNonContinuousDim()) + return ptr; + + auto gatherScatterPtr = + ptr.getDefiningOp(); + + // TODO: support unstructured mask for StructuredPtr + if (!gatherScatterPtr) + return nullptr; + + if (mstate.getRank() != 2 || mstate.getNonContinuousDim() != 0) + return nullptr; + + auto dims = mstate.dims; + auto nonContinuousDim = mstate.getNonContinuousDim().value(); + auto nonContinuousMask = dims[nonContinuousDim]; + + assert(nonContinuousDim == gatherScatterPtr.getGatherScatterDim()); + + // Clear the mask size for gather/scatter dim. + mstate.dims[nonContinuousDim] = OpFoldResult(builder.getI32IntegerAttr(0)); + return builder + .create( + loc, gatherScatterPtr.getBase(), + gatherScatterPtr.getGatherScatterOffset(), + dyn_cast(nonContinuousMask), + gatherScatterPtr.getGatherScatterDim(), gatherScatterPtr.getSizes(), + gatherScatterPtr.getMixedStrides(), + gatherScatterPtr.getMixedOffsets()) + .getResult(); +} + +bool isNotStructured(OpFoldResult offset) { + auto value = dyn_cast(offset); + return value && isa(value.getType()); +} + +bool PtrState::dimIsStructured(uint32_t dim) const { + assert(dim < getRank()); + + return !isNotStructured(offsets[dim]); +} + +int32_t PtrState::getNonStructuredDim() const { + SmallVector dims; + for (int32_t i = 0; i < getRank(); i++) { + if (dimIsStructured(i)) + continue; + dims.emplace_back(i); + } + assert(dims.size() == 1 && "must have single non-continuous dimension"); + return dims.front(); +} + +bool PtrState::isStructured() const { + return llvm::all_of( + offsets, [](OpFoldResult offset) { return !isNotStructured(offset); }); +} + +bool isSupportedStructuredPtr(Value ptr) { + + auto defOp = ptr.getDefiningOp(); + if (isa(defOp)) + return true; + + assert(isa(defOp)); + + auto unstructuredPtr = cast(defOp); + + auto ptrType = unstructuredPtr.getType(); + + auto tensorType = isa(ptrType) + ? cast(ptrType).getPointeeType() + : ptrType; + auto rank = cast(tensorType).getRank(); + + return rank == 2 && unstructuredPtr.getGatherScatterDim() == 0; +} + bool PtrState::isBlockPtr() const { return !order.empty(); } LogicalResult PtrState::addState(const PtrState &lhsState, @@ -110,9 +193,9 @@ LogicalResult PtrState::addState(const PtrState &lhsState, auto loc = op->getLoc(); if (lhsState.source && rhsState.source) { - op->emitRemark( + LLVM_DEBUG(op->emitRemark( "PtrAnalysis: do not support adding two pointer states that both " - "have base pointers"); + "have base pointers")); return failure(); } @@ -122,26 +205,75 @@ LogicalResult PtrState::addState(const PtrState &lhsState, auto addOp = builder.create(loc, lhsState.scalar, rhsState.scalar); scalar = addOp.getResult(); - } else if (lhsState.getRank() == 0) { // both lhs and rhs are scalars + } else if (lhsState.getRank() == 0) { + // both lhs and rhs are scalars. but only one of them has a scalar value + // Usually scalar: ptr add + // NOTE: rank > 0 && scalar: scalar has been added in offsets scalar = lhsState.scalar ? lhsState.scalar : rhsState.scalar; } + if (!lhsState.isStructured() && !rhsState.isStructured() && + lhsState.getNonStructuredDim() != rhsState.getNonStructuredDim()) { + LLVM_DEBUG( + op->emitRemark("PtrAnalysis: do not support adding two pointer states " + "that have different non-continuous dimension")); + return failure(); + } + for (uint64_t i = 0; i < lhsState.getRank(); i++) { - auto newOffset = - addOFRs(lhsState.offsets[i], rhsState.offsets[i], loc, builder); - offsets.push_back(newOffset); + sizes.push_back(lhsState.sizes[i]); - auto newStride = - addOFRs(lhsState.strides[i], rhsState.strides[i], loc, builder); - strides.push_back(newStride); + auto lhsStructured = lhsState.dimIsStructured(i); + auto rhsStructured = rhsState.dimIsStructured(i); + if (lhsStructured && rhsStructured) { + auto newOffset = + addOFRs(lhsState.offsets[i], rhsState.offsets[i], loc, builder); + offsets.push_back(newOffset); + auto newStride = + addOFRs(lhsState.strides[i], rhsState.strides[i], loc, builder); + strides.push_back(newStride); + continue; + } + assert(!(lhsState.hasModulo() && rhsState.hasModulo()) && + "PtrAnalysis: do not support adding two pointer states " + "that both have modulo"); + + strides.push_back(builder.getIndexAttr(1)); + if (!lhsStructured && !rhsStructured) { + auto newOffset = + builder + .create(loc, dyn_cast(lhsState.offsets[i]), + dyn_cast(rhsState.offsets[i])) + ->getResult(0); + offsets.push_back(newOffset); + continue; + } - sizes.push_back(lhsState.sizes[i]); + auto [lhsOffset, rhsOffset] = + lhsState.dimIsStructured(i) + ? std::tuple{ofrToIndexValue(lhsState.offsets[i], loc, + builder), + dyn_cast(rhsState.offsets[i])} + : std::tuple{ + ofrToIndexValue(rhsState.offsets[i], loc, builder), + dyn_cast(lhsState.offsets[i])}; + + auto tensorType = cast(rhsOffset.getType()); + lhsOffset = builder.create( + loc, tensorType.getElementType(), lhsOffset); + lhsOffset = builder.create(loc, tensorType, lhsOffset); + + auto newOffset = + builder.create(loc, lhsOffset, rhsOffset)->getResult(0); + offsets.push_back(newOffset); } // AddPtr where both lhs and rhs containing modulo operators not supported if (lhsState.hasModulo() && rhsState.hasModulo()) { - op->emitRemark("PtrAnalysis: do not support adding two pointer states " - "that both have modulo"); + // May modulo different value + LLVM_DEBUG( + op->emitRemark("PtrAnalysis: do not support adding two pointer states " + "that both have modulo")); return failure(); } @@ -188,17 +320,17 @@ LogicalResult PtrState::addState(const PtrState &lhsState, } else if (i == 0 && lhs->getRank() == 2 && rhs->scalar) { shape.push_back(lhs->shape[1]); shape.push_back(lhs->shape[0]); - op->emitWarning( + LLVM_DEBUG(op->emitWarning( "PtrAnalysis: allowing adding pointer state with modulo in dim 0 to " "another pointer state with offset in dim 0.\nPlease verify the " "operand that contains a scalar is meant to increment pointers in " "dim1. If that is not the case it WILL LEAD TO WRONG COMPILATION " "RESULTS.\n\nTo avoid this warning, use expand_dims (instead of " - "splat) to explicitly specify which dimension contains the scalar."); + "splat) to explicitly specify which dimension contains the scalar.")); break; } else { - op->emitRemark( - "PtrAnalysis: do not support adding to operand with modulo"); + LLVM_DEBUG(op->emitRemark( + "PtrAnalysis: do not support adding to operand with modulo")); return failure(); } } @@ -215,16 +347,21 @@ void PtrState::dump() const { llvm::dbgs() << "scalar: " << scalar << "\n"; } - llvm::dbgs() << "offsets: "; + llvm::dbgs() << "offsets:\n"; llvm::interleave(offsets, llvm::dbgs(), "\n"); - llvm::dbgs() << "\nstrides: "; + llvm::dbgs() << "\nstrides:\n"; llvm::interleave(strides, llvm::dbgs(), "\n"); - llvm::dbgs() << "\nsizes: "; + llvm::dbgs() << "\nsizes:\n"; llvm::interleave(sizes, llvm::dbgs(), "\n"); - llvm::dbgs() << "\nshape: "; + llvm::dbgs() << "\nshape:\n"; llvm::interleave(shape, llvm::dbgs(), "\n"); - llvm::dbgs() << "\norder: "; + llvm::dbgs() << "\norder:\n"; llvm::interleave(order, llvm::dbgs(), "\n"); + + isStructured() + ? llvm::dbgs() << "structured\n" + : llvm::dbgs() << "dim " << getNonStructuredDim() << " not structured\n"; + llvm::dbgs() << "\n"; } @@ -238,15 +375,16 @@ LogicalResult PtrState::mulState(const PtrState &lhsState, // neither lhs nor rhs should have source, since multiplying base pointer // does not make sense if (lhsState.source && rhsState.source) { - op->emitRemark("PtrAnalysis: do not support multiplying base pointers"); + LLVM_DEBUG(op->emitRemark( + "PtrAnalysis: do not support multiplying base pointers")); return failure(); } // currently do not support both tensors are effectively non-scalar if (!lhsState.scalar && !rhsState.scalar) { - op->emitRemark( + LLVM_DEBUG(op->emitRemark( "PtrAnalysis: only support multiplying pointer states when one of " - "them represent a scalar"); + "them represent a scalar")); return failure(); } @@ -262,23 +400,44 @@ LogicalResult PtrState::mulState(const PtrState &lhsState, builder.create(loc, lhsState.scalar, rhsState.scalar); } + assert(rhs->scalar && "rhs always scalar"); for (uint64_t i = 0; i < lhs->sizes.size(); i++) { - OpFoldResult newOffset = - mulOFRValue(lhs->offsets[i], rhs->scalar, loc, builder); - OpFoldResult newStride = - mulOFRValue(lhs->strides[i], rhs->scalar, loc, builder); OpFoldResult newShape = mulOFRValue(lhs->shape[i], rhs->scalar, loc, builder); - offsets.push_back(newOffset); - strides.push_back(newStride); shape.push_back(newShape); sizes.push_back(lhs->sizes[i]); + + if (lhs->dimIsStructured(i)) { + auto newOffset = mulOFRValue(lhs->offsets[i], rhs->scalar, loc, builder); + offsets.push_back(newOffset); + auto newStride = mulOFRValue(lhs->strides[i], rhs->scalar, loc, builder); + strides.push_back(newStride); + continue; + } + + assert(!lhs->hasModulo() && + "should not have non-structured dimension with modulo"); + + // NOTE: Can also consider mul scalar on stride + auto lhsOffset = dyn_cast(lhs->offsets[i]); + + auto tensorType = cast(lhsOffset.getType()); + auto rhsScalar = builder.create( + loc, tensorType.getElementType(), rhs->scalar); + auto rhsSplatTensor = + builder.create(loc, tensorType, rhsScalar); + OpFoldResult newOffset = + builder.create(loc, lhsOffset, rhsSplatTensor) + ->getResult(0); + + offsets.push_back(newOffset); + strides.push_back(lhs->strides[i]); } if (rhs->hasModulo()) { - op->emitRemark( + LLVM_DEBUG(op->emitRemark( "PtrAnalysis: do not support multiplying pointer states that has " - "modulos"); + "modulos")); return failure(); } @@ -292,26 +451,10 @@ LogicalResult PtrState::subState(const PtrState &lhsState, auto loc = op->getLoc(); if (lhsState.source && rhsState.source) { - if (lhsState.source != rhsState.source) { - op->emitRemark( - "PtrAnalysis: subtracting pointers from different bases is not supported"); - return failure(); - } - - if (lhsState.scalar && rhsState.scalar) { - auto subOp = builder.create(loc, lhsState.scalar, rhsState.scalar); - scalar = subOp.getResult(); - } else if (lhsState.scalar) { - scalar = lhsState.scalar; - } else if (rhsState.scalar) { - auto zero = builder.create(loc, 0, - rhsState.scalar.getType()); - auto negOp = builder.create(loc, zero, rhsState.scalar); - scalar = negOp.getResult(); - } - - source = nullptr; - return success(); + LLVM_DEBUG(op->emitRemark( + "PtrAnalysis: do not support subtract two pointer states that both " + "have base pointers")); + return failure(); } if (!lhsState.source && rhsState.source) { @@ -319,20 +462,28 @@ LogicalResult PtrState::subState(const PtrState &lhsState, return failure(); } - source = lhsState.source ? lhsState.source : rhsState.source; + // WORKAROUND: avoid other non-structured cases. + if (!lhsState.isStructured() || !rhsState.isStructured()) { + return failure(); + } + + source = lhsState.source; if (lhsState.scalar && rhsState.scalar) { - auto subOp = builder.create(loc, lhsState.scalar, rhsState.scalar); + auto subOp = + builder.create(loc, lhsState.scalar, rhsState.scalar); scalar = subOp.getResult(); } else if (lhsState.getRank() == 0) { // both lhs and rhs are scalars scalar = lhsState.scalar ? lhsState.scalar : rhsState.scalar; } for (uint64_t i = 0; i < lhsState.getRank(); i++) { - auto newOffset = subOFRs(lhsState.offsets[i], rhsState.offsets[i], loc, builder); + auto newOffset = + subOFRs(lhsState.offsets[i], rhsState.offsets[i], loc, builder); offsets.push_back(newOffset); - auto newStride = subOFRs(lhsState.strides[i], rhsState.strides[i], loc, builder); + auto newStride = + subOFRs(lhsState.strides[i], rhsState.strides[i], loc, builder); strides.push_back(newStride); sizes.push_back(lhsState.sizes[i]); @@ -345,8 +496,8 @@ LogicalResult PtrState::subState(const PtrState &lhsState, return success(); } -tts::MakeTensorPtrOp PtrState::createTTSMakeTensorPtrOp(OpBuilder &builder, - Location loc) { +Operation *PtrState::createTTSMakeTensorPtrOp(OpBuilder &builder, + Location loc) { SmallVector staticSizes; for (size_t i = 0; i < getRank(); i++) { auto s = getIntAttr(sizes[i]); @@ -354,11 +505,27 @@ tts::MakeTensorPtrOp PtrState::createTTSMakeTensorPtrOp(OpBuilder &builder, staticSizes.push_back(s.value()); } - auto op = builder.create( + if (!isStructured()) { + int nonContinuousDim = getNonStructuredDim(); + Value nonContinuousOffset = cast(offsets[nonContinuousDim]); + auto op = builder.create( + loc, source, nonContinuousOffset, nonContinuousDim, staticSizes, + strides, offsets); + + LLVM_DEBUG({ + llvm::dbgs() << "creating tts::make_gather_scatter_tensor_ptr:\n"; + op->dump(); + op->getParentOfType()->dump(); + }); + return op; + } + + auto op = builder.create( loc, source, staticSizes, strides, offsets, shape, order); LLVM_DEBUG({ llvm::dbgs() << "creating tts::make_tensor_ptr:\n"; op->dump(); + op->getParentOfType()->dump(); }); return op; @@ -379,8 +546,8 @@ LogicalResult PtrAnalysis::visitOperandAdd(arith::AddIOp addOp, PtrState &state, // Checking for higher dimension is done in addState below if ((lhsState.getRank() == 1 && lhsState.hasModulo()) || (rhsState.getRank() == 1 && rhsState.hasModulo())) { - addOp->emitRemark( - "PtrAnalysis: do not support this pattern: a + arange(0, K) % M"); + LLVM_DEBUG(addOp->emitRemark( + "PtrAnalysis: do not support this pattern: a + arange(0, K) % M")); return failure(); } @@ -401,7 +568,8 @@ LogicalResult PtrAnalysis::visitOperandSub(arith::SubIOp subOp, PtrState &state, } if (lhsState.hasModulo() || rhsState.hasModulo()) { - subOp->emitRemark("PtrAnalysis: do not support modulo for subi op\n"); + LLVM_DEBUG( + subOp->emitRemark("PtrAnalysis: do not support modulo for subi op\n")); return failure(); } @@ -432,6 +600,94 @@ LogicalResult PtrAnalysis::visitOperandMul(arith::MulIOp mulOp, PtrState &state, return state.mulState(lhsState, rhsState, mulOp, builder); } +LogicalResult PtrAnalysis::visitOperandDivSI(arith::DivSIOp divOp, + PtrState &state, + const Location loc, + OpBuilder &builder) { + assert(state.isEmpty()); + + // Divisor + PtrState rhsState; + if (visitOperand(divOp.getRhs(), rhsState, loc, builder).failed()) { + return failure(); + } + + if (!rhsState.scalar) { + LLVM_DEBUG(divOp->emitRemark( + "PtrAnalysis: only support cases when divisor contains scalar")); + return failure(); + } + + // Check if divisor is a constant value + auto rhsInt = getConstValue(rhsState.scalar); + if (!rhsInt.has_value()) { + LLVM_DEBUG(divOp->emitRemark( + "PtrAnalysis: only support cases when divisor is a constant")); + return failure(); + } + + int64_t divisor = rhsInt.value(); + assert(divisor != 0 && "Divisor cannot be 0"); + + if (visitOperand(divOp.getLhs(), state, loc, builder).failed()) { + return failure(); + } + + // Check if dividend is 1D or 2D with one dimension being 1 + auto resultType = cast(divOp.getType()); + if (resultType.getRank() > 2) { + LLVM_DEBUG(divOp->emitRemark( + "PtrAnalysis: only support 1D or 2D tensors for division")); + return failure(); + } + + if (state.hasModulo()) { + LLVM_DEBUG(divOp->emitRemark( + "PtrAnalysis: do not support division on a state with modulo (e.g. " + "(x % N) / M pattern)")); + return failure(); + } + // Apply division to offsets and strides + auto constDiv = + builder.create(loc, builder.getIndexAttr(divisor)); + + if (state.scalar) { + state.scalar = + builder + .create( + loc, ofrToIndexValue(state.scalar, loc, builder), constDiv) + .getResult(); + return failure(); + } + + auto shape = resultType.getShape(); + if (resultType.getRank() == 2 && shape[1] != 1) { + LLVM_DEBUG(divOp->emitRemark("PtrAnalysis: only support outernest dim " + "division, otherwise unstructured")); + return failure(); + } + + for (uint32_t i = 0; i < state.getRank(); i++) { + if (shape[i] == 1) { + assert(hasConstZero(state.offsets[i])); + state.offsets[i] = builder.getIndexAttr(0); + state.strides[i] = builder.getIndexAttr(0); + continue; + } + // Divide offsets by divisor + state.offsets[i] = + resultType.getRank() == 2 + ? builder + .create( + loc, divOp, SmallVector{{0, 1}}) + ->getResult(0) + : divOp->getResult(0); + state.strides[i] = builder.getIndexAttr(1); + } + + return success(); +} + LogicalResult PtrAnalysis::visitOperandRem(arith::RemSIOp remOp, PtrState &state, const Location loc, OpBuilder &builder) { @@ -443,8 +699,9 @@ LogicalResult PtrAnalysis::visitOperandRem(arith::RemSIOp remOp, } if (!rhsState.scalar) { - remOp->emitRemark("PtrAnalysis: only support cases when rhs of remainder " - "contains scalar"); + LLVM_DEBUG(remOp->emitRemark( + "PtrAnalysis: only support cases when rhs of remainder " + "contains scalar")); return failure(); } @@ -456,8 +713,8 @@ LogicalResult PtrAnalysis::visitOperandRem(arith::RemSIOp remOp, // would have already populated the modulo states after visiting the lhs. // Assert that all the modulo states are empty. if (state.hasModulo()) { - remOp->emitRemark( - "PtrAnalysis: do not support multiple modulo within an expression"); + LLVM_DEBUG(remOp->emitRemark( + "PtrAnalysis: do not support multiple modulo within an expression")); return failure(); } @@ -481,13 +738,13 @@ LogicalResult PtrAnalysis::visitOperandRem(arith::RemSIOp remOp, } else if (shape[1] == 1) { state.shape[0] = rhsState.scalar; } else { - remOp->emitRemark( + LLVM_DEBUG(remOp->emitRemark( "PtrAnalysis: taking modulo on a 2D tensor with no singleton " - "dimension not supported"); + "dimension not supported")); return failure(); } } else { - remOp->emitRemark("PtrAnalysis: unsupported modulo pattern"); + LLVM_DEBUG(remOp->emitRemark("PtrAnalysis: unsupported modulo pattern")); return failure(); } return success(); @@ -545,9 +802,9 @@ PtrAnalysis::visitOperandExpandDims(triton::ExpandDimsOp expandDimsOp, state.shape.insert(state.shape.begin() + axis, builder.getIndexAttr(0)); if (state.hasModulo() && state.getRank() > 2) { - expandDimsOp->emitRemark( + LLVM_DEBUG(expandDimsOp->emitRemark( "PtrAnalysis: unsupported scenario where expand_dims result " - "has modulo and rank > 2"); + "has modulo and rank > 2")); return failure(); } @@ -564,7 +821,8 @@ PtrAnalysis::visitOperandBroadcast(triton::BroadcastOp broadcastOp, auto dst = broadcastOp.getResult(); if (!isa(src.getType())) { - broadcastOp->emitRemark("PtrAnalysis: Unsupported broadcast source type"); + LLVM_DEBUG(broadcastOp->emitRemark( + "PtrAnalysis: Unsupported broadcast source type")); return failure(); } @@ -612,7 +870,7 @@ LogicalResult PtrAnalysis::visitOperandSplat(triton::SplatOp splatOp, state.shape.push_back(builder.getIndexAttr(0)); } } else { - splatOp->emitRemark("PtrAnalysis: unsupported splat pattern"); + LLVM_DEBUG(splatOp->emitRemark("PtrAnalysis: unsupported splat pattern")); return failure(); } @@ -622,8 +880,9 @@ LogicalResult PtrAnalysis::visitOperandSplat(triton::SplatOp splatOp, state.offsets[0] = state.scalar; if (state.hasModulo() && state.getRank() > 2) { - splatOp->emitRemark("PtrAnalysis: unsupported scenario where splat result " - "has modulo and rank > 2"); + LLVM_DEBUG(splatOp->emitRemark( + "PtrAnalysis: unsupported scenario where splat result " + "has modulo and rank > 2")); return failure(); } @@ -639,7 +898,6 @@ LogicalResult PtrAnalysis::visitOperandAddptr(triton::AddPtrOp addptrOp, PtrState ptrState; if (visitOperand(addptrOp.getPtr(), ptrState, addptrOp.getLoc(), builder) .failed()) { - // assert(0); return failure(); } @@ -716,8 +974,8 @@ PtrAnalysis::visitOperandMakeTensorPtr(triton::MakeTensorPtrOp makeTPtrOp, state.source = makeTPtrOp.getBase(); if (makeTPtrOp.getOrder().empty()) { - makeTPtrOp->emitRemark( - "PtrAnalysis: expect tt.make_tensor_ptr to have order field set"); + LLVM_DEBUG(makeTPtrOp->emitRemark( + "PtrAnalysis: expect tt.make_tensor_ptr to have order field set")); return failure(); } @@ -760,9 +1018,9 @@ LogicalResult PtrAnalysis::visitOperandForOp(scf::ForOp forOp, Value operand, auto newState = getLoopResultPtrState(forOp, index); if (failed(newState)) { - forOp.emitError( + LLVM_DEBUG(forOp.emitError( "Rewrite for-op failed. Could not find PtrState returned by " - "the loop."); + "the loop.")); return failure(); } @@ -770,6 +1028,14 @@ LogicalResult PtrAnalysis::visitOperandForOp(scf::ForOp forOp, Value operand, return success(); } +LogicalResult PtrAnalysis::visitOperandIntToPtr(triton::IntToPtrOp op, + PtrState &state, + const Location loc, + OpBuilder &builder) { + state.source = op.getResult(); + return success(); +} + LogicalResult PtrAnalysis::visitOperandBitcast(triton::BitcastOp op, PtrState &state, const Location loc, @@ -814,17 +1080,20 @@ LogicalResult PtrAnalysis::visitOperand(Value operand, PtrState &state, builder); } else if (auto castOp = dyn_cast(op)) { return visitOperandBitcast(castOp, state, loc, builder); + } else if (auto intToPtrOp = dyn_cast(op)) { + return visitOperandIntToPtr(intToPtrOp, state, loc, builder); } else if (auto makeTensorOp = dyn_cast(op)) { llvm_unreachable("Unexpected operand defining operation tts.make_tptr"); } else if (auto selectOp = dyn_cast(op)) { state.source = selectOp.getResult(); return success(); } else { - op->emitRemark("Unexpected operand defining operation"); + LLVM_DEBUG(op->emitRemark("Unexpected operand defining operation")); return failure(); } } else { state.source = operand; + state.scalar = builder.create(loc, 0); return success(); } } @@ -847,6 +1116,8 @@ LogicalResult PtrAnalysis::visitOperand(Value operand, PtrState &state, return visitOperandAddptr(op, state, loc, builder); } else if (auto op = operand.getDefiningOp()) { return visitOperandConstSplat(op, state, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return visitOperandDivSI(op, state, loc, builder); } else if (auto op = operand.getDefiningOp()) { return visitOperandRem(op, state, loc, builder); } else if (auto op = operand.getDefiningOp()) { @@ -864,9 +1135,24 @@ LogicalResult PtrAnalysis::visitOperand(Value operand, PtrState &state, state = knownPtrs[operand]; return success(); } else { - llvm::dbgs() << "PtrAnalysis: encountered addptr operand produced by an " - "unsupported operation\n"; - operand.dump(); + auto resultType = operand.getType(); + if (isa(resultType) && + cast(resultType).getRank() == 1) { + assert(state.isEmpty()); + LLVM_DEBUG(llvm::dbgs() << "\n unsupported: PtrAnalysis directly handle " + "ranked-1 tensors as offset\n";); + auto shape = cast(resultType).getShape(); + state.offsets.push_back(operand); + state.sizes.push_back(builder.getIndexAttr(shape.front())); + state.strides.push_back(builder.getIndexAttr(1)); + state.shape.push_back(builder.getIndexAttr(0)); + return success(); + } + + LLVM_DEBUG(llvm::dbgs() + << "PtrAnalysis: encountered addptr operand produced by an " + "unsupported operation\n"; + operand.dump()); return failure(); } } @@ -883,7 +1169,7 @@ LogicalResult PtrAnalysis::rewriteAddptrOp(triton::AddPtrOp op) { if (isa(op.getPtr().getType())) { auto maketptrOp = state.createTTSMakeTensorPtrOp(builder, op.getLoc()); - ptrMap.map(op.getResult(), maketptrOp.getResult()); + ptrMap.map(op.getResult(), maketptrOp->getResult(0)); } else { // record the ptr as we have visited and built up the state for this scalar // pointer, which may be used by rewriteForOp later. @@ -897,6 +1183,10 @@ LogicalResult PtrAnalysis::rewriteBitcastOp(triton::BitcastOp op) { llvm::dbgs() << "Rewriting bitcast op:\n"; op->dump(); }); + + if (!triton::isPtrTypeLike(op.getOperand().getType())) + return failure(); + OpBuilder builder(op); PtrState state; @@ -915,7 +1205,7 @@ LogicalResult PtrAnalysis::rewriteBitcastOp(triton::BitcastOp op) { builder.create(op->getLoc(), resType, state.source); state.source = newBitcast; auto maketptrOp = state.createTTSMakeTensorPtrOp(builder, op.getLoc()); - ptrMap.map(op.getResult(), maketptrOp.getResult()); + ptrMap.map(op.getResult(), maketptrOp->getResult(0)); } else { ptrMap.map(op.getResult(), op.getResult()); } @@ -932,7 +1222,7 @@ LogicalResult PtrAnalysis::rewriteMakeTensorPtrOp(triton::MakeTensorPtrOp op) { auto maketptrOp = state.createTTSMakeTensorPtrOp(builder, op.getLoc()); knownPtrs[op.getResult()] = state; - ptrMap.map(op.getResult(), maketptrOp.getResult()); + ptrMap.map(op.getResult(), maketptrOp->getResult(0)); return success(); } @@ -942,7 +1232,8 @@ LogicalResult PtrAnalysis::rewriteAdvanceOp(triton::AdvanceOp op) { PtrState state; if (visitOperand(op->getOperand(0), state, loc, builder).failed()) { - op->emitRemark("PtrAnalysis: Failed to analyze ptr of tt.advance"); + LLVM_DEBUG( + op->emitRemark("PtrAnalysis: Failed to analyze ptr of tt.advance")); return failure(); } assert(state.isBlockPtr() && @@ -974,7 +1265,7 @@ LogicalResult PtrAnalysis::rewriteAdvanceOp(triton::AdvanceOp op) { auto newOp = state.createTTSMakeTensorPtrOp(builder, loc); knownPtrs[op.getResult()] = state; - ptrMap.map(op.getResult(), newOp.getResult()); + ptrMap.map(op.getResult(), newOp->getResult(0)); return success(); } @@ -1048,13 +1339,30 @@ PtrState PtrAnalysis::reconcileLoopPtrState( PtrState newState = state; int cnt = iterArgIndex + 1; if (newState.getRank() == 0) { - assert(newState.scalar); - // for scalar pointers, the scalar contains the offset and is the only - // relevant newState that could be updated by the loop. + // rewriteGetStructuredStateOp will return a scalar constant 0 in the + // newState.scalar for rank 0 case. Therefore, we can always get the correct + // newState.scalar by calling getReplacementVal with the iterArgIndex of the + // current loop. + // For scalar pointers, the scalar contains the offset and is + // the only relevant newState that could be updated by the loop. newState.scalar = getReplacementVal(forOp, cnt); } else { - for (auto &offset : newState.offsets) { - offset = getReplacementVal(forOp, cnt++); + int rank = newState.getRank(); + + std::optional nonStructuredDim = + newState.isStructured() + ? std::nullopt + : std::optional(newState.getNonStructuredDim()); + for (int i = 0; i < rank; i++) { + // Unstructured dim will be tensor type which is not compatible with index + // type of tts::getStructuredStateop + if (nonStructuredDim && i == nonStructuredDim.value()) { + forOp.getRegionIterArg(cnt).setType( + dyn_cast(newState.offsets[i]).getType()); + forOp->getResult(cnt).setType( + dyn_cast(newState.offsets[i]).getType()); + } + newState.offsets[i] = getReplacementVal(forOp, cnt++); } for (auto &stride : newState.strides) { @@ -1072,6 +1380,13 @@ FailureOr PtrAnalysis::getLoopIterArgPtrState(scf::ForOp forOp, return failure(); } + // WORKAROUND: avoid other non-structured cases. Since unstructured offset + // need ShapeType which is not compatible with iterarg (index type) + if (!state->isStructured() && + (state->getRank() != 2 || state->getNonStructuredDim() != 0)) { + return failure(); + } + return reconcileLoopPtrState( forOp, index, state.value(), [](scf::ForOp op, size_t index) { return op.getRegionIterArg(index); }); @@ -1084,6 +1399,13 @@ FailureOr PtrAnalysis::getLoopResultPtrState(scf::ForOp forOp, return failure(); } + // WORKAROUND: avoid other non-structured cases. Since unstructured offset + // need ShapeType which is not compatible with result (index type) + if (!state->isStructured() && + (state->getRank() != 2 || state->getNonStructuredDim() != 0)) { + return failure(); + } + return reconcileLoopPtrState( forOp, index, state.value(), [](scf::ForOp op, size_t index) { return op->getResult(index); }); @@ -1101,9 +1423,9 @@ LogicalResult PtrAnalysis::rewriteForOp(scf::ForOp op) { // considered structured by PtrAnalysis, failing to retrieve the PtrState // should not fail the rewrite process. // We emit an error for diagnostics and debugging purposes. - op->emitWarning( + LLVM_DEBUG(op->emitWarning( "Rewrite for-op failed. Could not find PtrState for iter-arg index " + - std::to_string(i)); + std::to_string(i))); continue; } @@ -1135,19 +1457,17 @@ LogicalResult PtrAnalysis::rewriteForOp(scf::ForOp op) { // original pointer. // Note that there can be tensor of indices in iter-arg, so we only create // the make_tensor_ptr op when the arg is of pointer type. - if (isPointerType(arg.getType())) { - if (state->getRank() != 0) { - OpBuilder builder(op.getRegion()); - auto maketptrOp = state->createTTSMakeTensorPtrOp(builder, op.getLoc()); - ptrMap.map(arg, maketptrOp.getResult()); - } + if (isPointerType(arg.getType()) && state->getRank() != 0) { + OpBuilder builder(op.getRegion()); + auto maketptrOp = state->createTTSMakeTensorPtrOp(builder, op.getLoc()); + ptrMap.map(arg, maketptrOp->getResult(0)); } } // Recursively rewrite the inner ops if (rewriteOp(op).failed()) { - op->emitRemark( - "PtrAnalysis: update loop body failed when rewriting for op"); + LLVM_DEBUG(op->emitRemark( + "PtrAnalysis: update loop body failed when rewriting for op")); return failure(); } @@ -1164,8 +1484,8 @@ PtrAnalysis::rewriteGetStructuredStateOp(tts::GetStructuredStateOp op) { // analyze this pointer. In such cases, simply remap all uses of the // structured value back to its original triton value. if (!knownPtrs.contains(tritonValue)) { - op.emitRemark( - "Rewrite GetStructuredStateOp failed. Could not find PtrState."); + LLVM_DEBUG(op.emitRemark( + "Rewrite GetStructuredStateOp failed. Could not find PtrState.")); auto numResults = op.getNumResults(); SmallVector replacements( numResults, builder.create(op.getLoc(), @@ -1176,6 +1496,20 @@ PtrAnalysis::rewriteGetStructuredStateOp(tts::GetStructuredStateOp op) { } tts::PtrState state = knownPtrs[tritonValue]; + + // WORKAROUND: avoid other non-structured cases. Since unstructured offset + // need ShapeType which is not compatible with iterarg (index type) + if (!state.isStructured() && + (state.getRank() != 2 || state.getNonStructuredDim() != 0)) { + auto numResults = op.getNumResults(); + SmallVector replacements( + numResults, builder.create(op.getLoc(), + builder.getIndexAttr(0))); + replacements.front() = tritonValue; + op.getResults().replaceAllUsesWith(replacements); + return failure(); + } + Value remappedValue = ptrMap.contains(tritonValue) ? ptrMap.lookup(tritonValue) : tritonValue; @@ -1189,7 +1523,6 @@ PtrAnalysis::rewriteGetStructuredStateOp(tts::GetStructuredStateOp op) { } else { // This operand is a pointer directly from the kernel arguments. // Use offset 0. - assert(!tritonValue.getDefiningOp()); replacements.push_back(builder.create( op.getLoc(), builder.getIndexAttr(0))); } @@ -1230,14 +1563,20 @@ LogicalResult PtrAnalysis::rewriteLoadOp(triton::LoadOp op, auto loc = op.getLoc(); if (!ptr) { - op->emitRemark("PtrAnalysis: pointer is not replace with tts.make_tptr so " - "loadOp cannot be rewritten"); + LLVM_DEBUG(op->emitRemark( + "PtrAnalysis: pointer is not replace with tts.make_tptr so " + "loadOp cannot be rewritten")); return failure(); } auto ptrType = dyn_cast(ptr.getType()); if (ptrType && !isa(ptrType.getPointeeType())) { - op->emitRemark("PtrAnalysis: scalar loadOp will not be rewritten"); + LLVM_DEBUG( + op->emitRemark("PtrAnalysis: scalar loadOp will not be rewritten")); + return failure(); + } + + if (!isSupportedStructuredPtr(ptr)) { return failure(); } @@ -1250,14 +1589,21 @@ LogicalResult PtrAnalysis::rewriteLoadOp(triton::LoadOp op, // are moving. if (mask) { if (mstate.parse(mask, loc, builder).failed()) { - op->emitRemark("MaskAnalysis failed"); + LLVM_DEBUG(op->emitRemark("MaskAnalysis failed")); return failure(); } + + ptr = applyUnstructuredMask(op, ptr, mstate, loc, builder); + if (!ptr) + return failure(); + dims = mstate.dims; } auto boundaryCheck = op.getBoundaryCheck(); if (!boundaryCheck.empty()) { + assert(dims.empty() && "Mask and boundary check cannot be used together"); + boundaryCheckToMaskDim(builder, loc, knownPtrs.at(op.getPtr()), boundaryCheck, mstate); dims = mstate.dims; @@ -1268,8 +1614,8 @@ LogicalResult PtrAnalysis::rewriteLoadOp(triton::LoadOp op, scalarOther = triton::getScalarValue(other, loc, builder); if (!scalarOther) { - op->emitRemark("other value used in masked load produced by " - "unsupported instruction"); + LLVM_DEBUG(op->emitRemark("other value used in masked load produced by " + "unsupported instruction")); return failure(); } } @@ -1279,6 +1625,7 @@ LogicalResult PtrAnalysis::rewriteLoadOp(triton::LoadOp op, LLVM_DEBUG({ llvm::dbgs() << "creating tts::load:\n"; loadOp->dump(); + loadOp->getParentOfType()->dump(); }); op.replaceAllUsesWith(loadOp.getResult()); @@ -1374,14 +1721,20 @@ LogicalResult PtrAnalysis::rewriteStoreOp(triton::StoreOp op, auto loc = op.getLoc(); if (!ptr) { - op->emitRemark("PtrAnalysis: pointer is not replace with tts.make_tptr so " - "storeOp cannot be rewritten"); + LLVM_DEBUG(op->emitRemark( + "PtrAnalysis: pointer is not replace with tts.make_tptr so " + "storeOp cannot be rewritten")); return failure(); } auto ptrType = dyn_cast(ptr.getType()); if (ptrType && !isa(ptrType.getPointeeType())) { - op->emitRemark("PtrAnalysis: scalar storeOp will not be rewritten"); + LLVM_DEBUG( + op->emitRemark("PtrAnalysis: scalar storeOp will not be rewritten")); + return failure(); + } + + if (!isSupportedStructuredPtr(ptr)) { return failure(); } @@ -1394,14 +1747,21 @@ LogicalResult PtrAnalysis::rewriteStoreOp(triton::StoreOp op, // are moving. if (mask) { if (mstate.parse(mask, loc, builder).failed()) { - op->emitRemark("MaskAnalysis failed"); + LLVM_DEBUG(op->emitRemark("MaskAnalysis failed")); return failure(); } + + ptr = applyUnstructuredMask(op, ptr, mstate, loc, builder); + if (!ptr) + return failure(); + dims = mstate.dims; } auto boundaryCheck = op.getBoundaryCheck(); if (!boundaryCheck.empty()) { + assert(dims.empty() && "Mask and boundary check cannot be used together"); + boundaryCheckToMaskDim(builder, loc, knownPtrs.at(op.getPtr()), boundaryCheck, mstate); dims = mstate.dims; @@ -1412,6 +1772,7 @@ LogicalResult PtrAnalysis::rewriteStoreOp(triton::StoreOp op, LLVM_DEBUG({ llvm::dbgs() << "creating tts::store:\n"; storeOp->dump(); + storeOp->getParentOfType()->dump(); }); op->erase(); @@ -1439,6 +1800,10 @@ LogicalResult PtrAnalysis::rewriteAtomicRMWOp(triton::AtomicRMWOp op, return failure(); } + if (!isSupportedStructuredPtr(ptr)) { + return failure(); + } + ArrayRef dims; mlir::triton::MaskState mstate(useUnsafeMask); @@ -1451,6 +1816,11 @@ LogicalResult PtrAnalysis::rewriteAtomicRMWOp(triton::AtomicRMWOp op, LLVM_DEBUG(op->emitRemark("MaskAnalysis failed")); return failure(); } + + ptr = applyUnstructuredMask(op, ptr, mstate, loc, builder); + if (!ptr) + return failure(); + dims = mstate.dims; } @@ -1489,6 +1859,10 @@ LogicalResult PtrAnalysis::rewriteAtomicCASOp(triton::AtomicCASOp op) { return failure(); } + if (!isSupportedStructuredPtr(ptr)) { + return failure(); + } + OpBuilder builder(op); auto atomicCASOp = builder.create( @@ -1518,44 +1892,49 @@ LogicalResult PtrAnalysis::rewriteOp(Operation *rootOp, bool useUnsafeMask) { return TypeSwitch(op) .Case([&](auto addptr) { if (rewriteAddptrOp(addptr).failed()) { - addptr->emitRemark("PtrAnalysis: Failed to rewrite AddPtrOp"); + LLVM_DEBUG( + addptr->emitRemark("PtrAnalysis: Failed to rewrite AddPtrOp")); } return WalkResult::advance(); }) .Case([&](auto bitcast) { if (rewriteBitcastOp(bitcast).failed()) { - bitcast->emitRemark("PtrAnalysis: Failed to rewrite BitcastOp"); + LLVM_DEBUG(bitcast->emitRemark( + "PtrAnalysis: Failed to rewrite BitcastOp")); } return WalkResult::advance(); }) .Case([&](auto maketptr) { if (rewriteMakeTensorPtrOp(maketptr).failed()) { - maketptr->emitRemark( - "PtrAnalysis: Failed to rewrite MakeTensorPtrOp"); + LLVM_DEBUG(maketptr->emitRemark( + "PtrAnalysis: Failed to rewrite MakeTensorPtrOp")); } return WalkResult::advance(); }) .Case([&](auto advance) { if (rewriteAdvanceOp(advance).failed()) { - advance->emitRemark("PtrAnalysis: Failed to rewrite AdvanceOp"); + LLVM_DEBUG(advance->emitRemark( + "PtrAnalysis: Failed to rewrite AdvanceOp")); } return WalkResult::advance(); }) .Case([&](auto load) { if (rewriteLoadOp(load, useUnsafeMask).failed()) { - load->emitRemark("PtrAnalysis: Failed to rewrite LoadOp"); + LLVM_DEBUG( + load->emitRemark("PtrAnalysis: Failed to rewrite LoadOp")); return WalkResult::advance(); } return WalkResult::skip(); }) .Case([&](auto store) { if (rewriteStoreOp(store, useUnsafeMask).failed()) { - store->emitRemark("PtrAnalysis: Failed to rewrite StoreOp"); + LLVM_DEBUG( + store->emitRemark("PtrAnalysis: Failed to rewrite StoreOp")); return WalkResult::advance(); } return WalkResult::skip(); }) - .Case([&](auto atomicRMW) { + .Case([&](auto atomicRMW) { if (rewriteAtomicRMWOp(atomicRMW, useUnsafeMask).failed()) { LLVM_DEBUG(atomicRMW->emitRemark( "PtrAnalysis: Failed to rewrite AtomicRMWOp")); @@ -1577,7 +1956,8 @@ LogicalResult PtrAnalysis::rewriteOp(Operation *rootOp, bool useUnsafeMask) { // that the the walk does not visit the for-op's child operations // the second time. if (rewriteForOp(forOp).failed()) { - forOp->emitRemark("PtrAnalysis: Failed to rewrite ForOp"); + LLVM_DEBUG( + forOp->emitRemark("PtrAnalysis: Failed to rewrite ForOp")); } return WalkResult::skip(); }) @@ -1598,11 +1978,13 @@ LogicalResult PtrAnalysis::rewriteOp(Operation *rootOp, bool useUnsafeMask) { PtrState state; OpBuilder b(getStateOp); if (succeeded(visitOperand(tritonValue, state, - getStateOp->getLoc(), b))) { + getStateOp->getLoc(), b)) && + state.isStructured()) { knownPtrs[tritonValue] = state; } else { - getStateOp->emitRemark("PtrAnalysis: Failed to populate ptr " - "state for tensor of indices"); + LLVM_DEBUG(getStateOp->emitRemark( + "PtrAnalysis: Failed to populate ptr " + "state for tensor of indices")); } } diff --git a/lib/Conversion/MemrefCopyToDMA_FlagTree/MemrefCopyToDMAFlagTree.cpp b/lib/Conversion/MemrefCopyToDMA_FlagTree/MemrefCopyToDMAFlagTree.cpp index 34fbb5c4..d73f8546 100644 --- a/lib/Conversion/MemrefCopyToDMA_FlagTree/MemrefCopyToDMAFlagTree.cpp +++ b/lib/Conversion/MemrefCopyToDMA_FlagTree/MemrefCopyToDMAFlagTree.cpp @@ -25,9 +25,6 @@ using namespace mlir; -#define GEN_PASS_CLASSES -#include "triton-shared/Conversion/TritonArithToLinalg/Passes.h.inc" - namespace { struct CopyConverter : public OpConversionPattern { using OpConversionPattern::OpConversionPattern; diff --git a/lib/Conversion/NoBufferize_FlagTree/NoBufferizeFlagTree.cpp b/lib/Conversion/NoBufferize_FlagTree/NoBufferizeFlagTree.cpp index 227ee2ff..92f75cc1 100644 --- a/lib/Conversion/NoBufferize_FlagTree/NoBufferizeFlagTree.cpp +++ b/lib/Conversion/NoBufferize_FlagTree/NoBufferizeFlagTree.cpp @@ -27,7 +27,6 @@ using namespace mlir; #define GEN_PASS_CLASSES #include "triton-shared/Conversion/NoBufferize_FlagTree/Passes.h.inc" -#include "triton-shared/Conversion/TritonArithToLinalg/Passes.h.inc" namespace { struct NoBufferizeConverter : public RewritePattern { diff --git a/lib/Conversion/TritonToUnstructured/TritonToUnstructuredPass.cpp b/lib/Conversion/TritonToUnstructured/TritonToUnstructuredPass.cpp index 4e2d13cb..7533e1bc 100644 --- a/lib/Conversion/TritonToUnstructured/TritonToUnstructuredPass.cpp +++ b/lib/Conversion/TritonToUnstructured/TritonToUnstructuredPass.cpp @@ -479,13 +479,16 @@ class TritonToUnstructuredPass ptrUsers.push_back(op); return success(); }) + .Case( + [&](Operation *op) { return success(); }) .Case([&](scf::ForOp forOp) { // Index of the init-arg corresponding to this use, note that // we have to subtract by 3 from the operand number because // scf.for ops always have 3 leading operands for start, end, // and step. - auto argIndex = use.getOperandNumber() - 3; - auto init = forOp.getInitArgs()[argIndex]; + auto argIndex = use.getOperandNumber(); + auto iterArgIndex = argIndex - 3; + auto init = forOp.getInitArgs()[iterArgIndex]; auto offsetInfo = offsetMap.at(init); @@ -505,10 +508,10 @@ class TritonToUnstructuredPass // At this point, the IR is in an invalid state because the // init-args still have tt.ptr. But at the end, we will // replace all uses of the tt.ptr to offset values. - auto iterArg = forOp.getRegionIterArg(argIndex); + auto iterArg = forOp.getRegionIterArg(iterArgIndex); iterArg.setType(offsetType); - auto res = forOp.getResult(argIndex); + auto res = forOp.getResult(iterArgIndex); res.setType(offsetType); // For other ops, we only need to push the result into the @@ -621,8 +624,35 @@ class TritonToUnstructuredPass return success(); }) - .Case( - [](auto) { return success(); }) + .Case([&](scf::YieldOp yieldOp) { + auto parentOp = yieldOp->getParentOp(); + // Initialization seeds only ptr-like values into the + // worklist. Here scf.if consumes the condition, while the ptr + // is passed via scf.yield, so we match yield first. + auto ifOp = dyn_cast(parentOp); + if (!ifOp) + return success(); + + auto operandIndex = use.getOperandNumber(); + // Use the scf.if result directly as the merged value of then/else + auto ifResult = ifOp.getResult(operandIndex); + if (!triton::isPtrTypeLike(ifResult.getType()) || + offsetMap.contains(ifResult)) + return success(); + + OpBuilder builder(ifOp->getContext()); + builder.setInsertionPointAfter(ifOp); + Value zero = builder.create( + ifOp.getLoc(), + builder.getIntegerAttr( + IntegerType::get(&getContext(), defaultBitWidth), 0)); + PtrOffset newOffsetInfo{ifResult, ifResult.getType(), + defaultBitWidth, zero}; + offsetMap.insert({ifResult, newOffsetInfo}); + workList.push(ifResult); + return success(); + }) + .Case([](auto) { return success(); }) .Case([](triton::CatOp op) { op->emitError("Do not support gather / scatter with multiple " "bases yet"); diff --git a/lib/Dialect/TritonStructured/IR/CMakeLists.txt b/lib/Dialect/TritonStructured/IR/CMakeLists.txt index 27aac38f..7b24232b 100644 --- a/lib/Dialect/TritonStructured/IR/CMakeLists.txt +++ b/lib/Dialect/TritonStructured/IR/CMakeLists.txt @@ -8,4 +8,5 @@ add_triton_library(TritonStructuredIR LINK_LIBS PUBLIC TritonIR MLIRIR + TritonSharedAnalysis ) diff --git a/lib/Dialect/TritonStructured/IR/TritonStructuredDialect.cpp b/lib/Dialect/TritonStructured/IR/TritonStructuredDialect.cpp index 2af19b8a..531bdb49 100644 --- a/lib/Dialect/TritonStructured/IR/TritonStructuredDialect.cpp +++ b/lib/Dialect/TritonStructured/IR/TritonStructuredDialect.cpp @@ -1,4 +1,6 @@ #include "triton-shared/Dialect/TritonStructured/IR/TritonStructuredDialect.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "triton-shared/Analysis/OpFoldResultUtils.h" using namespace mlir; using namespace mlir::tts; @@ -12,6 +14,113 @@ void TritonStructuredDialect::initialize() { >(); } +struct TTSLoadContiguousPattern : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(tts::LoadOp op, + PatternRewriter &rewriter) const override { + + auto makeTensorPtrOp = + op->getOperand(0).getDefiningOp(); + if (!makeTensorPtrOp) + return failure(); + + auto tensorType = cast(op.getType()); + auto rank = tensorType.getRank(); + + if (rank != 2) { + // TODO: Make it work for higher rank. Here simplify for calculate + // strides. + return failure(); + } + + auto mixedStrides = makeTensorPtrOp.getMixedStrides(); + assert(mixedStrides.size() == 2); + + // Check if first dimension (dim0) stride is 1 + auto dim0Stride = mixedStrides[0]; + auto dim0StrideVal = getConstValue(dim0Stride); + if (!dim0StrideVal || dim0StrideVal.value() != 1) { + return failure(); + } + + // Prevent infinite loop: check that dim1 stride is not also 1 + auto dim1Stride = mixedStrides[1]; + auto dim1StrideVal = getConstValue(dim1Stride); + if (dim1StrideVal && dim1StrideVal.value() == 1) { + return failure(); + } + + // NOTE: In FlagGems outer op, MakeTensorPtrOp may has stride = 0. So we + // need to check the order to decide whether to transpose + // FIXME: How to handle the stride = 0 case? + auto order = makeTensorPtrOp.getOrder(); + if (!order.empty() && + llvm::is_sorted(makeTensorPtrOp.getOrder(), std::greater<>())) { + return failure(); + } + + SmallVector newStrides; + // New dim0 stride = original dim0 size + // newStrides.push_back(rewriter.getIndexAttr(sizes[0])); + newStrides.push_back(mixedStrides[1]); + // New dim1 stride = 1 + newStrides.push_back(rewriter.getIndexAttr(1)); + + // Create new strides array with dim1 stride set to 1 and dim0 stride set to + // size[0] + auto sizes = makeTensorPtrOp.getSizes(); + // Create new sizes array (swapped) + SmallVector newSizes; + newSizes.push_back(sizes[1]); // Original dim1 size becomes new dim0 size + newSizes.push_back(sizes[0]); // Original dim0 size becomes new dim1 size + + auto offset = makeTensorPtrOp.getMixedOffsets(); + // Shape is original source tensor shape, no need to swap + SmallVector mixedShape = makeTensorPtrOp.getMixedShape(); + SmallVector newShape = + makeTensorPtrOp.isSplitPtr() + ? SmallVector{mixedShape[1], mixedShape[0]} + : mixedShape; + + // Create new offsets array (swapped) + auto mixedOffsets = makeTensorPtrOp.getMixedOffsets(); + assert(mixedOffsets.size() == 2); + // dim0 stride is 1 + SmallVector newOffsets = {mixedOffsets[1], mixedOffsets[0]}; + + // TODO: Now not used order + // Create the new MakeTensorPtrOp with swapped dimensions + auto newMakeTensorPtrOp = rewriter.create( + op->getLoc(), makeTensorPtrOp.getBase(), newSizes, newStrides, + newOffsets, newShape, order); + + auto dims = op.getMixedMaskDims(); + assert(!op.hasMask() || dims.size() == 2); + SmallVector newDims = + op.hasMask() ? SmallVector{dims[1], dims[0]} : dims; + + auto newLoadOp = rewriter.create( + op->getLoc(), newMakeTensorPtrOp, newDims, op.getOther()); + + Value init = rewriter.create( + op->getLoc(), tensorType.getShape(), tensorType.getElementType()); + + auto transposeOp = rewriter.create( + op->getLoc(), newLoadOp, init, + rewriter.getDenseI64ArrayAttr( + {1, 0})); // Permutation to swap dimensions + rewriter.replaceOp(op, transposeOp); + + return success(); + } +}; + +void TritonStructuredDialect::getCanonicalizationPatterns( + RewritePatternSet &results) const { + results.add(getContext()); +} + //===----------------------------------------------------------------------===// // TableGen'd op method definitions //===----------------------------------------------------------------------===// diff --git a/lib/Dialect/TritonStructured/IR/TritonStructuredOps.cpp b/lib/Dialect/TritonStructured/IR/TritonStructuredOps.cpp index 85038cae..a48e8a30 100644 --- a/lib/Dialect/TritonStructured/IR/TritonStructuredOps.cpp +++ b/lib/Dialect/TritonStructured/IR/TritonStructuredOps.cpp @@ -131,6 +131,72 @@ void MakeTensorPtrOp::build(OpBuilder &b, OperationState &state, Value base, b.getDenseI64ArrayAttr(staticShape), order); } +void MakeGatherScatterTensorPtrOp::build(OpBuilder &b, OperationState &state, + Value base, Value gatherScatterOffset, + int gatherScatterDim, + ArrayRef sizes, + ArrayRef strides, + ArrayRef offsets) { + build(b, state, base, gatherScatterOffset, Value(), gatherScatterDim, sizes, + strides, offsets); +} + +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); + + auto basePtr = cast(base.getType()); + auto elemType = basePtr.getPointeeType(); + + Type resType = RankedTensorType::get(sizes, basePtr); + + 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(); + + // Verify that the gatherScatterOffset is a 1D tensor. + auto rankedTensorType = dyn_cast(offsetType); + if (!rankedTensorType) { + return emitError("gatherScatterOffset must be a 1D tensor"); + } + + Type offsetEltType = rankedTensorType.getElementType(); + + if (!offsetEltType.isIntOrIndex()) { + return emitError("gatherScatterOffset must be a 1D tensor of " + "int or index type"); + } + + return success(); +} + void LoadOp::build(OpBuilder &b, OperationState &state, Value ptr, ArrayRef dims, Value other) { SmallVector staticDims; diff --git a/lib/Utils/Utils.cpp b/lib/Utils/Utils.cpp index cb4d86a7..73ddd889 100644 --- a/lib/Utils/Utils.cpp +++ b/lib/Utils/Utils.cpp @@ -128,8 +128,12 @@ bool isOperandMemorySpaceSPM(Value operand) { operand = forOp.getInitArgs()[idx]; } } else if (auto ifOp = dyn_cast(op)) { - bool thenResult = isOperandMemorySpaceSPM(ifOp.thenYield().getOperand(0)); - bool elseResult = isOperandMemorySpaceSPM(ifOp.elseYield().getOperand(0)); + auto ifResults = ifOp.getResults(); + auto idx = std::distance(ifResults.begin(), + std::find(ifResults.begin(), ifResults.end(), + operand)); + bool thenResult = isOperandMemorySpaceSPM(ifOp.thenYield().getOperand(idx)); + bool elseResult = isOperandMemorySpaceSPM(ifOp.elseYield().getOperand(idx)); assert(thenResult == elseResult && "Inconsistent memory space for IfOp results: " "one branch uses SPM, another branch does not");