diff --git a/third_party/ascend/include/TritonControlFlowOpt/ControlFlowAnalysis.h b/third_party/ascend/include/TritonControlFlowOpt/ControlFlowAnalysis.h index 1b09ca56a7..60ec5225da 100644 --- a/third_party/ascend/include/TritonControlFlowOpt/ControlFlowAnalysis.h +++ b/third_party/ascend/include/TritonControlFlowOpt/ControlFlowAnalysis.h @@ -75,7 +75,6 @@ struct AnalyzedComponent { struct AnalyzedValue { Type originalType; SmallVector components; - SmallVector invariants; SmallVector attributes; }; diff --git a/third_party/ascend/include/TritonControlFlowOpt/ControlFlowRewrite.h b/third_party/ascend/include/TritonControlFlowOpt/ControlFlowRewrite.h index e77c2b104c..49d0110a53 100644 --- a/third_party/ascend/include/TritonControlFlowOpt/ControlFlowRewrite.h +++ b/third_party/ascend/include/TritonControlFlowOpt/ControlFlowRewrite.h @@ -32,19 +32,27 @@ #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/StringRef.h" namespace mlir::triton::controlflow { +/// Handoff marker for SCF loops whose pointer slots have already been expanded +/// into policy-owned descriptor components. Its DenseI32ArrayAttr value lists +/// the loop-carried init/result slots occupied by pointer descriptor state. +/// TritonToLinalg uses those slots to preserve only the required producer +/// chains and removes the marker after conversion. +inline constexpr llvm::StringLiteral kPointerDescriptorBoundaryAttr = + "PointerDescriptorBoundary"; + /// Policy-owned description of one value crossing a control-flow boundary. /// -/// `components` are runtime values that a policy may place in an expanded SCF -/// signature. `invariants` and `attributes` are public storage whose layout is -/// interpreted only by the policy that creates them. The shared rewrite treats -/// those fields as opaque and only accesses `components` directly. +/// `components` contain every runtime value needed to rebuild the original +/// value. A policy may place a selected subset in an expanded SCF signature. +/// `attributes` retain non-SSA metadata. Both layouts are private to the +/// policy; the shared rewrite never interprets pointer-specific fields. struct DecomposedValue { Type originalType; SmallVector components; - SmallVector invariants; SmallVector attributes; }; @@ -71,8 +79,9 @@ class ControlFlowRewriteContext { /// /// The policy decides how its value is decomposed and rebuilt, which components /// cross loop/if boundaries, and whether two decompositions share a compatible -/// invariant schema. It is not an IR marker and carries no state between -/// policy invocations. +/// non-carried schema. It carries no mutable state between policy invocations; +/// a capability hook tells the shared rewrite whether expanded loop slots must +/// be recorded for downstream conversion. class ControlFlowRewritePolicy : public ControlFlowAnalysisPolicy { public: virtual ~ControlFlowRewritePolicy() = default; @@ -81,6 +90,11 @@ class ControlFlowRewritePolicy : public ControlFlowAnalysisPolicy { /// after cloning so later operations can reuse their exact component state. virtual bool shouldDecomposeOperation(Operation *op) const = 0; + /// Whether rewritten loops owned by this policy must expose their descriptor + /// slots to downstream conversion. The shared rewrite owns the positional + /// marker because it alone knows both the previous and expanded signatures. + virtual bool requiresPointerDescriptorBoundaryMarker() const { return false; } + virtual FailureOr decompose(Value value, const ControlFlowRewriteContext &context, OpBuilder &builder, Location loc) const = 0; diff --git a/third_party/ascend/include/TritonToLinalg/BlockPtrAnalysis.h b/third_party/ascend/include/TritonToLinalg/BlockPtrAnalysis.h index 3b3b7ca5e6..445480a667 100644 --- a/third_party/ascend/include/TritonToLinalg/BlockPtrAnalysis.h +++ b/third_party/ascend/include/TritonToLinalg/BlockPtrAnalysis.h @@ -45,6 +45,16 @@ namespace triton { enum class MemAccVal { Undefined = 0, StrucMemAcc = 1, UnstrucMemAcc = 2 }; +/// Creates a verifier-valid HIVM pointer cast for a scalar Triton pointer +/// represented by an integer address. Triton scalar pointers do not carry an +/// extent, while their downstream carrier is normally `memref` and every +/// dynamic memref dimension requires a size operand. Use one element as the +/// conservative carrier extent; reinterpret-cast lowering replaces it with a +/// precise access range when a larger descriptor is materialized. +hivm::PointerCastOp createScalarPointerCast(OpBuilder &builder, Location loc, + MemRefType resultType, + Value address); + struct MemAccType { MemAccVal value; @@ -288,8 +298,8 @@ class BlockDataParser { ConversionPatternRewriter &rewriter, llvm::SmallDenseMap &known); - static void - rewriteMakeTensorPtrOp(triton::MakeTensorPtrOp op, Value base, + static LogicalResult + rewriteMakeTensorPtrOp(triton::MakeTensorPtrOp op, Value convertedBase, ConversionPatternRewriter &rewriter, llvm::SmallDenseMap &known); @@ -312,9 +322,9 @@ class BlockDataParser { /// @param known is mainly designed for `rewriteLoop`, and is just non-const /// in `rewriteLoop`, `rewriteAddPtr` and `rewriteAdvance` - static void rewriteLoopOp(LoopLikeOpInterface op, - ConversionPatternRewriter &rewriter, - llvm::SmallDenseMap &known); + static LogicalResult + rewriteLoopOp(LoopLikeOpInterface op, ConversionPatternRewriter &rewriter, + llvm::SmallDenseMap &known); static void rewriteAddPtrToUnstrucMemAcc(triton::AddPtrOp op, triton::AddPtrOp::Adaptor &adaptor, diff --git a/third_party/ascend/include/TritonToLinalg/TritonOpConverter.h b/third_party/ascend/include/TritonToLinalg/TritonOpConverter.h index 15f2d32f63..d6aa7bafdf 100644 --- a/third_party/ascend/include/TritonToLinalg/TritonOpConverter.h +++ b/third_party/ascend/include/TritonToLinalg/TritonOpConverter.h @@ -37,6 +37,7 @@ #include "mlir/Transforms/DialectConversion.h" #include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/StringRef.h" #include "llvm/ADT/TypeSwitch.h" #include "llvm/Support/Debug.h" @@ -576,6 +577,66 @@ class GatherConverter : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override; }; +// These predicates select only scalar !tt.ptr transports. A +// tensor<...x!tt.ptr> follows the separate tensor-pointer lowering. +bool hasScalarPointerResult(scf::IfOp op); +bool isScalarPointerSelect(arith::SelectOp op); + +// Marks an scf.if temporarily rebuilt by IfConverter. Its scalar-pointer +// results are represented as complete i64 addresses, so only its own yields +// require the matching pointer-to-address conversion. +inline constexpr llvm::StringLiteral kScalarPointerCarrierBoundaryAttr = + "ScalarPointerCarrierBoundary"; + +// Rebuild an scf.if with scalar-pointer results so the boundary carries +// complete i64 addresses and reconstructs memrefs only after the join. +// The original branch regions are moved into the new operation, preserving +// side effects and allowing the conversion driver to rewrite each scf.yield +// operand in place. +// +// Example: +// %base = scf.if %cond -> !tt.ptr { +// scf.yield %lhs : !tt.ptr +// } else { +// scf.yield %rhs : !tt.ptr +// } +// %ptr = tt.make_tensor_ptr %base, ... +// becomes an scf.if returning i64 plus one hivm.pointer_cast after the if. +class IfConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(scf::IfOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +// Convert a scalar-pointer select into a select over complete integer addresses +// and reconstruct one memref after the selection. This handles both BlockPtr +// bases and ordinary scalar pointers without asking the backend to merge two +// memory objects. +// +// Example: +// %base = arith.select %cond, %lhs, %rhs : !tt.ptr +// %ptr = tt.make_tensor_ptr %base, ... +// becomes: +// %lhs_addr = memref.extract_aligned_pointer_as_index %lhs +// %rhs_addr = memref.extract_aligned_pointer_as_index %rhs +// %selected_addr = arith.select %cond, %lhs_addr, %rhs_addr : i64 +// %base = hivm.pointer_cast %selected_addr : i64 to memref +class PointerSelectConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(arith::SelectOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +// Convert the yields of an IfConverter-created scf.if to its carrier result +// types. In particular, a yielded scalar pointer becomes its complete i64 +// address. Yields belonging to ordinary ifs or loops are intentionally left to +// their owning conversions. class YieldConverter : public OpConversionPattern { public: using OpConversionPattern::OpConversionPattern; @@ -596,11 +657,15 @@ class LoopConverter : public OpConversionPattern { matchAndRewrite(LoopOpTy op, typename OpConversionPattern::OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { + // CFO-expanded descriptor loops already carry pointer-free policy values + // and remain structurally unchanged. This legacy BlockData rewrite is only + // valid for explicitly marked loops. + if (!op->hasAttr("UnhandledLoopOp")) + return failure(); llvm::SmallDenseMap known; - op->removeAttr("UnhandledLoopOp"); - BlockDataParser::rewriteLoopOp(op, rewriter, known); - return success(); + rewriter.modifyOpInPlace(op, [&]() { op->removeAttr("UnhandledLoopOp"); }); + return BlockDataParser::rewriteLoopOp(op, rewriter, known); } }; @@ -736,6 +801,14 @@ class PtrToIntConverter : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override; }; +class IntToPtrConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(triton::IntToPtrOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + class IndexPutConverter : public OpConversionPattern { public: diff --git a/third_party/ascend/include/TritonToUnstructure/BubbleUpOperation.h b/third_party/ascend/include/TritonToUnstructure/BubbleUpOperation.h index 401b20915f..fe745a590a 100644 --- a/third_party/ascend/include/TritonToUnstructure/BubbleUpOperation.h +++ b/third_party/ascend/include/TritonToUnstructure/BubbleUpOperation.h @@ -72,6 +72,11 @@ class BubbleUpExtract : public OpRewritePattern { PatternRewriter &rewriter) const; void bubbleUpOperation(ExtractOpTy op, arith::CmpIOp parentOp, Location loc, PatternRewriter &rewriter) const; + // Pushes extract(select(condition, lhs, rhs)) through the select. A shaped + // condition is extracted at the same position while a scalar condition is + // reused directly. + void bubbleUpOperation(ExtractOpTy op, arith::SelectOp parentOp, Location loc, + PatternRewriter &rewriter) const; void bubbleUpOperation(ExtractOpTy op, arith::TruncFOp parentOp, Location loc, PatternRewriter &rewriter) const; void bubbleUpOperation(ExtractOpTy op, arith::ExtFOp parentOp, Location loc, diff --git a/third_party/ascend/lib/TritonControlFlowOpt/BlockPtrDecompose.cpp b/third_party/ascend/lib/TritonControlFlowOpt/BlockPtrDecompose.cpp index ae9f6a247f..86e1e8590b 100644 --- a/third_party/ascend/lib/TritonControlFlowOpt/BlockPtrDecompose.cpp +++ b/third_party/ascend/lib/TritonControlFlowOpt/BlockPtrDecompose.cpp @@ -35,15 +35,26 @@ using namespace mlir::triton::controlflow; namespace { +static constexpr unsigned kBaseComponent = 0; +static constexpr unsigned getShapeStart() { return kBaseComponent + 1; } +static constexpr unsigned getStrideStart(unsigned rank) { + return getShapeStart() + rank; +} +static constexpr unsigned getOffsetStart(unsigned rank) { + return getStrideStart(rank) + rank; +} + /// Block-pointer component layout used only by this policy: /// -/// components = [shape..., strides..., offsets...] -/// invariants = [base] +/// components = [base_address, shape..., strides..., offsets...] /// attributes = [order] /// -/// Loops currently carry only `offsets`; shape and strides must remain -/// invariant across a backedge. An scf.if may select any component whose SSA -/// value differs between its branches. +/// Every supported SCF boundary carries all components in this exact order. +/// The base is represented as an i64 address rather than a pointer/memref so a +/// control-flow merge remains an SSA value selection and cannot be lowered as +/// a memory-object copy by the backend. +/// `order` remains policy-owned static metadata and must agree on every +/// incoming path. /// /// Keep rank/layout checks local to this file: the generic control-flow /// machinery intentionally does not know the descriptor format of a policy. @@ -59,15 +70,41 @@ static FailureOr getRank(Type originalType) { return tensorType.getRank(); } +// Recovers the scalar pointer type accepted by tt.make_tensor_ptr from the +// BlockPtr result type. The address space is preserved across the temporary +// integer carrier. +static FailureOr getBasePointerType(Type originalType) { + auto pointerType = dyn_cast(originalType); + if (!pointerType) + return failure(); + auto tensorType = dyn_cast(pointerType.getPointeeType()); + if (!tensorType) + return failure(); + return triton::PointerType::get(tensorType.getElementType(), + pointerType.getAddressSpace()); +} + // Validates the common block-pointer schema for either state representation. // Their component element types differ, but this check only needs field sizes. template static bool hasValidLayout(const StateT &state) { FailureOr rank = getRank(state.originalType); - return succeeded(rank) && state.components.size() == 3 * *rank && - state.invariants.size() == 1 && state.attributes.size() == 1 && + return succeeded(rank) && state.components.size() == 1 + 3 * *rank && + state.attributes.size() == 1 && isa(state.attributes.front()); } +/// Returns the complete, ordered descriptor range used to expand one +/// block-pointer control-flow slot. Keeping this policy-local prevents the +/// generic SCF rewrite from depending on the BlockPtr field layout. +static SmallVector +getAllComponentIndices(const AnalyzedValue &value) { + SmallVector indices; + indices.reserve(value.components.size()); + for (unsigned index = 0; index < value.components.size(); ++index) + indices.push_back(index); + return indices; +} + class BlockPtrPolicy final : public ControlFlowRewritePolicy { public: bool matches(Type type) const override { @@ -93,7 +130,11 @@ class BlockPtrPolicy final : public ControlFlowRewritePolicy { // types and symbolic identities here; this phase must not create IR. AnalyzedValue result; result.originalType = value.getType(); - unsigned componentIndex = 0; + Value base = makePtr.getBase(); + result.components.push_back( + {IntegerType::get(value.getContext(), 64), + ComponentIdentity::fromValue(base, kBaseComponent)}); + unsigned componentIndex = getShapeStart(); auto appendComponents = [&](ValueRange values) { for (Value component : values) { result.components.push_back( @@ -104,7 +145,6 @@ class BlockPtrPolicy final : public ControlFlowRewritePolicy { appendComponents(makePtr.getShape()); appendComponents(makePtr.getStrides()); appendComponents(makePtr.getOffsets()); - result.invariants.push_back(makePtr.getBase()); result.attributes.push_back(makePtr.getOrderAttr()); if (!hasValidLayout(result)) return failure(); @@ -124,7 +164,7 @@ class BlockPtrPolicy final : public ControlFlowRewritePolicy { if (advance.getOffsets().size() != rank) return failure(); for (unsigned dimension = 0; dimension < rank; ++dimension) { - unsigned componentIndex = 2 * rank + dimension; + unsigned componentIndex = getOffsetStart(rank) + dimension; result->components[componentIndex].identity = ComponentIdentity::fromValue(value, componentIndex); } @@ -136,13 +176,9 @@ class BlockPtrPolicy final : public ControlFlowRewritePolicy { getLoopCandidateComponents(const AnalyzedValue &value) const override { if (!hasValidLayout(value)) return failure(); - unsigned rank = *getRank(value.originalType); - SmallVector indices; - // Only the final rank entries (offsets) are legal loop-carried state in the - // current block-pointer model. - for (unsigned dimension = 0; dimension < rank; ++dimension) - indices.push_back(2 * rank + dimension); - return indices; + // A BlockPtr never crosses a supported loop boundary directly. Its entire + // descriptor becomes ordinary loop-carried SSA state. + return getAllComponentIndices(value); } FailureOr> @@ -153,35 +189,25 @@ class BlockPtrPolicy final : public ControlFlowRewritePolicy { !hasValidLayout(next) || initial.originalType != regionArgument.originalType || initial.originalType != next.originalType || - initial.invariants != regionArgument.invariants || - initial.invariants != next.invariants || initial.attributes != regionArgument.attributes || initial.attributes != next.attributes) return failure(); - unsigned rank = *getRank(initial.originalType); - // The current implementation does not expand shape or stride iter_args. - // Reject a loop that changes either instead of silently reconstructing a - // descriptor with stale values. - for (unsigned index = 0; index < 2 * rank; ++index) { - if (initial.components[index].type != next.components[index].type || - initial.components[index].identity != next.components[index].identity) - return failure(); - } - - SmallVector transferred; - // An offset is carried only if the backedge state depends on the region - // argument. Constant/invariant offsets remain outside the loop signature. - for (unsigned dimension = 0; dimension < rank; ++dimension) { - unsigned index = 2 * rank + dimension; + // Full descriptor transfer permits every field to change, but each + // position must retain the strict type required by tt.make_tensor_ptr. + // In particular, tt.advance preserving its base does not make the base a + // loop invariant: a nested if may rebuild the descriptor from a different + // root base, and that if result becomes the next backedge value. A generic + // SCF canonicalization may remove exactly forwarded components from an + // advance-only loop later; this policy must retain changing-base semantics. + for (unsigned index = 0; index < initial.components.size(); ++index) { if (failed(joinComponentTypes(initial.components[index].type, + regionArgument.components[index].type)) || + failed(joinComponentTypes(initial.components[index].type, next.components[index].type))) return failure(); - if (regionArgument.components[index].identity != - next.components[index].identity) - transferred.push_back(index); } - return transferred; + return getAllComponentIndices(initial); } FailureOr> @@ -189,24 +215,18 @@ class BlockPtrPolicy final : public ControlFlowRewritePolicy { const AnalyzedValue &elseValue) const override { if (!hasValidLayout(thenValue) || !hasValidLayout(elseValue) || thenValue.originalType != elseValue.originalType || - thenValue.invariants != elseValue.invariants || - thenValue.attributes != elseValue.attributes || - thenValue.components.size() != elseValue.components.size()) + thenValue.attributes != elseValue.attributes) return failure(); - // Unlike loops, an if can select shape, stride, or offset components. Base - // and order remain invariants because AdapterIR cannot represent a runtime - // selection between heterogeneous pointer descriptors. - SmallVector transferred; + // Both branches yield a complete descriptor even when a field has the same + // symbolic identity. This gives every rewritten BlockPtr boundary one + // stable positional schema. for (unsigned index = 0; index < thenValue.components.size(); ++index) { if (failed(joinComponentTypes(thenValue.components[index].type, elseValue.components[index].type))) return failure(); - if (thenValue.components[index].identity != - elseValue.components[index].identity) - transferred.push_back(index); } - return transferred; + return getAllComponentIndices(thenValue); } FailureOr joinComponentTypes(Type lhs, Type rhs) const override { @@ -224,6 +244,8 @@ class BlockPtrPolicy final : public ControlFlowRewritePolicy { return isa(op); } + bool requiresPointerDescriptorBoundaryMarker() const override { return true; } + FailureOr decompose(Value value, const ControlFlowRewriteContext &context, OpBuilder &builder, @@ -241,13 +263,17 @@ class BlockPtrPolicy final : public ControlFlowRewritePolicy { // Materialize the concrete counterpart of analyzeValue's descriptor. DecomposedValue result; result.originalType = value.getType(); + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPoint(makePtr); + Value baseAddress = builder.create( + makePtr.getLoc(), builder.getI64Type(), makePtr.getBase()); + result.components.push_back(baseAddress); result.components.append(makePtr.getShape().begin(), makePtr.getShape().end()); result.components.append(makePtr.getStrides().begin(), makePtr.getStrides().end()); result.components.append(makePtr.getOffsets().begin(), makePtr.getOffsets().end()); - result.invariants.push_back(makePtr.getBase()); result.attributes.push_back(makePtr.getOrderAttr()); if (!hasValidLayout(result)) return failure(); @@ -271,7 +297,7 @@ class BlockPtrPolicy final : public ControlFlowRewritePolicy { OpBuilder::InsertionGuard guard(builder); builder.setInsertionPoint(advance); for (auto [dim, delta] : llvm::enumerate(advance.getOffsets())) { - unsigned component = 2 * *rank + dim; + unsigned component = getOffsetStart(*rank) + dim; Value currentOffset = result->components[component]; Value remappedDelta = context.remap(delta); if (!remappedDelta) @@ -295,11 +321,18 @@ class BlockPtrPolicy final : public ControlFlowRewritePolicy { return nullptr; unsigned rank = *getRank(value.originalType); auto order = cast(value.attributes.front()); + FailureOr basePointerType = + getBasePointerType(value.originalType); + if (failed(basePointerType) || + value.components[kBaseComponent].getType() != builder.getI64Type()) + return nullptr; + Value base = builder.create( + loc, *basePointerType, value.components[kBaseComponent]); return builder.create( - loc, value.originalType, value.invariants.front(), - ValueRange(value.components).take_front(rank), - ValueRange(value.components).slice(rank, rank), - ValueRange(value.components).take_back(rank), order); + loc, value.originalType, base, + ValueRange(value.components).slice(getShapeStart(), rank), + ValueRange(value.components).slice(getStrideStart(rank), rank), + ValueRange(value.components).slice(getOffsetStart(rank), rank), order); } }; @@ -308,8 +341,8 @@ class BlockPtrPolicy final : public ControlFlowRewritePolicy { namespace mlir::triton::controlflow { LogicalResult runBlockPtrDecompose(ModuleOp module) { - // Make the explicit descriptor carried by a block pointer cross each - // supported SCF boundary as ordinary SSA components. + // Expand every BlockPtr that crosses a supported SCF boundary into its full + // ordered descriptor, then rebuild the pointer at each use site. BlockPtrPolicy policy; return rewriteControlFlow(module, policy); } diff --git a/third_party/ascend/lib/TritonControlFlowOpt/ControlFlowAnalysis.cpp b/third_party/ascend/lib/TritonControlFlowOpt/ControlFlowAnalysis.cpp index 67ad7c7c31..68fd957696 100644 --- a/third_party/ascend/lib/TritonControlFlowOpt/ControlFlowAnalysis.cpp +++ b/third_party/ascend/lib/TritonControlFlowOpt/ControlFlowAnalysis.cpp @@ -44,7 +44,7 @@ static bool isSupportedControlFlow(Operation *op) { static void setResultIdentity(AnalyzedValue &value, Value result, ArrayRef componentIndices) { // A transferred component is represented by the SCF result outside the op; - // invariant components retain their incoming symbolic identities. + // non-transferred components retain their incoming symbolic identities. for (unsigned index : componentIndices) value.components[index].identity = ComponentIdentity::fromValue(result, index); @@ -120,8 +120,8 @@ void ControlFlowAnalysisContext::bindRegionArgument( Value argument, const AnalyzedValue &initial, ArrayRef componentIndices) { // Only candidate loop components acquire a new identity at region entry. - // Policy-owned invariants and non-carried components remain traceable to the - // initial descriptor and can therefore be checked at the backedge. + // Policy-owned non-carried components remain traceable to the initial + // descriptor and can therefore be checked at the backedge. AnalyzedValue argumentState = initial; for (unsigned index : componentIndices) argumentState.components[index].identity = diff --git a/third_party/ascend/lib/TritonControlFlowOpt/ControlFlowRewrite.cpp b/third_party/ascend/lib/TritonControlFlowOpt/ControlFlowRewrite.cpp index 0b92e8d72c..3bf3f22e53 100644 --- a/third_party/ascend/lib/TritonControlFlowOpt/ControlFlowRewrite.cpp +++ b/third_party/ascend/lib/TritonControlFlowOpt/ControlFlowRewrite.cpp @@ -34,6 +34,8 @@ #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallVector.h" +#include +#include #include using namespace mlir; @@ -43,6 +45,7 @@ using mlir::triton::controlflow::ControlFlowRewritePlan; using mlir::triton::controlflow::ControlFlowRewritePolicy; using mlir::triton::controlflow::ControlFlowSlotAnalysis; using mlir::triton::controlflow::DecomposedValue; +using mlir::triton::controlflow::kPointerDescriptorBoundaryAttr; namespace mlir::triton::controlflow { @@ -65,7 +68,7 @@ namespace { // handlers are mutually recursive through rewriteBodyOps(), share one // short-lived RewriteEnv, and must agree on signature expansion, nested-op // ordering and failure cleanup. Splitting them by op kind would expose those -// private invariants through additional internal headers without creating an +// private constraints through additional internal headers without creating an // independently reusable component. // //===----------------------------------------------------------------------===// @@ -84,8 +87,8 @@ namespace { /// offsets, a region-local environment can contain: /// valueMapping: %old_ptr_arg -> %rebuilt_ptr /// decomposedValues: %old_ptr_arg -> { -/// components = [%shape0, %stride0, %new_offset0], -/// invariants = [%base], attributes = [order] +/// components = [%base, %shape0, %stride0, %new_offset0], +/// attributes = [order] /// } /// Operations cloned into that region use the mapping, while pointer /// decomposition uses the stored components. @@ -152,9 +155,9 @@ struct RewriteEnv { /// %next = tt.advance %ptr, [%delta0, %delta1] /// /// decomposeValue(%next) -> DecomposedValue { - /// components = [%shape0, %shape1, %stride0, %stride1, + /// components = [%base, %shape0, %shape1, %stride0, %stride1, /// %offset0 + %delta0, %offset1 + %delta1], - /// invariants = [%base], attributes = [order] + /// attributes = [order] /// } /// The caller can put selected components into a new control-flow signature /// or pass the whole descriptor to policy.recompose(). Unsupported values @@ -219,6 +222,60 @@ struct IfPointerInfo { std::optional thenInfo; }; +// Updates the downstream descriptor marker after a loop signature expansion. +// Existing slots may have been recorded by an earlier pointer policy, so they +// are remapped through oldToNewStart before the current policy's expanded +// component slots are merged. The resulting DenseI32ArrayAttr is expressed in +// the replacement loop's iter-argument/result coordinate space. +// +// Example: +// old slots = [0], oldToNewStart = [0, 1], new component slots = [1, 2] +// result = [0, 1, 2] +static LogicalResult updatePointerDescriptorBoundaryMarker( + Operation *loop, ArrayRef pointerInfos, + ArrayRef oldToNewStart, const ControlFlowRewritePolicy &policy) { + SmallVector descriptorSlots; + llvm::SmallDenseSet seenSlots; + auto appendSlot = [&](unsigned slot) -> LogicalResult { + if (slot > static_cast(std::numeric_limits::max())) + return failure(); + if (!seenSlots.insert(slot).second) + return success(); + descriptorSlots.push_back(static_cast(slot)); + return success(); + }; + + if (Attribute oldMarker = loop->getAttr(kPointerDescriptorBoundaryAttr)) { + auto oldSlots = dyn_cast(oldMarker); + if (!oldSlots) + return failure(); + for (int32_t oldSlot : oldSlots.asArrayRef()) { + if (oldSlot < 0 || + static_cast(oldSlot) >= oldToNewStart.size() || + failed(appendSlot(oldToNewStart[oldSlot]))) + return failure(); + } + } + + if (policy.requiresPointerDescriptorBoundaryMarker()) { + for (const LoopPointerInfo &pointerInfo : pointerInfos) { + for (unsigned newSlot : pointerInfo.newIndices) { + if (failed(appendSlot(newSlot))) + return failure(); + } + } + } + + if (descriptorSlots.empty()) { + loop->removeAttr(kPointerDescriptorBoundaryAttr); + return success(); + } + llvm::sort(descriptorSlots); + loop->setAttr(kPointerDescriptorBoundaryAttr, + DenseI32ArrayAttr::get(loop->getContext(), descriptorSlots)); + return success(); +} + // Copies the values selected by indices into a new owning vector while // preserving their input order. Indices must be unique and in bounds; this // function reports failure instead of deduplicating or accessing invalid input. @@ -245,9 +302,9 @@ static FailureOr> gatherValues(ValueRange sourceValues, // replacement changes the component type; the input object remains unchanged. // // Example: -// decomposition.components = [shape, stride, originalOffset] -// componentIndices = [2], replacements = [nextOffset] -// result.components = [shape, stride, nextOffset] +// decomposition.components = [base, shape, stride, originalOffset] +// componentIndices = [3], replacements = [nextOffset] +// result.components = [base, shape, stride, nextOffset] static FailureOr withReplacedComponents(DecomposedValue decomposition, ArrayRef componentIndices, @@ -306,10 +363,12 @@ static auto findPointerInfoByOldIndex(InfoRange &pointerInfos, return nullptr; } -// A replacement loop carries selected scalar or tensor descriptor components -// instead of the original pointer iter-argument. Operations cloned from the -// original body still expect one pointer-typed block argument, so this function -// reconstructs that pointer at the replacement region entry. +// A replacement loop carries policy-selected scalar or tensor descriptor +// components instead of the original pointer iter-argument. BlockPtr selects +// its complete descriptor; TensorPtr currently selects complete offsets only. +// Operations cloned from the original body still expect one pointer-typed +// block argument, so this function reconstructs that pointer at the replacement +// region entry. // `pointerInfo.newIndices` selects the current component values from // `newRegionArguments`, while // `pointerInfo.componentIndices` identifies the descriptor fields that those @@ -323,17 +382,17 @@ static auto findPointerInfoByOldIndex(InfoRange &pointerInfos, // // Example: // oldRegionArgument = %old_ptr -// newRegionArguments = [%ordinary, %current_offset0, %current_offset1] -// pointerInfo.newIndices = [1, 2] -// pointerInfo.componentIndices = [4, 5] +// newRegionArguments = [%ordinary, %base, %shape, %stride, %offset] +// pointerInfo.newIndices = [1, 2, 3, 4] +// pointerInfo.componentIndices = [0, 1, 2, 3] // pointerInfo.initInfo.components = -// [shape0, shape1, stride0, stride1, initial_offset0, initial_offset1] +// [initial_base, initial_shape, initial_stride, initial_offset] // -// The rebuilt descriptor keeps shape and stride, replaces the final two -// components with the current offsets, and records `%old_ptr -> %rebuilt_ptr` -// in `regionEnv`. Invalid indices, incompatible component types, or a policy -// that cannot recompose the descriptor return failure without recording a -// partial binding; the enclosing loop rewrite owns cleanup of inserted IR. +// The rebuilt BlockPtr descriptor replaces all four fields with the current +// loop values and records `%old_ptr -> %rebuilt_ptr` in `regionEnv`. Invalid +// indices, incompatible component types, or a policy that cannot recompose the +// descriptor return failure without recording a partial binding; the enclosing +// loop rewrite owns cleanup of inserted IR. static LogicalResult bindLoopCarriedPointer(Value oldRegionArgument, const LoopPointerInfo &pointerInfo, ValueRange newRegionArguments, @@ -416,16 +475,19 @@ static LogicalResult bindLoopRegionArguments( // Example: // oldOperands = [%next_ptr, %sum] // pointerInfo = { -// oldIndex = 0, componentIndices = [4, 5], newIndices = [0, 1] +// oldIndex = 0, componentIndices = [0, 1, 2, 3], +// newIndices = [0, 1, 2, 3] // } -// currentRegionArguments = [%current_offset0, %current_offset1, %sum_arg] +// currentRegionArguments = +// [%current_base, %current_shape, %current_stride, %current_offset, +// %sum_arg] // // A valid `%next_ptr` decomposition produces -// `[%next_offset0, %next_offset1, %mapped_sum]`. If pointer decomposition or -// component normalization fails, the output instead uses -// `[%current_offset0, %current_offset1, %mapped_sum]`. The fallback keeps the -// temporary scf.yield/scf.condition structurally complete until the enclosing -// failed loop rewrite erases it. +// `[%next_base, %next_shape, %next_stride, %next_offset, %mapped_sum]`. If +// pointer decomposition or component normalization fails, the output instead +// uses the four current descriptor arguments followed by `%mapped_sum`. The +// fallback keeps the temporary scf.yield/scf.condition structurally complete +// until the enclosing failed loop rewrite erases it. // // The output vector is separate from the LogicalResult intentionally. The // function visits every old operand and fills all available fallback positions @@ -484,15 +546,17 @@ static LogicalResult rewriteLoopTerminatorOperands( // // Example: // oldResults = [%old_sum, %old_ptr, %old_flag] -// newResults = [%new_sum, %offset0, %offset1, %new_flag] +// newResults = +// [%new_sum, %base, %shape, %stride, %offset, %new_flag] // pointerInfo = { -// oldIndex = 1, componentIndices = [4, 5], newIndices = [1, 2] +// oldIndex = 1, componentIndices = [0, 1, 2, 3], +// newIndices = [1, 2, 3, 4] // } -// oldToNewStart = [0, 1, 3] +// oldToNewStart = [0, 1, 5] // // The function maps `%old_sum -> %new_sum` and `%old_flag -> %new_flag`. It -// inserts `%offset0` and `%offset1` into the pointer descriptor, rebuilds -// `%old_ptr`, and records `%old_ptr -> %rebuilt_ptr` plus that decomposition. +// inserts all four results into the pointer descriptor, rebuilds `%old_ptr`, +// and records `%old_ptr -> %rebuilt_ptr` plus that decomposition. // // The caller must set the builder insertion point after the replacement loop, // so any rebuilt pointer dominates later operations. On failure this function @@ -675,6 +739,12 @@ static LogicalResult rewriteForOp(scf::ForOp forOp, OpBuilder &builder, bodyBuilder.create(yieldOp.getLoc(), newYieldOperands); }); newForOp->setAttrs(forOp->getAttrs()); + if (analysis->rewritesOwnSignature() && + failed(updatePointerDescriptorBoundaryMarker( + newForOp, pointerInfos, oldToNewStart, env.policy))) { + newForOp.erase(); + return failure(); + } if (!bodyOk) { newForOp.erase(); @@ -797,6 +867,12 @@ static LogicalResult rewriteWhileOp(scf::WhileOp whileOp, OpBuilder &builder, bodyBuilder.create(yieldOp.getLoc(), newYieldOperands); }); newWhileOp->setAttrs(whileOp->getAttrs()); + if (analysis->rewritesOwnSignature() && + failed(updatePointerDescriptorBoundaryMarker( + newWhileOp, pointerInfos, oldToNewStart, env.policy))) { + newWhileOp.erase(); + return failure(); + } if (!bodyOk) { newWhileOp.erase(); diff --git a/third_party/ascend/lib/TritonControlFlowOpt/TensorPtrDecompose.cpp b/third_party/ascend/lib/TritonControlFlowOpt/TensorPtrDecompose.cpp index c55a80186e..76be21761a 100644 --- a/third_party/ascend/lib/TritonControlFlowOpt/TensorPtrDecompose.cpp +++ b/third_party/ascend/lib/TritonControlFlowOpt/TensorPtrDecompose.cpp @@ -34,8 +34,8 @@ using namespace mlir::triton::controlflow; namespace { -static constexpr unsigned kOffsetsComponent = 0; -static constexpr unsigned kBaseInvariant = 0; +static constexpr unsigned kBaseComponent = 0; +static constexpr unsigned kCompleteOffsetsComponent = 1; static constexpr unsigned kBaseIsScalarAttribute = 0; /// Identifies tensor-of-pointers handled by this stage: @@ -48,14 +48,13 @@ static bool isTensorPointerType(Type type) { /// Tensor-pointer state used only by this policy: /// -/// control-flow components = [complete_offsets] -/// rewrite-only invariants = [common_base] +/// components = [common_base, complete_offsets] /// attributes = [base_is_scalar] /// -/// Only `components` expand an SCF signature. The common base is deliberately -/// kept out of iter-args and results; it must be identical at every incoming -/// edge and is used to rebuild the tensor-of-pointers inside and after the -/// rewritten control-flow operation. +/// The complete offsets are the only transfer candidate in this NFC stage. +/// The common base is now an ordinary component, but it must remain identical +/// at every incoming edge and therefore stays outside expanded SCF signatures. +/// static RankedTensorType getDefaultOffsetsType(Type pointerType) { auto pointerTensor = cast(pointerType); return RankedTensorType::get(pointerTensor.getShape(), @@ -64,14 +63,12 @@ static RankedTensorType getDefaultOffsetsType(Type pointerType) { } /// Validates the policy-owned component layout, including the offsets shape -/// and the representation selected for the invariant base. -static bool hasValidSchema(Type originalType, Type offsetsType, - ArrayRef invariants, +/// and the scalar-or-tensor representation selected for the base component. +static bool hasValidSchema(Type originalType, Type baseType, Type offsetsType, ArrayRef attributes) { auto pointerTensor = dyn_cast(originalType); auto offsetsTensor = dyn_cast(offsetsType); - if (!pointerTensor || !offsetsTensor || invariants.size() != 1 || - attributes.size() != 1 || + if (!pointerTensor || !offsetsTensor || attributes.size() != 1 || !isa(pointerTensor.getElementType()) || !isa(offsetsTensor.getElementType()) || pointerTensor.getShape() != offsetsTensor.getShape() || @@ -83,33 +80,39 @@ static bool hasValidSchema(Type originalType, Type offsetsType, return false; Type expectedBaseType = baseIsScalar.getValue() ? pointerTensor.getElementType() : originalType; - return invariants[kBaseInvariant].getType() == expectedBaseType; + return baseType == expectedBaseType; } /// Validates the concrete Value-based state created while rewriting IR. static bool hasValidLayout(const DecomposedValue &value) { - return value.components.size() == 1 && + return value.components.size() == 2 && hasValidSchema(value.originalType, - value.components[kOffsetsComponent].getType(), - value.invariants, value.attributes); + value.components[kBaseComponent].getType(), + value.components[kCompleteOffsetsComponent].getType(), + value.attributes); } /// Validates the read-only counterpart without materializing component Values. static bool hasValidLayout(const AnalyzedValue &value) { - return value.components.size() == 1 && + return value.components.size() == 2 && hasValidSchema(value.originalType, - value.components[kOffsetsComponent].type, - value.invariants, value.attributes); + value.components[kBaseComponent].type, + value.components[kCompleteOffsetsComponent].type, + value.attributes); } -template -static bool haveSameTensorPtrBaseSchema(const StateT &lhs, const StateT &rhs) { +static bool haveSameTensorPtrBaseSchema(const AnalyzedValue &lhs, + const AnalyzedValue &rhs) { return hasValidLayout(lhs) && hasValidLayout(rhs) && lhs.originalType == rhs.originalType && - lhs.invariants == rhs.invariants && lhs.attributes == rhs.attributes; + lhs.components[kBaseComponent].type == + rhs.components[kBaseComponent].type && + lhs.components[kBaseComponent].identity == + rhs.components[kBaseComponent].identity && + lhs.attributes == rhs.attributes; } -/// Whether the invariant base must be broadcast before rebuilding addptr. +/// Whether the scalar base component must be broadcast before rebuilding. static bool hasScalarBase(const DecomposedValue &value) { return cast(value.attributes[kBaseIsScalarAttribute]).getValue(); } @@ -206,14 +209,15 @@ class TensorPtrDecomposePolicy final : public ControlFlowRewritePolicy { FailureOr result = context.analyzeValue(addPtr.getPtr()); if (failed(result) || !hasValidLayout(*result)) return failure(); - FailureOr offsetsType = - getWiderOffsetsType(result->components[kOffsetsComponent].type, - addPtr.getOffset().getType()); + FailureOr offsetsType = getWiderOffsetsType( + result->components[kCompleteOffsetsComponent].type, + addPtr.getOffset().getType()); if (failed(offsetsType)) return failure(); result->originalType = value.getType(); - result->components[kOffsetsComponent] = { - *offsetsType, ComponentIdentity::fromValue(value, kOffsetsComponent)}; + result->components[kCompleteOffsetsComponent] = { + *offsetsType, + ComponentIdentity::fromValue(value, kCompleteOffsetsComponent)}; return *result; } @@ -223,10 +227,12 @@ class TensorPtrDecomposePolicy final : public ControlFlowRewritePolicy { if (!isa(splat.getSrc().getType())) return failure(); Type offsetsType = getDefaultOffsetsType(value.getType()); - return AnalyzedValue{value.getType(), - {{offsetsType, ComponentIdentity::zero()}}, - {splat.getSrc()}, - {BoolAttr::get(value.getContext(), true)}}; + return AnalyzedValue{ + value.getType(), + {{splat.getSrc().getType(), + ComponentIdentity::fromValue(splat.getSrc(), kBaseComponent)}, + {offsetsType, ComponentIdentity::zero(kCompleteOffsetsComponent)}}, + {BoolAttr::get(value.getContext(), true)}}; } // An otherwise opaque tensor-of-pointers is treated as an already-vector @@ -235,38 +241,40 @@ class TensorPtrDecomposePolicy final : public ControlFlowRewritePolicy { if (!matches(value.getType())) return failure(); Type offsetsType = getDefaultOffsetsType(value.getType()); - return AnalyzedValue{value.getType(), - {{offsetsType, ComponentIdentity::zero()}}, - {value}, - {BoolAttr::get(value.getContext(), false)}}; + return AnalyzedValue{ + value.getType(), + {{value.getType(), ComponentIdentity::fromValue(value, kBaseComponent)}, + {offsetsType, ComponentIdentity::zero(kCompleteOffsetsComponent)}}, + {BoolAttr::get(value.getContext(), false)}}; } /// Tensor pointers have exactly one loop-transfer candidate: the complete - /// per-lane offsets tensor at component index 0. + /// per-lane offsets tensor at component index 1. FailureOr> getLoopCandidateComponents(const AnalyzedValue &value) const override { if (!hasValidLayout(value)) return failure(); - return SmallVector{kOffsetsComponent}; + return SmallVector{kCompleteOffsetsComponent}; } /// Classifies the loop offsets as transferred only when the backedge changes /// their symbolic identity. The original pointer type, common base, and base - /// representation must remain invariant for reconstruction to be valid. + /// representation must remain unchanged for reconstruction to be valid. FailureOr> getLoopTransferredComponents(const AnalyzedValue &initial, const AnalyzedValue ®ionArgument, const AnalyzedValue &next) const override { if (!haveSameTensorPtrBaseSchema(initial, regionArgument) || !haveSameTensorPtrBaseSchema(initial, next) || - failed(joinComponentTypes(initial.components[kOffsetsComponent].type, - next.components[kOffsetsComponent].type))) + failed(joinComponentTypes( + initial.components[kCompleteOffsetsComponent].type, + next.components[kCompleteOffsetsComponent].type))) return failure(); // Yielding the region argument unchanged requires no new SCF iter-arg. - if (regionArgument.components[kOffsetsComponent].identity == - next.components[kOffsetsComponent].identity) + if (regionArgument.components[kCompleteOffsetsComponent].identity == + next.components[kCompleteOffsetsComponent].identity) return SmallVector{}; - return SmallVector{kOffsetsComponent}; + return SmallVector{kCompleteOffsetsComponent}; } /// Merges the two `scf.if` pointer states. Different complete-offset @@ -276,15 +284,15 @@ class TensorPtrDecomposePolicy final : public ControlFlowRewritePolicy { getIfTransferredComponents(const AnalyzedValue &thenValue, const AnalyzedValue &elseValue) const override { if (!haveSameTensorPtrBaseSchema(thenValue, elseValue) || - failed( - joinComponentTypes(thenValue.components[kOffsetsComponent].type, - elseValue.components[kOffsetsComponent].type))) + failed(joinComponentTypes( + thenValue.components[kCompleteOffsetsComponent].type, + elseValue.components[kCompleteOffsetsComponent].type))) return failure(); - // Identical symbolic offsets are available outside the if as an invariant. - if (thenValue.components[kOffsetsComponent].identity == - elseValue.components[kOffsetsComponent].identity) + // Identical symbolic offsets remain available outside the if. + if (thenValue.components[kCompleteOffsetsComponent].identity == + elseValue.components[kCompleteOffsetsComponent].identity) return SmallVector{}; - return SmallVector{kOffsetsComponent}; + return SmallVector{kCompleteOffsetsComponent}; } /// Chooses the offsets type carried by the replacement control-flow op. @@ -302,6 +310,8 @@ class TensorPtrDecomposePolicy final : public ControlFlowRewritePolicy { return isa(op); } + bool requiresPointerDescriptorBoundaryMarker() const override { return true; } + /// Materializes the same decomposition described by analyzeValue. This may /// create zero constants and integer additions, so it is called only after /// the complete control-flow subtree has passed read-only analysis. @@ -331,18 +341,19 @@ class TensorPtrDecomposePolicy final : public ControlFlowRewritePolicy { // ControlFlowRewrite will use this Value in the expanded SCF signature. OpBuilder::InsertionGuard guard(builder); builder.setInsertionPoint(addPtr); - Value offsets = createOffsetsAdd(builder, addPtr.getLoc(), - result->components[kOffsetsComponent], - context.remap(addPtr.getOffset())); + Value offsets = + createOffsetsAdd(builder, addPtr.getLoc(), + result->components[kCompleteOffsetsComponent], + context.remap(addPtr.getOffset())); if (!offsets) return failure(); result->originalType = value.getType(); - result->components[kOffsetsComponent] = offsets; + result->components[kCompleteOffsetsComponent] = offsets; return *result; } - // A scalar pointer splat materializes zero offsets while retaining the - // scalar source as the common-base invariant. + // A scalar pointer splat materializes zero offsets and records the + // scalar source as component 0. if (auto splat = value.getDefiningOp()) { if (!isa(splat.getSrc().getType())) return failure(); @@ -353,13 +364,12 @@ class TensorPtrDecomposePolicy final : public ControlFlowRewritePolicy { if (!offsets) return failure(); return DecomposedValue{value.getType(), - {offsets}, - {splat.getSrc()}, + {splat.getSrc(), offsets}, {builder.getBoolAttr(true)}}; } - // Fallback for an opaque tensor base: use the entire tensor-of-pointers as - // the invariant base and represent only subsequent displacement in offsets. + // Fallback for an opaque tensor base: keep the entire tensor-of-pointers + // as component 0 and represent subsequent displacement in component 1. if (!matches(value.getType())) return failure(); @@ -373,22 +383,23 @@ class TensorPtrDecomposePolicy final : public ControlFlowRewritePolicy { if (!offsets) return failure(); return DecomposedValue{ - value.getType(), {offsets}, {value}, {builder.getBoolAttr(false)}}; + value.getType(), {value, offsets}, {builder.getBoolAttr(false)}}; } - /// Rebuilds the original tensor-of-pointers from the invariant base and the + /// Rebuilds the original tensor-of-pointers from the base component and /// complete offsets selected/carried by the rewritten control flow. Value recompose(const DecomposedValue &value, OpBuilder &builder, Location loc) const override { if (!hasValidLayout(value)) return nullptr; - Value base = value.invariants[kBaseInvariant]; + Value base = value.components[kBaseComponent]; // addptr requires matching tensor lanes; broadcast a scalar common base // only when the decomposition recorded `base_is_scalar = true`. if (hasScalarBase(value)) base = builder.create(loc, value.originalType, base); return builder.create( - loc, value.originalType, base, value.components[kOffsetsComponent]); + loc, value.originalType, base, + value.components[kCompleteOffsetsComponent]); } }; @@ -400,10 +411,9 @@ namespace mlir::triton::controlflow { /// handling. The shared driver analyzes each outermost SCF root, then rewrites /// only the complete-offset components selected by this policy. LogicalResult runTensorPtrDecompose(ModuleOp module) { - // Carry only complete per-lane offsets through SCF. The common scalar base - // remains a rewrite invariant and is used to rebuild tensor-of-pointers at - // each region boundary. This decomposition is independent of - // BlockPtrDecompose. + // Carry only complete per-lane offsets through SCF. The base is a + // non-carried component used to rebuild tensor-of-pointers at each region + // boundary. This decomposition is independent of BlockPtrDecompose. // TODO: Replace this local extraction with TritonToUnstructure's common-base // analysis. Different or lane-wise bases must become explicit diagnostics // instead of pattern misses. diff --git a/third_party/ascend/lib/TritonToLinalg/BlockPtrAnalysis.cpp b/third_party/ascend/lib/TritonToLinalg/BlockPtrAnalysis.cpp index 3d21e90027..7544188b14 100644 --- a/third_party/ascend/lib/TritonToLinalg/BlockPtrAnalysis.cpp +++ b/third_party/ascend/lib/TritonToLinalg/BlockPtrAnalysis.cpp @@ -59,6 +59,25 @@ namespace mlir { namespace triton { +hivm::PointerCastOp createScalarPointerCast(OpBuilder &builder, Location loc, + MemRefType resultType, + Value address) { + SmallVector dynamicSizes; + if (resultType.getNumDynamicDims() != 0) { + Value defaultSize = builder.create(loc, 1); + dynamicSizes.assign(resultType.getNumDynamicDims(), defaultSize); + } + return builder.create( + loc, resultType, ValueRange{address}, ValueRange{dynamicSizes}); +} + +// Recognize original scalar-pointer producers whose converted result may act +// as a memref carrier. The parse site still requires BaseMemRefType, so merely +// appearing in this list never makes an unconverted pointer an opaque source. +static bool isScalarPointerTransport(Operation *op) { + return op && isa(op); +} + // MemAccType selectMaxMemAccTy(const MemAccType &v1, const MemAccType &v2) { // return (v1 > v2) ? v1 : v2; // } @@ -421,8 +440,12 @@ void BlockDataParser::parse( // if (isa(operand.getType())) { // Just consider two state: ptr and ptr> - auto remappedPtr = rewriter.getRemappedValue(operand); - assert(remappedPtr); + Value remappedPtr = rewriter.getRemappedValue(operand); + if (!remappedPtr) { + if (Operation *definingOp = operand.getDefiningOp()) + definingOp->emitError("scalar pointer has no converted value"); + return; + } if (auto op = operand.getDefiningOp()) { if (auto addPtrOp = dyn_cast(op)) { parseAddPtr(addPtrOp, data, loc, rewriter, known); @@ -439,11 +462,18 @@ void BlockDataParser::parse( data.setSource(remappedPtr); } else if (isDistributedTypeCustomOp(op)) { data.setSource(remappedPtr); + } else if (isScalarPointerTransport(op)) { + if (!isa(remappedPtr.getType())) { + op->emitError("scalar pointer transport did not convert to a memref"); + return; + } + data.setSource(remappedPtr); } else { - LLVM_DEBUG({ llvm::dbgs() << operand << "\n"; }); - llvm_unreachable("Unexpected operand defining operation, a scalar " - "pointer can only be produced by AddPtrOp or direct " - "block ptr or hivm CustomOp"); + op->emitError() << "unsupported scalar pointer producer '" + << op->getName() << "' with original type " + << operand.getType() << " and converted type " + << remappedPtr.getType(); + return; } } else { data.setSource(remappedPtr); @@ -1358,8 +1388,8 @@ void BlockDataParser::rewriteAddPtr( auto rtype = cast(intToPtrOp.getResult().getType()); auto memrefType = MemRefType::get({ShapedType::kDynamic}, rtype.getPointeeType()); - auto hivmPointCastOp = rewriter.create( - intToPtrOp.getLoc(), memrefType, ValueRange{intToPtrOp.getSrc()}); + auto hivmPointCastOp = createScalarPointerCast( + rewriter, intToPtrOp.getLoc(), memrefType, intToPtrOp.getSrc()); data.setSource(hivmPointCastOp.getResult()); } @@ -1389,19 +1419,28 @@ void BlockDataParser::rewriteAddPtr( rewriter.restoreInsertionPoint(insertPoint); } -OpFoldResult -accumulatePotentialOffsetOnBase(triton::MakeTensorPtrOp op, Value base, - OpFoldResult offset, - ConversionPatternRewriter &rewriter) { - if (auto baseRecast = base.getDefiningOp()) { - assert(isa(op.getBase().getDefiningOp()) && - "base of MakeTensorPtrOp only comes from native ptr or AddPtrOp"); - - return addOpFoldResult(offset, baseRecast.getConstifiedMixedOffset(), - op.getLoc(), rewriter, rewriter.getIndexType()); - } - - return offset; +static FailureOr +getBaseMemRefOffset(Value convertedBase, ConversionPatternRewriter &rewriter) { + auto memrefType = dyn_cast(convertedBase.getType()); + if (!memrefType) + return failure(); + + // Preserve the existing foldable path for a directly converted tt.addptr. + // Reinterpret-casting a reinterpret cast does not compose offsets, so the + // first cast's absolute offset must be carried into the new descriptor. + if (auto baseRecast = + convertedBase.getDefiningOp()) + return baseRecast.getConstifiedMixedOffset(); + + auto stridedLayout = memrefType.getStridesAndOffset(); + int64_t staticOffset = stridedLayout.second; + if (!ShapedType::isDynamic(staticOffset)) + return OpFoldResult(rewriter.getIndexAttr(staticOffset)); + + // A control-flow-carried BlockPtr base uses the canonical identity-layout + // memref and therefore has static offset zero. Dynamic hidden offsets are not + // valid BlockPtr bases; their displacement belongs in descriptor offsets. + return failure(); } void BlockDataParser::rewriteCustomOp( @@ -1419,8 +1458,8 @@ void BlockDataParser::rewriteCustomOp( auto rtype = cast(intToPtrOp.getResult().getType()); auto memrefType = MemRefType::get({ShapedType::kDynamic}, rtype.getPointeeType()); - auto hivmPointCastOp = rewriter.create( - intToPtrOp.getLoc(), memrefType, ValueRange{intToPtrOp.getSrc()}); + auto hivmPointCastOp = createScalarPointerCast( + rewriter, intToPtrOp.getLoc(), memrefType, intToPtrOp.getSrc()); if (data.getSizesRef().size() == 0) { data.getSizesRef().push_back(rewriter.getIndexAttr(1)); if (data.getScalarRef().isNull()) { @@ -1488,6 +1527,7 @@ void BlockDataParser::rewriteCustomOp( // Design for load/store boundary_check. memref::ReinterpretCastOp createRedundantOp(triton::MakeTensorPtrOp op, + OpFoldResult sourceBaseOffset, ConversionPatternRewriter &rewriter, BlockData &data) { auto loc = op.getLoc(); @@ -1507,10 +1547,10 @@ memref::ReinterpretCastOp createRedundantOp(triton::MakeTensorPtrOp op, // dim offset from base is initialized as zero. SmallVector curOffsets(op.getOffsets().size(), rewriter.getIndexAttr(0)); - // Just accumulate base potential offset - curOffsets.front() = accumulatePotentialOffsetOnBase( - op, rewriter.getRemappedValue(op.getBase()), curOffsets.front(), - rewriter); + // Both the full-shape descriptor and the final block descriptor use the + // same absolute source offset. Reusing this value avoids both dropping it + // across SCF and accidentally composing it twice. + curOffsets.front() = sourceBaseOffset; for (auto offset : curOffsets) { data.getOffsetsRef().push_back(offset); @@ -1532,25 +1572,39 @@ memref::ReinterpretCastOp createRedundantOp(triton::MakeTensorPtrOp op, return castOp; } -void BlockDataParser::rewriteMakeTensorPtrOp( - triton::MakeTensorPtrOp op, Value base, ConversionPatternRewriter &rewriter, +LogicalResult BlockDataParser::rewriteMakeTensorPtrOp( + triton::MakeTensorPtrOp op, Value convertedBase, + ConversionPatternRewriter &rewriter, llvm::SmallDenseMap &known) { + if (!convertedBase || !isa(convertedBase.getType())) { + op.emitOpError("expected the converted base to be a memref descriptor"); + return failure(); + } Location loc = op.getLoc(); BlockData data; - auto orderSize = op.getOrder().size(); - - // Handle base is defined by tt.bitcast + // Parse the original producer only for semantic information such as a + // bitcast element type. The runtime source always comes from the conversion + // adaptor so SCF-selected memref descriptors are not bypassed. BlockDataParser::parse(op.getBase(), data, loc, rewriter, known); + if (!data.hasSource()) { + op.emitOpError("failed to resolve the converted scalar base"); + return failure(); + } if (data.hasResElemTy()) { - auto memrefType = dyn_cast(data.getSourceRef().getType()) - .cloneWith(std::nullopt, data.getResElemTyRef()); + auto sourceType = dyn_cast(data.getSourceRef().getType()); + if (!sourceType) { + op.emitOpError("bitcast base did not resolve to a memref descriptor"); + return failure(); + } + auto memrefType = + sourceType.cloneWith(std::nullopt, data.getResElemTyRef()); UnrealizedConversionCastOp castOp = rewriter.create(loc, memrefType, data.getSourceRef()); data.setSource(castOp.getOutputs()[0]); } else { - data.setSource(rewriter.getRemappedValue(op.getBase())); + data.setSource(convertedBase); } data.getOffsetsRef() = @@ -1571,20 +1625,19 @@ void BlockDataParser::rewriteMakeTensorPtrOp( newOffsets.push_back(mulOpFoldResult(offset, stride, loc, rewriter, rewriter.getIndexType())); - // 1. Consider that current base ptr may comes from `triton::AddPtrOp`, - // which have been converted to `memref::ReinterpretCastOp` with 1D - // shape([1,]) by `AddPtrConverter`. - // 2. While here would also convert `triton::MakeTensorPtrOp` to - // `memref::ReinterpretCastOp`, it will create use-def on double recast - // which means offset&size&stride info of first one will be dropped in terms - // of memref recast op fold specification. - // - // Conclusion with above two: - // Base of MakeTensorPtrOp has been seen as origin base, so it should - // reserve offset of first recast if it exists. - // Here extract the offset of first recast and add it to highest dimension - newOffsets.front() = - accumulatePotentialOffsetOnBase(op, base, newOffsets.front(), rewriter); + if (newOffsets.empty()) { + op.emitOpError("expected at least one block pointer dimension"); + return failure(); + } + + FailureOr sourceBaseOffset = + getBaseMemRefOffset(convertedBase, rewriter); + if (failed(sourceBaseOffset)) { + op.emitOpError("could not extract the converted base offset"); + return failure(); + } + newOffsets.front() = addOpFoldResult(newOffsets.front(), *sourceBaseOffset, + loc, rewriter, rewriter.getIndexType()); data.getOffsetsRef().clear(); @@ -1610,7 +1663,7 @@ void BlockDataParser::rewriteMakeTensorPtrOp( // special handling for davinci // create redundant reinterpret_cast op for record shape info - auto redundantOp = createRedundantOp(op, rewriter, data); + auto redundantOp = createRedundantOp(op, *sourceBaseOffset, rewriter, data); redundantOp->setAttr("tensor_ptr_full_shape", rewriter.getUnitAttr()); // create reinterpret_cast op for the target block @@ -1700,6 +1753,8 @@ void BlockDataParser::rewriteMakeTensorPtrOp( } rewriter.create(loc, funcName, dstElemTy, args); } + + return success(); } void BlockDataParser::rewriteAdvanceOp( @@ -1967,6 +2022,26 @@ bool isUsedWithCondition(Value v, std::function cond, return false; } +// A loop-carried value may be consumed through a region argument, through a +// while after-argument, or only after the loop result. Check every semantic +// view of the same carried slot so an identity tensor.cast after the loop +// cannot hide an AddPtr/load/store use from the decomposition decision. +bool isLoopCarriedValueUsedWithCondition( + LoopLikeOpInterface loopOp, unsigned index, + const std::function &condition) { + if (index >= loopOp.getRegionIterArgs().size() || + index >= loopOp->getNumResults()) + return false; + if (isUsedWithCondition(loopOp.getRegionIterArgs()[index], condition)) + return true; + if (auto whileOp = dyn_cast(loopOp.getOperation())) { + if (index < whileOp.getAfterArguments().size() && + isUsedWithCondition(whileOp.getAfterArguments()[index], condition)) + return true; + } + return isUsedWithCondition(loopOp->getResult(index), condition); +} + // This function is util function for rewriteLoopOp that create value from data. // Assume data is structured, and from regionIterArg from LoopLikeOpInterface. // @@ -2030,9 +2105,10 @@ Value createFromData(RankedTensorType resType, const BlockData &data, return newRes; } -void BlockDataParser::rewriteLoopOp( - LoopLikeOpInterface op, ConversionPatternRewriter &rewriter, - llvm::SmallDenseMap &known) { +LogicalResult +BlockDataParser::rewriteLoopOp(LoopLikeOpInterface op, + ConversionPatternRewriter &rewriter, + llvm::SmallDenseMap &known) { SmallVector newInitArgs; SmallVector iterArgIdxMap; SmallVector maskIterArgs; @@ -2080,7 +2156,7 @@ void BlockDataParser::rewriteLoopOp( isa(cast(arg.getType()).getElementType()) && cast(cast(arg.getType()).getElementType()) .getWidth() != 1 && - isUsedWithCondition(op.getRegionIterArgs()[i], [](OpOperand *use) { + isLoopCarriedValueUsedWithCondition(op, i, [](OpOperand *use) { auto *user = use->getOwner(); return isa(user) || (isa(user) && use->getOperandNumber() == 1) || @@ -2102,7 +2178,7 @@ void BlockDataParser::rewriteLoopOp( maskIterArgs[i] = indexTensor && - isUsedWithCondition(op.getRegionIterArgs()[i], [](OpOperand *use) { + isLoopCarriedValueUsedWithCondition(op, i, [](OpOperand *use) { auto *user = use->getOwner(); return (isa(user) && use->getOperandNumber() == 1) || (isa(user) && use->getOperandNumber() == 2); @@ -2363,21 +2439,19 @@ void BlockDataParser::rewriteLoopOp( auto indexTensor = isa(resType) && isa(cast(resType).getElementType()) && - isUsedWithCondition(whileOp.getAfterArguments()[i], - [](OpOperand *use) { - auto *user = use->getOwner(); - return isa(user) || - (isa(user) && - use->getOperandNumber() == 1) || - (isa(user) && - use->getOperandNumber() == 2); - }); + isLoopCarriedValueUsedWithCondition(whileOp, i, [](OpOperand *use) { + auto *user = use->getOwner(); + return isa(user) || + (isa(user) && + use->getOperandNumber() == 1) || + (isa(user) && use->getOperandNumber() == 2); + }); if (indexTensor) { indexCnt += 2 * cast(resType).getRank(); usedForAfterRegionArgs.push_back(false); iterArgIdxMapForAfter.push_back(-1); - maskIterArgsForAfter[i] = isUsedWithCondition( - whileOp.getAfterArguments()[i], [](OpOperand *use) { + maskIterArgsForAfter[i] = + isLoopCarriedValueUsedWithCondition(whileOp, i, [](OpOperand *use) { auto *user = use->getOwner(); return (isa(user) && use->getOperandNumber() == 1) || @@ -2435,6 +2509,10 @@ void BlockDataParser::rewriteLoopOp( iterArgIdxMapForAfter, known); } + if (!newOp || newResults.size() != op->getNumResults()) + return op->emitError( + "loop rewrite produced a result list with incompatible arity"); + // Copy all attributes from op to newOp newOp->setAttrs(op->getAttrs()); rewriter.replaceOp(op, newResults); @@ -2456,17 +2534,20 @@ void BlockDataParser::rewriteLoopOp( dyn_cast(bodyOp)) { ConversionPatternRewriter::InsertionGuard guard(rewriter); rewriter.setInsertionPoint(makeTensorPtrOp); - rewriteMakeTensorPtrOp( - makeTensorPtrOp, - rewriter.getRemappedValue(makeTensorPtrOp.getBase()), rewriter, - known); + if (failed(rewriteMakeTensorPtrOp( + makeTensorPtrOp, + rewriter.getRemappedValue(makeTensorPtrOp.getBase()), rewriter, + known))) + return failure(); } else if (auto loopOp = dyn_cast(bodyOp); loopOp && !loopOp->hasAttr("ExtractedLoadOrStore")) { ConversionPatternRewriter::InsertionGuard guard(rewriter); rewriter.setInsertionPoint(loopOp); // Remove UnhandledLoopOp attr before process - loopOp->removeAttr("UnhandledLoopOp"); - rewriteLoopOp(loopOp, rewriter, known); + rewriter.modifyOpInPlace( + loopOp, [&]() { loopOp->removeAttr("UnhandledLoopOp"); }); + if (failed(rewriteLoopOp(loopOp, rewriter, known))) + return failure(); } } } @@ -2483,6 +2564,7 @@ void BlockDataParser::rewriteLoopOp( OpPrintingFlags().printGenericOpForm()); llvm::dbgs() << "\n"; }); + return success(); } /// @brief Rewrite the triton::AddPtrOp to handle unstructured memory access. diff --git a/third_party/ascend/lib/TritonToLinalg/TritonOpConverter.cpp b/third_party/ascend/lib/TritonToLinalg/TritonOpConverter.cpp index c36809373f..6b67f114e3 100644 --- a/third_party/ascend/lib/TritonToLinalg/TritonOpConverter.cpp +++ b/third_party/ascend/lib/TritonToLinalg/TritonOpConverter.cpp @@ -154,10 +154,294 @@ TransposeConverter::matchAndRewrite(triton::TransOp op, OpAdaptor adaptor, return success(); } +bool hasScalarPointerResult(scf::IfOp op) { + bool hasPointerResult = false; + for (Type resultType : op.getResultTypes()) { + auto pointerType = dyn_cast(resultType); + if (!pointerType) + continue; + if (isa(pointerType.getPointeeType())) + return false; + hasPointerResult = true; + } + return hasPointerResult; +} + +bool isScalarPointerSelect(arith::SelectOp op) { + auto pointerType = dyn_cast(op.getType()); + return pointerType && !isa(pointerType.getPointeeType()); +} + +// Convert a scalar Triton pointer type to the canonical memref descriptor used +// by TritonTypeConverter. It is reconstructed only after a complete byte +// address has crossed control flow, so it needs no dynamic layout. +static FailureOr +getScalarPointerCarrierType(Type originalType, + const TypeConverter &typeConverter) { + auto pointerType = dyn_cast(originalType); + if (!pointerType || isa(pointerType.getPointeeType())) + return failure(); + + Type convertedType = typeConverter.convertType(originalType); + if (!convertedType) + return failure(); + auto memrefType = dyn_cast(convertedType); + if (!memrefType) + return failure(); + + return memrefType; +} + +static FailureOr +getIfResultCarrierType(Type originalType, const TypeConverter &typeConverter) { + if (isa(originalType)) + return IntegerType::get(originalType.getContext(), 64); + + Type convertedType = typeConverter.convertType(originalType); + if (!convertedType) + return failure(); + return convertedType; +} + +// Materialize a no-op-compatible memref cast to the common carrier. Returning +// failure for non-memref or incompatible values prevents the pointer transport +// pattern from silently changing element, shape, rank, or memory-space types. +static FailureOr +castToMemRefCarrier(Value value, MemRefType carrierType, Location loc, + ConversionPatternRewriter &rewriter) { + if (value.getType() == carrierType) + return value; + + auto sourceType = dyn_cast(value.getType()); + if (!sourceType || + !memref::CastOp::areCastCompatible(sourceType, carrierType)) + return failure(); + + return rewriter.create(loc, carrierType, value).getResult(); +} + +// Dialect conversion may adapt a lane-local pointer descriptor such as +// `memref<1xT, strided<[1], offset: ?>>` to the canonical scalar-pointer +// carrier `memref`. Address materialization must inspect the original +// descriptor: its dynamic layout offset is part of the represented address and +// cannot be recovered from the shape-only carrier type. +// +// Only unwrap a one-to-one memref materialization that preserves rank, element +// type, and memory space. Other unrealized casts may represent a real element +// or address-space conversion and must remain visible to their converters. +static Value unwrapPointerDescriptorMaterialization(Value value) { + auto materialization = value.getDefiningOp(); + if (!materialization || materialization.getInputs().size() != 1 || + materialization.getOutputs().size() != 1) + return value; + + Value source = materialization.getInputs().front(); + auto sourceType = dyn_cast(source.getType()); + auto targetType = dyn_cast(value.getType()); + if (!sourceType || !targetType || + sourceType.getRank() != targetType.getRank() || + sourceType.getElementType() != targetType.getElementType() || + sourceType.getMemorySpace() != targetType.getMemorySpace()) + return value; + + return source; +} + +// Converts a memref descriptor into the complete byte address represented by +// that descriptor. extract_aligned_pointer_as_index yields the aligned buffer +// pointer; the descriptor's element offset must therefore be converted to +// bytes and added explicitly before the address crosses control flow. +static FailureOr +materializePointerAddress(Value value, Location loc, + ConversionPatternRewriter &rewriter) { + value = unwrapPointerDescriptorMaterialization(value); + auto memrefType = dyn_cast(value.getType()); + if (!memrefType) + return failure(); + Type elementType = memrefType.getElementType(); + if (!elementType.isIntOrFloat()) + return failure(); + + // A canonical scalar PointerCast already stores the complete byte address + // that IntToPtr received. Extracting its aligned pointer immediately and + // casting it back to i64 is an identity round trip. In current failing + // kernels this chain also survives into InjectSync, so folding it keeps a + // redundant region-local form out of the backend. Only fold the canonical + // one-dimensional carrier; descriptors with a non-zero offset or a non-unit + // stride still require the general address materialization below. + if (auto pointerCast = value.getDefiningOp()) { + auto [strides, offset] = memrefType.getStridesAndOffset(); + if (pointerCast.getAddrs().size() == 1 && memrefType.getRank() == 1 && + offset == 0 && strides.size() == 1 && strides.front() == 1) { + Value address = pointerCast.getAddrs().front(); + if (address.getType().isInteger(64)) + return address; + if (address.getType().isIndex()) + return rewriter + .create(loc, rewriter.getI64Type(), address) + .getResult(); + } + } + + Value address = + rewriter.create(loc, value); + int64_t staticOffset = memrefType.getStridesAndOffset().second; + if (staticOffset != 0) { + Value elementOffset; + if (ShapedType::isDynamic(staticOffset)) { + elementOffset = + rewriter.create(loc, value) + .getOffset(); + } else { + elementOffset = + rewriter.create(loc, staticOffset); + } + + int64_t elementBytes = (elementType.getIntOrFloatBitWidth() + 7) / 8; + if (elementBytes != 1) { + Value scale = rewriter.create(loc, elementBytes); + elementOffset = rewriter.create(loc, elementOffset, scale); + } + address = rewriter.create(loc, address, elementOffset); + } + return rewriter + .create(loc, rewriter.getI64Type(), address) + .getResult(); +} + +LogicalResult +IfConverter::matchAndRewrite(scf::IfOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + const TypeConverter *typeConverter = getTypeConverter(); + if (!typeConverter) + return rewriter.notifyMatchFailure(op, "requires a type converter"); + if (!hasScalarPointerResult(op)) + return failure(); + + SmallVector convertedResultTypes; + convertedResultTypes.reserve(op.getNumResults()); + for (Type resultType : op.getResultTypes()) { + FailureOr convertedType = + getIfResultCarrierType(resultType, *typeConverter); + if (failed(convertedType)) + return rewriter.notifyMatchFailure(op, + "could not convert an if result type"); + convertedResultTypes.push_back(*convertedType); + } + + auto newIfOp = rewriter.create(op.getLoc(), convertedResultTypes, + adaptor.getCondition(), + /*withElseRegion=*/true); + newIfOp->setAttrs(op->getAttrs()); + newIfOp->setAttr(kScalarPointerCarrierBoundaryAttr, + UnitAttr::get(rewriter.getContext())); + + // Move the original regions instead of cloning them. Besides preserving + // side effects, this keeps nested operations in the conversion driver's + // worklist so their operands are remapped normally. + rewriter.eraseBlock(newIfOp.thenBlock()); + rewriter.eraseBlock(newIfOp.elseBlock()); + rewriter.inlineRegionBefore(op.getThenRegion(), newIfOp.getThenRegion(), + newIfOp.getThenRegion().end()); + rewriter.inlineRegionBefore(op.getElseRegion(), newIfOp.getElseRegion(), + newIfOp.getElseRegion().end()); + + SmallVector replacementResults; + replacementResults.reserve(op.getNumResults()); + rewriter.setInsertionPointAfter(newIfOp); + for (auto [originalResultType, newResult] : + llvm::zip(op.getResultTypes(), newIfOp.getResults())) { + if (!isa(originalResultType)) { + replacementResults.push_back(newResult); + continue; + } + FailureOr resultType = + getScalarPointerCarrierType(originalResultType, *typeConverter); + if (failed(resultType)) + return rewriter.notifyMatchFailure( + op, "could not reconstruct a scalar pointer result"); + replacementResults.push_back( + createScalarPointerCast(rewriter, op.getLoc(), *resultType, newResult) + .getResult()); + } + rewriter.replaceOp(op, replacementResults); + return success(); +} + +LogicalResult PointerSelectConverter::matchAndRewrite( + arith::SelectOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + if (!isScalarPointerSelect(op)) + return failure(); + + const TypeConverter *typeConverter = getTypeConverter(); + if (!typeConverter) + return rewriter.notifyMatchFailure(op, "requires a type converter"); + + FailureOr resultType = + getScalarPointerCarrierType(op.getType(), *typeConverter); + if (failed(resultType)) + return rewriter.notifyMatchFailure( + op, "could not build a scalar pointer memref carrier"); + + FailureOr trueAddress = + materializePointerAddress(adaptor.getTrueValue(), op.getLoc(), rewriter); + FailureOr falseAddress = + materializePointerAddress(adaptor.getFalseValue(), op.getLoc(), rewriter); + if (failed(trueAddress) || failed(falseAddress)) + return rewriter.notifyMatchFailure( + op, "selected pointer values have no complete integer address"); + + auto selectedAddress = rewriter.create( + op.getLoc(), rewriter.getI64Type(), adaptor.getCondition(), *trueAddress, + *falseAddress); + selectedAddress->setAttrs(op->getAttrs()); + auto pointerCast = createScalarPointerCast(rewriter, op.getLoc(), *resultType, + selectedAddress.getResult()); + rewriter.replaceOp(op, pointerCast.getResult()); + return success(); +} + LogicalResult YieldConverter::matchAndRewrite(scf::YieldOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const { - rewriter.replaceOpWithNewOp(op, adaptor.getOperands()); + auto parentIf = dyn_cast(op->getParentOp()); + if (!parentIf || !parentIf->hasAttr(kScalarPointerCarrierBoundaryAttr)) + return failure(); + + SmallVector convertedOperands(adaptor.getOperands()); + + if (parentIf.getNumResults() != convertedOperands.size()) + return rewriter.notifyMatchFailure(op, "yield/result arity does not match"); + + for (auto [index, targetType] : llvm::enumerate(parentIf.getResultTypes())) { + Value &operand = convertedOperands[index]; + if (operand.getType() == targetType) + continue; + + if (targetType.isInteger(64) && isa(operand.getType())) { + FailureOr address = + materializePointerAddress(operand, op.getLoc(), rewriter); + if (failed(address)) + return rewriter.notifyMatchFailure( + op, "could not materialize a yielded pointer address"); + operand = *address; + continue; + } + auto targetMemrefType = dyn_cast(targetType); + if (!targetMemrefType) + return rewriter.notifyMatchFailure( + op, "converted yield operand is incompatible with if result"); + + FailureOr casted = + castToMemRefCarrier(operand, targetMemrefType, op.getLoc(), rewriter); + if (failed(casted)) + return rewriter.notifyMatchFailure( + op, "converted yield operand is incompatible with if result"); + operand = *casted; + } + + rewriter.replaceOpWithNewOp(op, convertedOperands); return success(); } @@ -178,9 +462,8 @@ LogicalResult MakeTensorPtrConverter::matchAndRewrite( triton::MakeTensorPtrOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const { llvm::SmallDenseMap known; - BlockDataParser::rewriteMakeTensorPtrOp(op, adaptor.getBase(), rewriter, - known); - return success(); + return BlockDataParser::rewriteMakeTensorPtrOp(op, adaptor.getBase(), + rewriter, known); } LogicalResult PreciseDivConverter::matchAndRewrite( @@ -3102,26 +3385,41 @@ DotScaledConverter::matchAndRewrite(triton::DotScaledOp op, OpAdaptor adaptor, } LogicalResult -PtrToIntConverter::matchAndRewrite(triton::PtrToIntOp op, OpAdaptor adaptor, +IntToPtrConverter::matchAndRewrite(triton::IntToPtrOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const { - auto loc = op.getLoc(); - Value ptr = adaptor.getSrc(); + auto pointerType = dyn_cast(op.getType()); + if (!pointerType || isa(pointerType.getPointeeType())) + return rewriter.notifyMatchFailure( + op, "only scalar pointer reconstruction is supported"); + + const TypeConverter *typeConverter = getTypeConverter(); + if (!typeConverter) + return rewriter.notifyMatchFailure(op, "requires a type converter"); + auto resultType = + dyn_cast(typeConverter->convertType(op.getType())); + if (!resultType) + return rewriter.notifyMatchFailure( + op, "pointer result did not convert to a memref type"); + + // Rebuild the memref only after the integer address has crossed control + // flow. Selecting an i64 address is a pure SSA operation and avoids the + // backend interpreting a memref-valued merge as a GM-to-UB copy. + auto pointerCast = createScalarPointerCast(rewriter, op.getLoc(), resultType, + adaptor.getSrc()); + rewriter.replaceOp(op, pointerCast.getResult()); + return success(); +} - if (!mlir::isa(ptr.getType())) { +LogicalResult +PtrToIntConverter::matchAndRewrite(triton::PtrToIntOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + FailureOr address = + materializePointerAddress(adaptor.getSrc(), op.getLoc(), rewriter); + if (failed(address)) { return rewriter.notifyMatchFailure(op, "input is not a memref type"); } - auto resultType = op.getType(); - - // memref.extract_aligned_pointer_as_index is used to obtain the integer - // representation of the base address. - auto ptrToIndexOp = - rewriter.create(loc, ptr); - - Value intResult = - rewriter.create(loc, resultType, ptrToIndexOp); - - rewriter.replaceOp(op, intResult); + rewriter.replaceOp(op, *address); return success(); } diff --git a/third_party/ascend/lib/TritonToLinalg/TritonToLinalgPass.cpp b/third_party/ascend/lib/TritonToLinalg/TritonToLinalgPass.cpp index 6d055cc01a..ff73da21b1 100644 --- a/third_party/ascend/lib/TritonToLinalg/TritonToLinalgPass.cpp +++ b/third_party/ascend/lib/TritonToLinalg/TritonToLinalgPass.cpp @@ -23,6 +23,7 @@ #include +#include "TritonControlFlowOpt/ControlFlowRewrite.h" #include "TritonToLinalg/BlockPtrAnalysis.h" #include "ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.h" #include "ascend/include/TritonToLinalg/ArgMinMaxConverter.h" @@ -74,7 +75,9 @@ #include "mlir/Transforms/Passes.h" #include "llvm/ADT/BitVector.h" +#include "llvm/ADT/DenseSet.h" #include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/SmallVectorExtras.h" #include "llvm/ADT/Twine.h" @@ -85,6 +88,7 @@ #include #include +#include #include #define DEBUG_TYPE "triton-to-linalg" @@ -96,6 +100,258 @@ int nd2nzFlag = 0; bool compileOn91095Flag = false; bool existDotFlag = false; +static bool containsTritonPointer(Type type) { + if (isa(type)) + return true; + auto shapedType = dyn_cast(type); + return shapedType && isa(shapedType.getElementType()); +} + +static bool hasPointerFreeBoundary(Operation *op) { + return llvm::none_of(op->getOperandTypes(), containsTritonPointer) && + llvm::none_of(op->getResultTypes(), containsTritonPointer); +} + +// Adds only the values that can produce one SSA result. Structured control +// flow is followed by result index so preserving one pointer descriptor does +// not retain unrelated values yielded by the same operation. +static LogicalResult +appendProducerOperands(Value value, SmallVectorImpl &producerWorklist) { + if (auto argument = dyn_cast(value)) { + Operation *parent = argument.getOwner()->getParentOp(); + unsigned argumentIndex = argument.getArgNumber(); + if (auto forOp = dyn_cast_or_null(parent)) { + if (argumentIndex == 0) + return success(); + unsigned iterIndex = argumentIndex - 1; + if (iterIndex >= forOp.getInitArgs().size()) + return failure(); + producerWorklist.push_back(forOp.getInitArgs()[iterIndex]); + producerWorklist.push_back(forOp.getYieldedValues()[iterIndex]); + return success(); + } + if (auto whileOp = dyn_cast_or_null(parent)) { + if (argumentIndex >= whileOp.getInits().size()) + return failure(); + if (argument.getOwner() == whileOp.getBeforeBody()) { + producerWorklist.push_back(whileOp.getInits()[argumentIndex]); + producerWorklist.push_back( + whileOp.getYieldOp().getOperand(argumentIndex)); + return success(); + } + if (argument.getOwner() == whileOp.getAfterBody()) { + producerWorklist.push_back( + whileOp.getConditionOp().getArgs()[argumentIndex]); + return success(); + } + } + return success(); + } + + Operation *producer = value.getDefiningOp(); + if (!producer) + return success(); + + auto result = dyn_cast(value); + if (auto forOp = dyn_cast(producer)) { + if (!result || result.getResultNumber() >= forOp.getNumResults()) + return failure(); + unsigned index = result.getResultNumber(); + producerWorklist.push_back(forOp.getInitArgs()[index]); + producerWorklist.push_back(forOp.getYieldedValues()[index]); + return success(); + } + if (auto whileOp = dyn_cast(producer)) { + if (!result || result.getResultNumber() >= whileOp.getNumResults()) + return failure(); + unsigned index = result.getResultNumber(); + producerWorklist.push_back(whileOp.getInits()[index]); + producerWorklist.push_back(whileOp.getConditionOp().getArgs()[index]); + producerWorklist.push_back(whileOp.getYieldOp().getOperand(index)); + return success(); + } + if (auto ifOp = dyn_cast(producer)) { + if (!result || result.getResultNumber() >= ifOp.getNumResults() || + !ifOp.elseBlock()) + return failure(); + unsigned index = result.getResultNumber(); + producerWorklist.push_back(ifOp.thenYield().getOperand(index)); + producerWorklist.push_back(ifOp.elseYield().getOperand(index)); + return success(); + } + + producerWorklist.append(producer->operand_begin(), producer->operand_end()); + return success(); +} + +// PointerDescriptorBoundary records exactly which loop-carried slots rebuild +// Triton pointers. Preserve only those init/condition/yield producer chains so +// MetaUseEraser cannot leave dangling descriptor operands without retaining +// unrelated accumulators, masks, bounds, or ordinary tensor computations. +static LogicalResult preservePointerDescriptorComputations(ModuleOp moduleOp) { + bool valid = true; + moduleOp.walk([&](LoopLikeOpInterface loopOp) { + Operation *loop = loopOp.getOperation(); + auto descriptorSlots = loop->getAttrOfType( + controlflow::kPointerDescriptorBoundaryAttr); + if (!loop->hasAttr(controlflow::kPointerDescriptorBoundaryAttr)) + return; + if (!descriptorSlots) { + valid = false; + return; + } + + loop->removeAttr("MetaUse"); + + SmallVector producerWorklist; + auto appendSlotValues = [&](ValueRange values, int32_t slot) { + if (slot < 0 || static_cast(slot) >= values.size()) { + valid = false; + return; + } + producerWorklist.push_back(values[slot]); + }; + + auto appendLoopSlotValues = [&](int32_t slot) { + if (auto forOp = dyn_cast(loop)) { + appendSlotValues(forOp.getInitArgs(), slot); + appendSlotValues(forOp.getYieldedValues(), slot); + return; + } + if (auto whileOp = dyn_cast(loop)) { + appendSlotValues(whileOp.getInits(), slot); + appendSlotValues(whileOp.getConditionOp().getArgs(), slot); + appendSlotValues(whileOp.getYieldOp().getOperands(), slot); + return; + } + valid = false; + }; + + llvm::SmallDenseSet seenSlots; + for (int32_t slot : descriptorSlots.asArrayRef()) { + if (!seenSlots.insert(slot).second) { + valid = false; + continue; + } + appendLoopSlotValues(slot); + } + + // A mixed loop can contain both CFO-expanded descriptor slots and a + // pointer slot that was invariant for a later policy. The latter still + // belongs to the legacy loop conversion and its producer must survive + // MetaUse erasure. Preserve exactly those residual pointer slots, without + // adding unrelated tensor/scalar loop state to the descriptor marker. + for (auto [slot, type] : llvm::enumerate(loop->getResultTypes())) { + if (!containsTritonPointer(type)) + continue; + if (slot > static_cast(std::numeric_limits::max())) { + valid = false; + continue; + } + appendLoopSlotValues(static_cast(slot)); + } + + llvm::DenseSet visitedValues; + + while (!producerWorklist.empty()) { + Value value = producerWorklist.pop_back_val(); + if (!value || !visitedValues.insert(value).second) + continue; + Operation *producer = value.getDefiningOp(); + if (producer) + producer->removeAttr("MetaUse"); + if (failed(appendProducerOperands(value, producerWorklist))) + valid = false; + } + }); + return success(valid); +} + +// Recomputes the result layouts of subviews whose source descriptor has been +// rebased. A memref.subview result layout is derived from both its mixed +// offsets/strides and its source layout. For example, changing the source from +// `memref<32xf32, strided<[1], offset: ?>>` to the equivalent rebased +// `memref<32xf32, strided<[1]>>` changes a zero-offset subview result from a +// dynamic offset to offset zero. Merely changing the source SSA value leaves +// the old result type behind and makes SubViewOp verification fail. +// +// Subviews may be chained, so every updated result becomes a new worklist +// source. Rank-reduced subviews retain their existing result shape while their +// layout is inferred again from the rebased source. +static LogicalResult propagateRebasedSubviewTypes(Value rebasedSource, + IRRewriter &rewriter) { + SmallVector sources{rebasedSource}; + llvm::SmallPtrSet visited; + + while (!sources.empty()) { + Value source = sources.pop_back_val(); + for (Operation *user : source.getUsers()) { + auto subview = dyn_cast(user); + if (!subview || !visited.insert(user).second) + continue; + + auto sourceType = dyn_cast(subview.getSource().getType()); + auto oldResultType = dyn_cast(subview.getResult().getType()); + if (!sourceType || !oldResultType) + return failure(); + + Type inferredType; + if (sourceType.getRank() == oldResultType.getRank()) { + inferredType = memref::SubViewOp::inferResultType( + sourceType, subview.getMixedOffsets(), subview.getMixedSizes(), + subview.getMixedStrides()); + } else { + inferredType = memref::SubViewOp::inferRankReducedResultType( + oldResultType.getShape(), sourceType, subview.getMixedOffsets(), + subview.getMixedSizes(), subview.getMixedStrides()); + } + + auto inferredMemRefType = dyn_cast(inferredType); + if (!inferredMemRefType) + return failure(); + if (inferredMemRefType != oldResultType) { + rewriter.modifyOpInPlace( + subview, [&] { subview.getResult().setType(inferredMemRefType); }); + } + sources.push_back(subview.getResult()); + } + } + return success(); +} + +// Returns true when rebasing a descriptor would reach a user whose type +// contract is owned elsewhere. Only a subview chain ending in direct +// memref.load/store operations is proven local here. Every other user, +// including SCF/CFG terminators, function calls/returns and dialect-specific +// fixed-type operations, conservatively keeps the original descriptor layout. +static bool reachesLayoutSensitiveBoundary(Value root) { + SmallVector worklist{root}; + llvm::DenseSet visited; + while (!worklist.empty()) { + Value value = worklist.pop_back_val(); + if (!visited.insert(value).second) + continue; + for (Operation *user : value.getUsers()) { + if (auto subview = dyn_cast(user)) { + if (subview.getSource() == value) { + worklist.push_back(subview.getResult()); + continue; + } + } + if (auto load = dyn_cast(user)) { + if (load.getMemRef() == value) + continue; + } + if (auto store = dyn_cast(user)) { + if (store.getMemRef() == value) + continue; + } + return true; + } + } + return false; +} + // Convert structured custom ops after operand type converted, // for example tt.ptr converted to memref. template @@ -563,35 +819,58 @@ void TritonToLinalgPass::addDynamicLegal( return true; }); - target.addDynamicallyLegalOp([](Operation *op) { + target.addDynamicallyLegalOp( + [](scf::IfOp op) { return !TTOpConverters::hasScalarPointerResult(op); }); + + auto controlFlowTerminatorLegal = [](Operation *op) { + Operation *parent = op->getParentOp(); + if (parent && + parent->hasAttr(TTOpConverters::kScalarPointerCarrierBoundaryAttr)) { + auto parentIf = cast(parent); + return llvm::equal(op->getOperandTypes(), parentIf.getResultTypes()); + } + + if (parent && parent->hasAttr(controlflow::kPointerDescriptorBoundaryAttr)) + return hasPointerFreeBoundary(op); + return llvm::all_of(op->getOperandTypes(), [](Type t) { - if (isa(t)) { + if (isa(t)) return false; - } - if (auto shapedType = dyn_cast(t)) { + if (auto shapedType = dyn_cast(t)) return shapedType.getElementType().isIntOrFloat(); - } assert(t.isIntOrIndexOrFloat()); return true; }); - }); + }; - target.addDynamicallyLegalDialect( - [this](Operation *op) { - if (op->hasAttr("MetaUse")) { - return false; - } + target.addDynamicallyLegalOp( + controlFlowTerminatorLegal); - if (isa(op)) { - return true; - } + auto isArithOrMathOpLegal = [this](Operation *op) { + if (op->hasAttr("MetaUse")) + return false; + + if (isa(op)) + return true; + + bool operateOnTensors = llvm::all_of(op->getOperandTypes(), [](Type type) { + return isa(type); + }); - bool operateOnTensors = - llvm::all_of(op->getOperandTypes(), - [](Type type) { return isa(type); }); + return this->namedOps || !operateOnTensors; + }; - return this->namedOps || !operateOnTensors; + // Numeric selects retain the existing Arith legality. Every scalar-pointer + // select uses the integer-address converter so no memref object crosses it. + target.addDynamicallyLegalOp( + [isArithOrMathOpLegal](arith::SelectOp op) { + if (TTOpConverters::isScalarPointerSelect(op)) + return false; + return isArithOrMathOpLegal(op); }); + + target.addDynamicallyLegalDialect( + isArithOrMathOpLegal); } void TritonToLinalgPass::populateTritonToLinalgCanonicalizationPatterns( @@ -710,11 +989,16 @@ void TritonToLinalgPass::populateTritonToLinalgConversionPatterns( patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); - patterns.add>( - patterns.getContext()); + patterns.add>(patterns.getContext(), + PatternBenefit(2)); patterns.add>( - patterns.getContext()); - patterns.add(patterns.getContext()); + patterns.getContext(), PatternBenefit(2)); + patterns.add(patterns.getContext(), + PatternBenefit(2)); + patterns.add( + typeConverter, patterns.getContext(), PatternBenefit(2)); + patterns.add(typeConverter, + patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); @@ -722,6 +1006,8 @@ void TritonToLinalgPass::populateTritonToLinalgConversionPatterns( patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); + patterns.add(typeConverter, + patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); @@ -1043,6 +1329,12 @@ void TritonToLinalgPass::runOnOperation() { } }); + if (failed(preservePointerDescriptorComputations(moduleOp))) { + moduleOp->emitError("invalid PointerDescriptorBoundary slot metadata"); + signalPassFailure(); + return; + } + RewritePatternSet patterns(&getContext()); ConversionTarget target(getContext()); TritonTypeConverter tritonTypeConverter{}; @@ -1051,8 +1343,11 @@ void TritonToLinalgPass::runOnOperation() { this->addDynamicLegal(target, tritonTypeConverter); // 4. Mark ops that must be converted explicitly (e.g. tt.scan). - auto loopOpLegalFn = [](LoopLikeOpInterface op) { - return !op.getOperation()->hasAttr("UnhandledLoopOp"); + auto loopOpLegalFn = [](LoopLikeOpInterface loopOp) { + Operation *op = loopOp.getOperation(); + if (op->hasAttr(controlflow::kPointerDescriptorBoundaryAttr)) + return hasPointerFreeBoundary(op); + return !op->hasAttr("UnhandledLoopOp"); }; target.addIllegalOp(); @@ -1076,9 +1371,21 @@ void TritonToLinalgPass::runOnOperation() { moduleOp.walk([this](LoopLikeOpInterface loopOp) { auto *op = loopOp.getOperation(); - if (!op->hasAttr("ExtractedLoadOrStore")) + // A marker identifies the slots already expanded by CFO. Skip the legacy + // loop rewrite only when that expansion made the complete boundary + // pointer-free. A mixed loop may still contain an invariant pointer slot; + // the legacy converter must process that remaining slot without losing the + // marker's precise producer-preservation roots. + bool hasPointerDescriptorSlots = + op->hasAttr(mlir::triton::controlflow::kPointerDescriptorBoundaryAttr); + bool hasCompletePointerFreeBoundary = + hasPointerDescriptorSlots && hasPointerFreeBoundary(op); + if (!op->hasAttr("ExtractedLoadOrStore") && !hasCompletePointerFreeBoundary) op->setAttr("UnhandledLoopOp", UnitAttr::get(op->getContext())); + if (hasCompletePointerFreeBoundary) + return; + for (auto res : loopOp->getResults()) { if (auto tensorType = dyn_cast(res.getType()); tensorType && @@ -1093,7 +1400,18 @@ void TritonToLinalgPass::runOnOperation() { }); // 7. Convert ops. - if (failed(applyPartialConversion(moduleOp, target, std::move(patterns)))) { + LogicalResult conversionResult = + applyPartialConversion(moduleOp, target, std::move(patterns)); + + // The marker is a contract between CFO and this conversion pass; no + // downstream dialect should observe this implementation detail. + moduleOp.walk([](LoopLikeOpInterface loopOp) { + loopOp->removeAttr(controlflow::kPointerDescriptorBoundaryAttr); + }); + moduleOp.walk([](scf::IfOp ifOp) { + ifOp->removeAttr(TTOpConverters::kScalarPointerCarrierBoundaryAttr); + }); + if (failed(conversionResult)) { moduleOp->emitError("failed to apply Conversion Patterns"); signalPassFailure(); } @@ -1141,34 +1459,31 @@ void TritonToLinalgPass::runOnOperation() { moduleOp.walk([&](hivm::PointerCastOp op) { castOps.push_back(op); }); for (auto op : castOps) { - SmallVector userOps(op->getUsers().begin(), - op->getUsers().end()); + SmallVector reinterpretCastOps; + for (Operation *user : op->getUsers()) { + if (auto reinterpretCast = dyn_cast(user)) + reinterpretCastOps.push_back(reinterpretCast); + } + if (reinterpretCastOps.empty()) + continue; + IRRewriter rewriter(&getContext()); rewriter.setInsertionPointAfter(op); Value addr = op.getAddrs()[0]; auto elementType = cast(op.getResult().getType()).getElementType(); - Value elementTypeSize; - if (auto intType = dyn_cast(elementType)) { - elementTypeSize = rewriter.create( - op.getLoc(), - rewriter.getIntegerAttr(addr.getType(), intType.getWidth() / 8)); - } else if (auto floatType = dyn_cast(elementType)) { - elementTypeSize = rewriter.create( - op.getLoc(), - rewriter.getIntegerAttr(addr.getType(), floatType.getWidth() / 8)); - } else { - llvm_unreachable("Cannot get memory size"); - } - for (auto userOp : userOps) { - auto reinterpretCastOp = cast(userOp); + for (memref::ReinterpretCastOp reinterpretCastOp : reinterpretCastOps) { auto sizes = reinterpretCastOp.getStaticSizes(); auto staticStrides = reinterpretCastOp.getStaticStrides(); auto strides = reinterpretCastOp.getStrides(); - if (reinterpretCastOp.getStaticOffsets().size() != 1) - userOp->emitError("IntToPtrOp must converted to PointerCastOp of " - "memref type"); + if (reinterpretCastOp.getStaticOffsets().size() != 1) { + reinterpretCastOp->emitError( + "IntToPtrOp must converted to PointerCastOp of " + "memref type"); + signalPassFailure(); + return; + } int64_t castOpSize = 0; SmallVector dynamicSizes; for (const auto &[size, stride] : llvm::zip_equal(sizes, staticStrides)) { @@ -1190,53 +1505,108 @@ void TritonToLinalgPass::runOnOperation() { dynamicSize = rewriter.create(op.getLoc(), dynamicSize, axisSize); } - Value offsetValue; + Value offsetElements; auto staticOffset = reinterpretCastOp.getStaticOffsets()[0]; if (ShapedType::isDynamic(staticOffset)) { - offsetValue = reinterpretCastOp.getOffsets()[0]; - if (offsetValue.getType() != addr.getType()) - offsetValue = rewriter.create( - op.getLoc(), addr.getType(), offsetValue); + offsetElements = reinterpretCastOp.getOffsets()[0]; } else { - offsetValue = rewriter.create( - op.getLoc(), rewriter.getIntegerAttr(addr.getType(), staticOffset)); + offsetElements = + rewriter.create(op.getLoc(), staticOffset); } - offsetValue = rewriter.create(op.getLoc(), offsetValue, - elementTypeSize); - Value realAddr = - rewriter.create(op.getLoc(), addr, offsetValue); + auto memrefType = MemRefType::get({ShapedType::kDynamic}, elementType); - auto newCastOp = rewriter.create( - op.getLoc(), memrefType, realAddr, dynamicSize); - auto markOp = rewriter.create(op.getLoc(), - newCastOp.getResult()); - markOp->setAttr(hivm::AddressSpaceAttr::getMnemonic(), + auto createPointerCast = [&](Value address, Value capacity) { + auto cast = rewriter.create( + op.getLoc(), memrefType, address, capacity); + auto mark = + rewriter.create(op.getLoc(), cast.getResult()); + mark->setAttr(hivm::AddressSpaceAttr::getMnemonic(), {hivm::AddressSpaceAttr::get(rewriter.getContext(), hivm::AddressSpace::GM)}); - - // update result offset - auto origResultType = - cast(reinterpretCastOp.getResult().getType()); - MemRefType newResultType = origResultType; - if (auto stridedLayout = - dyn_cast(origResultType.getLayout())) { - int64_t offset = stridedLayout.getOffset(); - if (!ShapedType::isDynamic(offset)) { - auto newLayout = StridedLayoutAttr::get(rewriter.getContext(), 0, - stridedLayout.getStrides()); - newResultType = MemRefType::get( - origResultType.getShape(), origResultType.getElementType(), - newLayout, origResultType.getMemorySpace()); + return cast; + }; + + if (reachesLayoutSensitiveBoundary(reinterpretCastOp.getResult())) { + // Keep the original descriptor type and offset at externally typed + // boundaries. Because the PointerCast still starts at the old base, + // its element capacity must cover both the leading offset and view. + Value pointerCapacity = dynamicSize; + if (offsetElements.getType() != pointerCapacity.getType()) + offsetElements = rewriter.create( + op.getLoc(), pointerCapacity.getType(), offsetElements); + Value leadingExtent = offsetElements; + if (ShapedType::isDynamic(staticOffset)) { + Value zero = rewriter.create(op.getLoc(), 0); + leadingExtent = rewriter.create(op.getLoc(), + offsetElements, zero); } + if (ShapedType::isDynamic(staticOffset) || staticOffset > 0) + pointerCapacity = rewriter.create( + op.getLoc(), pointerCapacity, leadingExtent); + auto newCastOp = createPointerCast(addr, pointerCapacity); + rewriter.modifyOpInPlace(reinterpretCastOp, [&] { + reinterpretCastOp.getSourceMutable().assign(newCastOp.getResult()); + }); + continue; } - rewriter.replaceOpWithNewOp( - reinterpretCastOp, newResultType, newCastOp, ValueRange({}), - reinterpretCastOp.getSizes(), reinterpretCastOp.getStrides(), - SmallVector({0}), reinterpretCastOp.getStaticSizes(), - reinterpretCastOp.getStaticStrides()); + Value offsetValue = offsetElements; + if (offsetValue.getType() != addr.getType()) + offsetValue = rewriter.create( + op.getLoc(), addr.getType(), offsetValue); + Value elementTypeSize; + if (auto intType = dyn_cast(elementType)) { + elementTypeSize = rewriter.create( + op.getLoc(), + rewriter.getIntegerAttr(addr.getType(), intType.getWidth() / 8)); + } else if (auto floatType = dyn_cast(elementType)) { + elementTypeSize = rewriter.create( + op.getLoc(), + rewriter.getIntegerAttr(addr.getType(), floatType.getWidth() / 8)); + } else { + llvm_unreachable("Cannot get memory size"); + } + offsetValue = rewriter.create(op.getLoc(), offsetValue, + elementTypeSize); + Value realAddr = + rewriter.create(op.getLoc(), addr, offsetValue); + auto newCastOp = createPointerCast(realAddr, dynamicSize); + // realAddr already includes the old reinterpret-cast offset in bytes. + // The replacement view therefore starts at offset zero, and its result + // type must describe the same rebased layout. Reusing the old type here + // would combine static_offsets=[0] with (for example) a type-level + // offset of 1, which is rejected by the ReinterpretCast verifier. + auto oldResultType = + cast(reinterpretCastOp.getResult().getType()); + SmallVector rebasedStrides( + oldResultType.getStridesAndOffset().first); + auto rebasedResultType = MemRefType::get( + oldResultType.getShape(), oldResultType.getElementType(), + StridedLayoutAttr::get(oldResultType.getContext(), /*offset=*/0, + rebasedStrides), + oldResultType.getMemorySpace()); + + // Keep the old result and replacement types equal while RAUW updates all + // users to the rebased descriptor type. + rewriter.modifyOpInPlace(reinterpretCastOp, [&] { + reinterpretCastOp.getResult().setType(rebasedResultType); + }); + auto rebasedReinterpretCast = + rewriter.replaceOpWithNewOp( + reinterpretCastOp, rebasedResultType, newCastOp, ValueRange({}), + reinterpretCastOp.getSizes(), reinterpretCastOp.getStrides(), + SmallVector({0}), reinterpretCastOp.getStaticSizes(), + reinterpretCastOp.getStaticStrides()); + if (failed(propagateRebasedSubviewTypes( + rebasedReinterpretCast.getResult(), rewriter))) { + rebasedReinterpretCast.emitError( + "failed to propagate rebased layout through subview users"); + signalPassFailure(); + return; + } } - rewriter.eraseOp(op); + if (op->use_empty()) + rewriter.eraseOp(op); } // Try interleave optimization diff --git a/third_party/ascend/lib/TritonToUnstructure/BubbleUpOperation.cpp b/third_party/ascend/lib/TritonToUnstructure/BubbleUpOperation.cpp index 3dd7d2fa1b..cb955c9f1d 100644 --- a/third_party/ascend/lib/TritonToUnstructure/BubbleUpOperation.cpp +++ b/third_party/ascend/lib/TritonToUnstructure/BubbleUpOperation.cpp @@ -87,6 +87,8 @@ BubbleUpExtract::matchAndRewrite(ExtractOpTy op, bubbleUpIntBinaryOp(op, orIOp, loc, rewriter); } else if (auto cmpIOp = dyn_cast(parentOp)) { bubbleUpOperation(op, cmpIOp, loc, rewriter); + } else if (auto selectOp = dyn_cast(parentOp)) { + bubbleUpOperation(op, selectOp, loc, rewriter); } else if (auto truncFOp = dyn_cast(parentOp)) { bubbleUpOperation(op, truncFOp, loc, rewriter); } else if (auto extFOp = dyn_cast(parentOp)) { @@ -219,6 +221,20 @@ void BubbleUpExtract::bubbleUpOperation( lhs, rhs); } +template +void BubbleUpExtract::bubbleUpOperation( + ExtractOpTy op, arith::SelectOp parentOp, Location loc, + PatternRewriter &rewriter) const { + Value condition = parentOp.getCondition(); + if (isa(condition.getType())) + condition = createExtractOp(op, condition, loc, rewriter); + Value trueValue = createExtractOp(op, parentOp.getTrueValue(), loc, rewriter); + Value falseValue = + createExtractOp(op, parentOp.getFalseValue(), loc, rewriter); + rewriter.replaceOpWithNewOp(op, condition, trueValue, + falseValue); +} + template <> void BubbleUpExtract::bubbleUpOperation( tensor::ExtractOp op, triton::BroadcastOp parentOp, Location loc, diff --git a/third_party/ascend/lib/TritonToUnstructure/OffsetAnalysis.cpp b/third_party/ascend/lib/TritonToUnstructure/OffsetAnalysis.cpp index 898c15070b..84cf9dbccf 100644 --- a/third_party/ascend/lib/TritonToUnstructure/OffsetAnalysis.cpp +++ b/third_party/ascend/lib/TritonToUnstructure/OffsetAnalysis.cpp @@ -207,6 +207,46 @@ PtrOffsetInfo combineInfo(const PtrOffsetInfo &lhs, const PtrOffsetInfo &rhs) { return info; } +namespace { + +bool isScalarPointer(Value value) { + auto pointerType = dyn_cast(value.getType()); + return pointerType && !isa(pointerType.getPointeeType()); +} + +bool isTensorPointer(Value value) { + auto tensorType = dyn_cast(value.getType()); + return tensorType && isa(tensorType.getElementType()); +} + +// A scalar pointer selected or carried by structured control flow is already a +// complete runtime address. Recording the SSA value itself as the source avoids +// choosing one incoming source and losing the other branch or loop iteration. +// Later addptr users can still accumulate a separate offset from this address. +void recordOpaqueScalarPointer( + Value pointer, llvm::DenseMap &offsetMap) { + offsetMap[pointer] = PtrOffsetInfo(); + offsetMap[pointer].setPtr(pointer); + offsetMap[pointer].setZeroOffset(); + offsetMap[pointer].setScalarLike(true); +} + +// A per-lane tensor-pointer select may choose a different source for every +// element, so it cannot be represented by one incoming base plus a structured +// offset. Preserve the selected pointer tensor itself as the complete opaque +// base and attach a zero displacement. Later addptr parsing can accumulate an +// additional offset without discarding the select's lane-wise source choice. +void recordOpaqueTensorPointer( + Value pointerTensor, llvm::DenseMap &offsetMap) { + auto tensorType = cast(pointerTensor.getType()); + offsetMap[pointerTensor] = PtrOffsetInfo(); + offsetMap[pointerTensor].setPtr(pointerTensor); + offsetMap[pointerTensor].setZeroOffset(); + offsetMap[pointerTensor].setUnstructured(tensorType.getRank()); +} + +} // namespace + void parse(Value operand, const Location &loc, RewriterBase &rewriter, llvm::DenseMap &offsetMap) { if (offsetMap.contains(operand)) { @@ -296,6 +336,11 @@ void parseLoopRegionIterArg(LoopLikeOpInterface loopOp, const Location &loc, RewriterBase &rewriter, llvm::DenseMap &offsetMap, BlockArgument regionIterArg) { + if (isScalarPointer(regionIterArg)) { + recordOpaqueScalarPointer(regionIterArg, offsetMap); + return; + } + if (auto whileOp = dyn_cast(loopOp.getOperation()); whileOp && whileOp.getAfterBody() == regionIterArg.getOwner()) { auto argNum = regionIterArg.getArgNumber(); @@ -787,6 +832,16 @@ void parseClampF(triton::ClampFOp op, const Location &loc, void parseSelect(arith::SelectOp op, const Location &loc, RewriterBase &rewriter, llvm::DenseMap &offsetMap) { + if (isScalarPointer(op.getResult())) { + recordOpaqueScalarPointer(op.getResult(), offsetMap); + return; + } + + if (isTensorPointer(op.getResult())) { + recordOpaqueTensorPointer(op.getResult(), offsetMap); + return; + } + // Get select condition auto condition = op.getCondition(); parse(condition, op.getLoc(), rewriter, offsetMap); @@ -967,6 +1022,11 @@ void parseReduceReturn(triton::ReduceReturnOp op, const Location &loc, void parseIf(scf::IfOp op, const Location &loc, RewriterBase &rewriter, llvm::DenseMap &offsetMap, Value dst) { + if (isScalarPointer(dst)) { + recordOpaqueScalarPointer(dst, offsetMap); + return; + } + const unsigned int index = cast(dst).getResultNumber(); // Get if then region Block &thenBlock = op.getThenRegion().front(); @@ -1022,6 +1082,11 @@ void parseYield(scf::YieldOp op, const Location &loc, RewriterBase &rewriter, void parseLoopOp(LoopLikeOpInterface op, const Location &loc, RewriterBase &rewriter, llvm::DenseMap &offsetMap, Value dst) { + if (isScalarPointer(dst)) { + recordOpaqueScalarPointer(dst, offsetMap); + return; + } + auto resNum = cast(dst).getResultNumber(); Value yieldedValue = nullptr; if (auto whileOp = dyn_cast(op.getOperation())) { diff --git a/third_party/ascend/lib/TritonToUnstructure/ReplaceArguments.cpp b/third_party/ascend/lib/TritonToUnstructure/ReplaceArguments.cpp index 3640f61157..0d2310253d 100644 --- a/third_party/ascend/lib/TritonToUnstructure/ReplaceArguments.cpp +++ b/third_party/ascend/lib/TritonToUnstructure/ReplaceArguments.cpp @@ -50,6 +50,11 @@ void replaceOperands(MutableArrayRef oprs, RewriterBase &rewriter, } --it; } else { + // source == operand marks a complete scalar address whose source may be + // selected or loop-carried at runtime. Keep it as a pointer instead of + // replacing it with an offset relative to one statically chosen base. + if (offsetMap.at(operand).getPtr() == operand) + continue; opr.set(offsetMap.at(operand).getOffset()); } } @@ -82,6 +87,11 @@ void replaceArgs(ValueRange args, RewriterBase &rewriter, rewriter.replaceOpWithNewOp( tempVar.getDefiningOp(), tempVar.getType(), src, arg); } else if (auto ptrType = dyn_cast(arg.getType())) { + parse(arg, arg.getLoc(), rewriter, offsetMap); + if (!isa(ptrType.getPointeeType()) && + offsetMap.at(arg).getPtr() == arg) + continue; + RewriterBase::InsertionGuard guard(rewriter); if (auto blockArg = dyn_cast(arg)) { rewriter.setInsertionPointToStart(blockArg.getOwner()); @@ -92,7 +102,6 @@ void replaceArgs(ValueRange args, RewriterBase &rewriter, .create( arg.getLoc(), arg.getType(), ValueRange({})) ->getResult(0); - parse(arg, arg.getLoc(), rewriter, offsetMap); rewriter.replaceAllUsesWith(arg, tempVar); if (auto tensorType = dyn_cast(ptrType.getPointeeType())) { diff --git a/third_party/ascend/lib/TritonToUnstructure/UnstructureConversionPass.cpp b/third_party/ascend/lib/TritonToUnstructure/UnstructureConversionPass.cpp index cfa066238e..1d1f1b936b 100644 --- a/third_party/ascend/lib/TritonToUnstructure/UnstructureConversionPass.cpp +++ b/third_party/ascend/lib/TritonToUnstructure/UnstructureConversionPass.cpp @@ -52,6 +52,11 @@ bool forceSimtTemplateFlag = false; namespace { +static bool isTensorOfPointers(Type type) { + auto tensorType = dyn_cast(type); + return tensorType && isa(tensorType.getElementType()); +} + constexpr int64_t kBitsPerByte = 8; static constexpr const char *kRouteDiscreteMaskToSimtAttrName = @@ -412,6 +417,12 @@ LogicalResult tryRewriteIndirectFastPath(MemAccOpTy op, Location loc, Value srcPtr, Value ptrOffset, ArrayRef resultShape, PatternRewriter &rewriter) { + // Indirect backend operations accept one scalar base plus lane offsets. An + // opaque tensor base can contain a different pointer in every lane, so it + // must be handled by the scalar-loop fallback below. + if (!isa(srcPtr.getType())) + return failure(); + bool rankWithinIndirectLoadStoreFastPathLimit = resultShape.size() <= 5; if (!canUseIndirectFastPath(srcPtr, ptrOffset)) { @@ -767,6 +778,10 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( auto srcPtr = ptrOffsetInfo.getPtr(); auto ptrOffset = ptrOffsetInfo.getOffset(); + if (!isa(srcPtr.getType()) && + !isTensorOfPointers(srcPtr.getType())) + return rewriter.notifyMatchFailure( + op, "expected a scalar pointer or a tensor of scalar pointers"); // LoadLike is operation with result bool isLoadLike = !op->use_empty(); @@ -943,8 +958,20 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( os << extractedOffset << "\n"; }); - assert(isa(srcPtr.getType()) && "src must be ptr type"); - if (!fullyUnstructured) { + // A tensor-of-pointers base is opaque: each lane may select a different + // allocation. Extract the base with the same scalar or slice coordinates as + // the offset, then form the access pointer lane by lane. + if (isTensorOfPointers(srcPtr.getType())) { + if (fullyUnstructured) + srcPtr = createExtractOp(loc, srcPtr, rewriter, offsets); + else + srcPtr = createExtractOp(loc, srcPtr, rewriter, offsets, sizes, strides); + } + + assert((isa(srcPtr.getType()) || + isTensorOfPointers(srcPtr.getType())) && + "src must be a scalar pointer or tensor of pointers"); + if (!fullyUnstructured && isa(srcPtr.getType())) { srcPtr = rewriter.create( loc, RankedTensorType::get(extractedShape, srcPtr.getType()), srcPtr); } diff --git a/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/block_ptr_addptr_base.mlir b/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/block_ptr_addptr_base.mlir new file mode 100644 index 0000000000..d7213d1de7 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/block_ptr_addptr_base.mlir @@ -0,0 +1,95 @@ +// RUN: triton-opt --triton-control-flow-opt --split-input-file %s | FileCheck %s +// RUN: triton-opt --triton-control-flow-opt --triton-to-linalg --split-input-file %s -verify-each | FileCheck %s --check-prefix=LINALG + +module attributes {hacc.target = #hacc.target<"Ascend910B2">} { + tt.func public @local_addptr_base(%base: !tt.ptr) -> tensor<16xf32> { + %c0_i32 = arith.constant 0 : i32 + %c3_i32 = arith.constant 3 : i32 + %c1_i64 = arith.constant 1 : i64 + %c16_i64 = arith.constant 16 : i64 + %shifted = tt.addptr %base, %c3_i32 : !tt.ptr, i32 + %ptr = tt.make_tensor_ptr %shifted, [%c16_i64], [%c1_i64], [%c0_i32] {order = array} : !tt.ptr> + %value = tt.load %ptr : !tt.ptr> + tt.return %value : tensor<16xf32> + } +} + +// CHECK-LABEL: tt.func public @local_addptr_base +// CHECK: %[[SHIFTED:.*]] = tt.addptr +// CHECK: tt.make_tensor_ptr %[[SHIFTED]], +// CHECK-NOT: tt.ptr_to_int + +// LINALG-LABEL: func.func @local_addptr_base +// LINALG: return + +// ----- + +module attributes {hacc.target = #hacc.target<"Ascend910B2">} { + tt.func public @region_local_addptr_base(%base: !tt.ptr, %cond: i1) -> tensor<16xf32> { + %result = scf.if %cond -> (tensor<16xf32>) { + %c0_i32 = arith.constant 0 : i32 + %c3_i32 = arith.constant 3 : i32 + %c1_i64 = arith.constant 1 : i64 + %c16_i64 = arith.constant 16 : i64 + %shifted = tt.addptr %base, %c3_i32 : !tt.ptr, i32 + %ptr = tt.make_tensor_ptr %shifted, [%c16_i64], [%c1_i64], [%c0_i32] {order = array} : !tt.ptr> + %loaded = tt.load %ptr : !tt.ptr> + scf.yield %loaded : tensor<16xf32> + } else { + %zero = arith.constant dense<0.000000e+00> : tensor<16xf32> + scf.yield %zero : tensor<16xf32> + } + tt.return %result : tensor<16xf32> + } +} + +// CHECK-LABEL: tt.func public @region_local_addptr_base +// CHECK: scf.if +// CHECK: %[[SHIFTED:.*]] = tt.addptr +// CHECK: tt.make_tensor_ptr %[[SHIFTED]], +// CHECK-NOT: tt.ptr_to_int + +// LINALG-LABEL: func.func @region_local_addptr_base +// LINALG: return + +// ----- + +module attributes {hacc.target = #hacc.target<"Ascend910B2">} { + tt.func public @loop_carried_addptr_base(%base: !tt.ptr, %output: !tt.ptr) { + %c0_i32 = arith.constant 0 : i32 + %c1_i32 = arith.constant 1 : i32 + %c3_i32 = arith.constant 3 : i32 + %c1_i64 = arith.constant 1 : i64 + %c16_i64 = arith.constant 16 : i64 + %c0 = arith.constant 0 : index + %c2 = arith.constant 2 : index + %c1 = arith.constant 1 : index + %shifted = tt.addptr %base, %c3_i32 : !tt.ptr, i32 + %initial = tt.make_tensor_ptr %shifted, [%c16_i64], [%c1_i64], [%c0_i32] {order = array} : !tt.ptr> + %final = scf.for %iv = %c0 to %c2 step %c1 iter_args(%ptr = %initial) -> (!tt.ptr>) { + %next = tt.advance %ptr, [%c1_i32] : !tt.ptr> + scf.yield %next : !tt.ptr> + } + %value = tt.load %final : !tt.ptr> + %output_ptr = tt.make_tensor_ptr %output, [%c16_i64], [%c1_i64], [%c0_i32] {order = array} : !tt.ptr> + tt.store %output_ptr, %value : !tt.ptr> + tt.return + } +} + +// CHECK-LABEL: tt.func public @loop_carried_addptr_base +// CHECK: %[[SHIFTED:.*]] = tt.addptr +// CHECK: %[[ADDRESS:.*]] = tt.ptr_to_int %[[SHIFTED]] +// CHECK: %[[FOR:.*]]:4 = scf.for +// CHECK-SAME: iter_args(%{{.*}} = %[[ADDRESS]], %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) +// CHECK-SAME: -> (i64, i64, i64, i32) +// CHECK: %[[BASE:.*]] = tt.int_to_ptr %[[FOR]]#0 +// CHECK: tt.make_tensor_ptr %[[BASE]], [%[[FOR]]#1], [%[[FOR]]#2], [%[[FOR]]#3] + +// LINALG-LABEL: func.func @loop_carried_addptr_base +// LINALG: scf.for +// LINALG-NOT: tt.addptr +// LINALG-NOT: tt.make_tensor_ptr +// LINALG-NOT: tt.ptr_to_int +// LINALG-NOT: tt.int_to_ptr +// LINALG: return diff --git a/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/block_ptr_different_base_invalid.mlir b/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/block_ptr_different_base.mlir similarity index 50% rename from third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/block_ptr_different_base_invalid.mlir rename to third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/block_ptr_different_base.mlir index ef8cc4d24b..99322f1be9 100644 --- a/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/block_ptr_different_base_invalid.mlir +++ b/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/block_ptr_different_base.mlir @@ -1,4 +1,4 @@ -// RUN: not triton-opt --triton-control-flow-opt %s 2>&1 | FileCheck %s +// RUN: triton-opt --triton-control-flow-opt %s | FileCheck %s module { tt.func public @if_block_ptr_different_base(%base0: !tt.ptr, %base1: !tt.ptr, %cond: i1) -> !tt.ptr> { @@ -20,4 +20,16 @@ module { } } -// CHECK: error: failed to analyze pointer components across control flow +// CHECK-LABEL: tt.func public @if_block_ptr_different_base( +// CHECK-SAME: %[[BASE0:[^ ,]+]]: !tt.ptr, %[[BASE1:[^ ,]+]]: !tt.ptr +// CHECK: %[[BASE0_ADDR:.*]] = tt.ptr_to_int %[[BASE0]] +// CHECK: %[[BASE1_ADDR:.*]] = tt.ptr_to_int %[[BASE1]] +// CHECK: %[[SELECTED:[^ :]+]]:4 = scf.if %{{[^ ]+}} -> (i64, i64, i64, i32) { +// CHECK: scf.yield %[[BASE0_ADDR]], %{{[^ ,]+}}, %{{[^ ,]+}}, %{{[^ ,]+}} : i64, i64, i64, i32 +// CHECK: } else { +// CHECK: scf.yield %[[BASE1_ADDR]], %{{[^ ,]+}}, %{{[^ ,]+}}, %{{[^ ,]+}} : i64, i64, i64, i32 +// CHECK: } +// CHECK: %[[SELECTED_BASE:.*]] = tt.int_to_ptr %[[SELECTED]]#0 +// CHECK: %[[REBUILT:.*]] = tt.make_tensor_ptr %[[SELECTED_BASE]], +// CHECK-SAME: [%[[SELECTED]]#1], [%[[SELECTED]]#2], [%[[SELECTED]]#3] +// CHECK: tt.return %[[REBUILT]] : !tt.ptr> diff --git a/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/pointer_descriptor_boundary_marker.mlir b/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/pointer_descriptor_boundary_marker.mlir new file mode 100644 index 0000000000..867c8c88d5 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/pointer_descriptor_boundary_marker.mlir @@ -0,0 +1,74 @@ +// RUN: triton-opt --triton-control-flow-opt --split-input-file %s | FileCheck %s + +module { + tt.func public @block_ptr_loop_marker(%base: !tt.ptr, %upper: index) -> !tt.ptr> { + %c0_i32 = arith.constant 0 : i32 + %c1_i32 = arith.constant 1 : i32 + %c1_i64 = arith.constant 1 : i64 + %c16_i64 = arith.constant 16 : i64 + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %initial = tt.make_tensor_ptr %base, [%c16_i64], [%c1_i64], [%c0_i32] {order = array} : !tt.ptr> + %final = scf.for %iv = %c0 to %upper step %c1 iter_args(%ptr = %initial) -> (!tt.ptr>) { + %next = tt.advance %ptr, [%c1_i32] : !tt.ptr> + scf.yield %next : !tt.ptr> + } + tt.return %final : !tt.ptr> + } +} + +// CHECK-LABEL: tt.func public @block_ptr_loop_marker +// CHECK: scf.for +// CHECK: PointerDescriptorBoundary = array + +// ----- + +module { + tt.func public @tensor_ptr_loop_marker(%base: !tt.ptr, %upper: index) -> tensor<4x!tt.ptr> { + %c0_i32 = arith.constant 0 : i32 + %c1_i32 = arith.constant 1 : i32 + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %zero = tt.splat %c0_i32 : i32 -> tensor<4xi32> + %delta = tt.splat %c1_i32 : i32 -> tensor<4xi32> + %base_tensor = tt.splat %base : !tt.ptr -> tensor<4x!tt.ptr> + %initial = tt.addptr %base_tensor, %zero : tensor<4x!tt.ptr>, tensor<4xi32> + %final = scf.for %iv = %c0 to %upper step %c1 iter_args(%ptr = %initial) -> (tensor<4x!tt.ptr>) { + %next = tt.addptr %ptr, %delta : tensor<4x!tt.ptr>, tensor<4xi32> + scf.yield %next : tensor<4x!tt.ptr> + } + tt.return %final : tensor<4x!tt.ptr> + } +} + +// CHECK-LABEL: tt.func public @tensor_ptr_loop_marker +// CHECK: scf.for +// CHECK: PointerDescriptorBoundary = array + +// ----- + +module { + tt.func public @mixed_pointer_loop_marker(%block_base: !tt.ptr, %tensor_base: !tt.ptr, %upper: index) -> tensor<4x!tt.ptr> { + %c0_i32 = arith.constant 0 : i32 + %c1_i32 = arith.constant 1 : i32 + %c1_i64 = arith.constant 1 : i64 + %c16_i64 = arith.constant 16 : i64 + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %zero = tt.splat %c0_i32 : i32 -> tensor<4xi32> + %delta = tt.splat %c1_i32 : i32 -> tensor<4xi32> + %block = tt.make_tensor_ptr %block_base, [%c16_i64], [%c1_i64], [%c0_i32] {order = array} : !tt.ptr> + %tensor_base_splat = tt.splat %tensor_base : !tt.ptr -> tensor<4x!tt.ptr> + %tensor = tt.addptr %tensor_base_splat, %zero : tensor<4x!tt.ptr>, tensor<4xi32> + %results:4 = scf.for %iv = %c0 to %upper step %c1 iter_args(%acc = %c0_i32, %block_arg = %block, %ordinary = %c1_i32, %tensor_arg = %tensor) -> (i32, !tt.ptr>, i32, tensor<4x!tt.ptr>) { + %next_block = tt.advance %block_arg, [%c1_i32] : !tt.ptr> + %next_tensor = tt.addptr %tensor_arg, %delta : tensor<4x!tt.ptr>, tensor<4xi32> + scf.yield %acc, %next_block, %ordinary, %next_tensor : i32, !tt.ptr>, i32, tensor<4x!tt.ptr> + } + tt.return %results#3 : tensor<4x!tt.ptr> + } +} + +// CHECK-LABEL: tt.func public @mixed_pointer_loop_marker +// CHECK: scf.for +// CHECK: PointerDescriptorBoundary = array diff --git a/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/pointer_descriptor_boundary_mixed_downstream.mlir b/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/pointer_descriptor_boundary_mixed_downstream.mlir new file mode 100644 index 0000000000..e67c79e546 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/pointer_descriptor_boundary_mixed_downstream.mlir @@ -0,0 +1,38 @@ +// RUN: triton-opt --triton-control-flow-opt %s | FileCheck %s --check-prefix=CFO +// RUN: triton-opt --triton-control-flow-opt --triton-to-linalg %s -verify-each | FileCheck %s --check-prefix=LINALG --implicit-check-not='!tt.ptr' --implicit-check-not=unrealized_conversion_cast --implicit-check-not=PointerDescriptorBoundary + +module attributes {hacc.target = #hacc.target<"Ascend910B2">} { + tt.func public @expanded_block_ptr_with_invariant_tensor_ptr( + %block_base: !tt.ptr, %tensor_base: !tt.ptr, + %output: !tt.ptr, %upper: index) { + %c0_i32 = arith.constant 0 : i32 + %c1_i32 = arith.constant 1 : i32 + %c1_i64 = arith.constant 1 : i64 + %c4_i64 = arith.constant 4 : i64 + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %range = tt.make_range {end = 4 : i32, start = 0 : i32} : tensor<4xi32> + %block = tt.make_tensor_ptr %block_base, [%c4_i64], [%c1_i64], [%c0_i32] {order = array} : !tt.ptr> + %tensor_base_splat = tt.splat %tensor_base : !tt.ptr -> tensor<4x!tt.ptr> + %tensor = tt.addptr %tensor_base_splat, %range : tensor<4x!tt.ptr>, tensor<4xi32> + %results:2 = scf.for %iv = %c0 to %upper step %c1 iter_args(%block_arg = %block, %tensor_arg = %tensor) -> (!tt.ptr>, tensor<4x!tt.ptr>) { + %next_block = tt.advance %block_arg, [%c1_i32] : !tt.ptr> + scf.yield %next_block, %tensor_arg : !tt.ptr>, tensor<4x!tt.ptr> + } + %value = tt.load %results#1 : tensor<4x!tt.ptr> + %output_splat = tt.splat %output : !tt.ptr -> tensor<4x!tt.ptr> + %output_ptrs = tt.addptr %output_splat, %range : tensor<4x!tt.ptr>, tensor<4xi32> + tt.store %output_ptrs, %value : tensor<4x!tt.ptr> + tt.return + } +} + +// CFO-LABEL: tt.func public @expanded_block_ptr_with_invariant_tensor_ptr +// CFO: %[[LOOP:.*]]:5 = scf.for +// CFO-SAME: -> (i64, i64, i64, i32, tensor<4x!tt.ptr>) { +// CFO: } {PointerDescriptorBoundary = array} +// CFO: tt.load %[[LOOP]]#4 : tensor<4x!tt.ptr> + +// LINALG-LABEL: func.func @expanded_block_ptr_with_invariant_tensor_ptr +// LINALG: scf.for +// LINALG: return diff --git a/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/scf_pointer_decouple.mlir b/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/scf_pointer_decouple.mlir index 22446047fa..4d2c60455c 100644 --- a/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/scf_pointer_decouple.mlir +++ b/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/scf_pointer_decouple.mlir @@ -22,13 +22,13 @@ module { } // CHECK-LABEL: tt.func public @for_block_ptr_dynamic_step -// CHECK: %[[FOR:.*]] = scf.for {{.*}} iter_args(%[[OFF:.*]] = %{{.*}}) -> (i32) { +// CHECK: %[[FOR:.*]]:4 = scf.for {{.*}} iter_args(%[[BASE:.*]] = %{{.*}}, %[[SHAPE:.*]] = %{{.*}}, %[[STRIDE:.*]] = %{{.*}}, %[[OFF:.*]] = %{{.*}}) -> (i64, i64, i64, i32) { // CHECK-NOT: arith.muli -// CHECK: tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%[[OFF]]] {order = array} : > +// CHECK: tt.make_tensor_ptr %{{.*}}, [%[[SHAPE]]], [%[[STRIDE]]], [%[[OFF]]] {order = array} : > // CHECK: %[[NEXT:.*]] = arith.addi %[[OFF]], %{{.*}} : i32 -// CHECK: scf.yield %[[NEXT]] : i32 +// CHECK: scf.yield %[[BASE]], %[[SHAPE]], %[[STRIDE]], %[[NEXT]] : i64, i64, i64, i32 // CHECK: } -// CHECK: tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%[[FOR]]] {order = array} : > +// CHECK: tt.make_tensor_ptr %{{.*}}, [%[[FOR]]#1], [%[[FOR]]#2], [%[[FOR]]#3] {order = array} : > // ----- @@ -52,16 +52,16 @@ module { } // CHECK-LABEL: tt.func public @for_block_ptr_invariant_delta_carried -// CHECK: %[[FOR:[^:]+]]:2 = scf.for -// CHECK-SAME: {{.*}} iter_args(%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) -> (i32, tensor<32xf16>) { +// CHECK: %[[FOR:[^:]+]]:5 = scf.for +// CHECK-SAME: {{.*}} -> (i64, i64, i64, i32, tensor<32xf16>) { // CHECK-NOT: arith.index_cast // CHECK-NOT: arith.muli // CHECK: %[[NEXT:.*]] = arith.addi %{{.*}}, %{{.*}} : i32 // CHECK: %[[LOAD_PTR:.*]] = tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%[[NEXT]]] {order = array} : > // CHECK: %[[LOADED:.*]] = tt.load %[[LOAD_PTR]] : !tt.ptr> -// CHECK: scf.yield %[[NEXT]], %[[LOADED]] : i32, tensor<32xf16> +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %[[NEXT]], %[[LOADED]] : i64, i64, i64, i32, tensor<32xf16> // CHECK: } -// CHECK: tt.return %[[FOR]]#1 : tensor<32xf16> +// CHECK: tt.return %[[FOR]]#4 : tensor<32xf16> // ----- @@ -80,12 +80,12 @@ module { } // CHECK-LABEL: tt.func public @for_block_ptr_dynamic_bounds_and_step -// CHECK: %[[FOR:.*]] = scf.for {{.*}} iter_args(%[[OFF:.*]] = %{{.*}}) -> (i32) { +// CHECK: %[[FOR:.*]]:4 = scf.for {{.*}} iter_args(%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %[[OFF:.*]] = %{{.*}}) -> (i64, i64, i64, i32) { // CHECK-NOT: arith.muli // CHECK: %[[NEXT:.*]] = arith.addi %[[OFF]], %{{.*}} : i32 -// CHECK: scf.yield %[[NEXT]] : i32 +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %[[NEXT]] : i64, i64, i64, i32 // CHECK: } -// CHECK: tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%[[FOR]]] {order = array} : > +// CHECK: tt.make_tensor_ptr %{{.*}}, [%[[FOR]]#1], [%[[FOR]]#2], [%[[FOR]]#3] {order = array} : > // ----- @@ -108,12 +108,12 @@ module { } // CHECK-LABEL: tt.func public @for_block_ptr_iter_arg_delta -// CHECK: %[[FOR:.*]]:2 = scf.for {{.*}} iter_args(%[[OFF:.*]] = %{{.*}}, %[[DELTA:.*]] = %{{.*}}) -> (i32, i32) { +// CHECK: %[[FOR:.*]]:5 = scf.for {{.*}} iter_args(%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %[[OFF:.*]] = %{{.*}}, %[[DELTA:.*]] = %{{.*}}) -> (i64, i64, i64, i32, i32) { // CHECK: %[[NEXT_PTR:.*]] = arith.addi %[[OFF]], %[[DELTA]] : i32 // CHECK: %[[NEXT_DELTA:.*]] = arith.addi %[[DELTA]], %{{.*}} : i32 -// CHECK: scf.yield %[[NEXT_PTR]], %[[NEXT_DELTA]] : i32, i32 +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %[[NEXT_PTR]], %[[NEXT_DELTA]] : i64, i64, i64, i32, i32 // CHECK: } -// CHECK: tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%[[FOR]]#0] {order = array} : > +// CHECK: tt.make_tensor_ptr %{{.*}}, [%[[FOR]]#1], [%[[FOR]]#2], [%[[FOR]]#3] {order = array} : > // ----- @@ -139,12 +139,13 @@ module { } // CHECK-LABEL: tt.func public @for_block_ptr_multidim_delta -// CHECK: %[[FOR:.*]]:2 = scf.for {{.*}} iter_args(%[[OFF0:.*]] = %{{.*}}, %[[OFF1:.*]] = %{{.*}}) -> (i32, i32) { +// CHECK: %[[FOR:.*]]:7 = scf.for {{.*}} iter_args(%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %[[OFF0:.*]] = %{{.*}}, %[[OFF1:.*]] = %{{.*}}) +// CHECK-SAME: -> (i64, i64, i64, i64, i64, i32, i32) { // CHECK: %[[NEXT0:.*]] = arith.addi %[[OFF0]], %{{.*}} : i32 // CHECK: %[[NEXT1:.*]] = arith.addi %[[OFF1]], %{{.*}} : i32 -// CHECK: scf.yield %[[NEXT0]], %[[NEXT1]] : i32, i32 +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[NEXT0]], %[[NEXT1]] : i64, i64, i64, i64, i64, i32, i32 // CHECK: } -// CHECK: tt.make_tensor_ptr %{{.*}}, [%{{.*}}, %{{.*}}], [%{{.*}}, %{{.*}}], [%[[FOR]]#0, %[[FOR]]#1] {order = array} : > +// CHECK: tt.make_tensor_ptr %{{.*}}, [%[[FOR]]#1, %[[FOR]]#2], [%[[FOR]]#3, %[[FOR]]#4], [%[[FOR]]#5, %[[FOR]]#6] {order = array} : > // ----- @@ -167,11 +168,11 @@ module { } // CHECK-LABEL: tt.func public @for_block_ptr_zero_trip -// CHECK: %[[FOR:.*]] = scf.for {{.*}} iter_args(%[[OFF:.*]] = %{{.*}}) -> (i32) { +// CHECK: %[[FOR:.*]]:4 = scf.for {{.*}} iter_args(%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %[[OFF:.*]] = %{{.*}}) -> (i64, i64, i64, i32) { // CHECK: %[[NEXT:.*]] = arith.addi %[[OFF]], %{{.*}} : i32 -// CHECK: scf.yield %[[NEXT]] : i32 +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %[[NEXT]] : i64, i64, i64, i32 // CHECK: } -// CHECK: tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%[[FOR]]] {order = array} : > +// CHECK: tt.make_tensor_ptr %{{.*}}, [%[[FOR]]#1], [%[[FOR]]#2], [%[[FOR]]#3] {order = array} : > // ----- @@ -196,13 +197,14 @@ module { } // CHECK-LABEL: tt.func public @while_block_ptr_basic -// CHECK: %[[WHILE:.*]]:2 = scf.while (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) : (i32, i32) -> (i32, i32) { -// CHECK: scf.condition(%{{.*}}) %{{.*}}, %{{.*}} : i32, i32 +// CHECK: %[[WHILE:.*]]:5 = scf.while +// CHECK-SAME: : (i32, i64, i64, i64, i32) -> (i32, i64, i64, i64, i32) { +// CHECK: scf.condition(%{{.*}}) %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i32, i64, i64, i64, i32 // CHECK: } do { // CHECK: tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%{{.*}}] {order = array} : > -// CHECK: scf.yield %{{.*}}, %{{.*}} : i32, i32 +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i32, i64, i64, i64, i32 // CHECK: } -// CHECK: tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%[[WHILE]]#1] {order = array} : > +// CHECK: tt.make_tensor_ptr %{{.*}}, [%[[WHILE]]#2], [%[[WHILE]]#3], [%[[WHILE]]#4] {order = array} : > // ----- @@ -226,12 +228,12 @@ module { } // CHECK-LABEL: tt.func public @if_block_ptr_same_base_offsets -// CHECK: %[[OFF:.*]] = scf.if %{{.*}} -> (i32) { -// CHECK: scf.yield %{{.*}} : i32 +// CHECK: %[[DESC:.*]]:4 = scf.if %{{.*}} -> (i64, i64, i64, i32) { +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i64, i64, i64, i32 // CHECK: } else { -// CHECK: scf.yield %{{.*}} : i32 +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i64, i64, i64, i32 // CHECK: } -// CHECK: tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%[[OFF]]] {order = array} : > +// CHECK: tt.make_tensor_ptr %{{.*}}, [%[[DESC]]#1], [%[[DESC]]#2], [%[[DESC]]#3] {order = array} : > // ----- @@ -288,9 +290,11 @@ module { // CHECK-LABEL: tt.func public @while_block_ptr_large_step // CHECK-DAG: %[[C37:.*]] = arith.constant 37 : i32 +// CHECK: %[[WHILE:.*]]:5 = scf.while +// CHECK-SAME: : (i32, i64, i64, i64, i32) -> (i32, i64, i64, i64, i32) // CHECK: } do { // CHECK: %[[NEXT_OFF:.*]] = arith.addi %{{.*}}, %[[C37]] : i32 -// CHECK: scf.yield %{{.*}}, %[[NEXT_OFF]] : i32, i32 +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[NEXT_OFF]] : i32, i64, i64, i64, i32 // ----- @@ -382,18 +386,18 @@ module { } // CHECK-LABEL: tt.func public @for_if_block_ptr_same_base_post_advance -// CHECK: %[[FOR:.*]] = scf.for {{.*}} iter_args(%[[OFF:.*]] = %{{.*}}) -> (i32) { -// CHECK: %[[SELECTED:.*]] = scf.if %{{.*}} -> (i32) { +// CHECK: %[[FOR:.*]]:4 = scf.for {{.*}} iter_args(%[[BASE:.*]] = %{{.*}}, %[[SHAPE:.*]] = %{{.*}}, %[[STRIDE:.*]] = %{{.*}}, %[[OFF:.*]] = %{{.*}}) -> (i64, i64, i64, i32) { +// CHECK: %[[SELECTED:.*]]:4 = scf.if %{{.*}} -> (i64, i64, i64, i32) { // CHECK: arith.addi %[[OFF]], %{{.*}} : i32 -// CHECK: scf.yield %{{.*}} : i32 +// CHECK: scf.yield %[[BASE]], %[[SHAPE]], %[[STRIDE]], %{{.*}} : i64, i64, i64, i32 // CHECK: } else { // CHECK: arith.addi %[[OFF]], %{{.*}} : i32 -// CHECK: scf.yield %{{.*}} : i32 +// CHECK: scf.yield %[[BASE]], %[[SHAPE]], %[[STRIDE]], %{{.*}} : i64, i64, i64, i32 // CHECK: } -// CHECK: %[[NEXT:.*]] = arith.addi %[[SELECTED]], %{{.*}} : i32 -// CHECK: scf.yield %[[NEXT]] : i32 +// CHECK: %[[NEXT:.*]] = arith.addi %[[SELECTED]]#3, %{{.*}} : i32 +// CHECK: scf.yield %[[SELECTED]]#0, %[[SELECTED]]#1, %[[SELECTED]]#2, %[[NEXT]] : i64, i64, i64, i32 // CHECK: } -// CHECK: tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%[[FOR]]] {order = array} : > +// CHECK: tt.make_tensor_ptr %{{.*}}, [%[[FOR]]#1], [%[[FOR]]#2], [%[[FOR]]#3] {order = array} : > // ----- @@ -428,18 +432,18 @@ module { } // CHECK-LABEL: tt.func public @while_if_block_ptr_same_base_post_advance -// CHECK: %[[WHILE:.*]]:2 = scf.while -// CHECK-SAME: (i32, i32) -> (i32, i32) +// CHECK: %[[WHILE:.*]]:5 = scf.while +// CHECK-SAME: (i32, i64, i64, i64, i32) -> (i32, i64, i64, i64, i32) // CHECK: } do { -// CHECK: %[[SELECTED:.*]] = scf.if %{{.*}} -> (i32) { -// CHECK: scf.yield %{{.*}} : i32 +// CHECK: %[[SELECTED:.*]]:4 = scf.if %{{.*}} -> (i64, i64, i64, i32) { +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i64, i64, i64, i32 // CHECK: } else { -// CHECK: scf.yield %{{.*}} : i32 +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i64, i64, i64, i32 // CHECK: } -// CHECK: %[[NEXT:.*]] = arith.addi %[[SELECTED]], %{{.*}} : i32 -// CHECK: scf.yield %{{.*}}, %[[NEXT]] : i32, i32 +// CHECK: %[[NEXT:.*]] = arith.addi %[[SELECTED]]#3, %{{.*}} : i32 +// CHECK: scf.yield %{{.*}}, %[[SELECTED]]#0, %[[SELECTED]]#1, %[[SELECTED]]#2, %[[NEXT]] : i32, i64, i64, i64, i32 // CHECK: } -// CHECK: tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%[[WHILE]]#1] {order = array} : > +// CHECK: tt.make_tensor_ptr %{{.*}}, [%[[WHILE]]#2], [%[[WHILE]]#3], [%[[WHILE]]#4] {order = array} : > // ----- @@ -518,20 +522,21 @@ module { } // CHECK-LABEL: tt.func public @for_if_block_ptr_load_after_post_advance -// CHECK: %[[FOR:.*]]:2 = scf.for {{.*}} iter_args(%[[OFF:.*]] = %{{.*}}, %{{.*}} = %{{.*}}) -> (i32, tensor<32xf16>) { -// CHECK: %[[SELECTED:.*]] = scf.if %{{.*}} -> (i32) { +// CHECK: %[[FOR:.*]]:5 = scf.for {{.*}} iter_args(%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %[[OFF:.*]] = %{{.*}}, %{{.*}} = %{{.*}}) +// CHECK-SAME: -> (i64, i64, i64, i32, tensor<32xf16>) { +// CHECK: %[[SELECTED:.*]]:4 = scf.if %{{.*}} -> (i64, i64, i64, i32) { // CHECK: arith.addi %[[OFF]], %{{.*}} : i32 -// CHECK: scf.yield %{{.*}} : i32 +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i64, i64, i64, i32 // CHECK: } else { // CHECK: arith.addi %[[OFF]], %{{.*}} : i32 -// CHECK: scf.yield %{{.*}} : i32 +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i64, i64, i64, i32 // CHECK: } -// CHECK: %[[NEXT:.*]] = arith.addi %[[SELECTED]], %{{.*}} : i32 +// CHECK: %[[NEXT:.*]] = arith.addi %[[SELECTED]]#3, %{{.*}} : i32 // CHECK: %[[LOAD_PTR:.*]] = tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%[[NEXT]]] {order = array} : > // CHECK: %[[LOADED:.*]] = tt.load %[[LOAD_PTR]] : !tt.ptr> -// CHECK: scf.yield %[[NEXT]], %[[LOADED]] : i32, tensor<32xf16> +// CHECK: scf.yield %[[SELECTED]]#0, %[[SELECTED]]#1, %[[SELECTED]]#2, %[[NEXT]], %[[LOADED]] : i64, i64, i64, i32, tensor<32xf16> // CHECK: } -// CHECK: tt.return %[[FOR]]#1 : tensor<32xf16> +// CHECK: tt.return %[[FOR]]#4 : tensor<32xf16> // ----- @@ -581,28 +586,29 @@ module { } // CHECK-LABEL: tt.func public @for_nested_if_block_ptr_load_after_post_advance -// CHECK: %[[FOR:.*]]:2 = scf.for {{.*}} iter_args(%[[OFF:.*]] = %{{.*}}, %{{.*}} = %{{.*}}) -> (i32, tensor<16xf32>) { -// CHECK: %[[SELECTED:.*]] = scf.if %{{.*}} -> (i32) { -// CHECK: %[[THEN_INNER:.*]] = scf.if %{{.*}} -> (i32) { -// CHECK: scf.yield %{{.*}} : i32 +// CHECK: %[[FOR:.*]]:5 = scf.for {{.*}} iter_args(%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %[[OFF:.*]] = %{{.*}}, %{{.*}} = %{{.*}}) +// CHECK-SAME: -> (i64, i64, i64, i32, tensor<16xf32>) { +// CHECK: %[[SELECTED:.*]]:4 = scf.if %{{.*}} -> (i64, i64, i64, i32) { +// CHECK: %[[THEN_INNER:.*]]:4 = scf.if %{{.*}} -> (i64, i64, i64, i32) { +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i64, i64, i64, i32 // CHECK: } else { -// CHECK: scf.yield %{{.*}} : i32 +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i64, i64, i64, i32 // CHECK: } -// CHECK: scf.yield %[[THEN_INNER]] : i32 +// CHECK: scf.yield %[[THEN_INNER]]#0, %[[THEN_INNER]]#1, %[[THEN_INNER]]#2, %[[THEN_INNER]]#3 : i64, i64, i64, i32 // CHECK: } else { -// CHECK: %[[ELSE_INNER:.*]] = scf.if %{{.*}} -> (i32) { -// CHECK: scf.yield %{{.*}} : i32 +// CHECK: %[[ELSE_INNER:.*]]:4 = scf.if %{{.*}} -> (i64, i64, i64, i32) { +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i64, i64, i64, i32 // CHECK: } else { -// CHECK: scf.yield %{{.*}} : i32 +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i64, i64, i64, i32 // CHECK: } -// CHECK: scf.yield %[[ELSE_INNER]] : i32 +// CHECK: scf.yield %[[ELSE_INNER]]#0, %[[ELSE_INNER]]#1, %[[ELSE_INNER]]#2, %[[ELSE_INNER]]#3 : i64, i64, i64, i32 // CHECK: } -// CHECK: %[[NEXT:.*]] = arith.addi %[[SELECTED]], %{{.*}} : i32 +// CHECK: %[[NEXT:.*]] = arith.addi %[[SELECTED]]#3, %{{.*}} : i32 // CHECK: %[[LOAD_PTR:.*]] = tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%[[NEXT]]] {order = array} : > // CHECK: %[[LOADED:.*]] = tt.load %[[LOAD_PTR]] : !tt.ptr> -// CHECK: scf.yield %[[NEXT]], %[[LOADED]] : i32, tensor<16xf32> +// CHECK: scf.yield %[[SELECTED]]#0, %[[SELECTED]]#1, %[[SELECTED]]#2, %[[NEXT]], %[[LOADED]] : i64, i64, i64, i32, tensor<16xf32> // CHECK: } -// CHECK: tt.return %[[FOR]]#1 : tensor<16xf32> +// CHECK: tt.return %[[FOR]]#4 : tensor<16xf32> // ----- @@ -698,15 +704,15 @@ module { } // CHECK-LABEL: tt.func public @if_block_ptr_same_base_dynamic_shape -// CHECK: %[[SHAPE:.*]] = scf.if %{{.*}} -> (i64) { -// CHECK: scf.yield %{{.*}} : i64 +// CHECK: %[[DESC:.*]]:4 = scf.if %{{.*}} -> (i64, i64, i64, i32) { +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i64, i64, i64, i32 // CHECK: } else { -// CHECK: scf.yield %{{.*}} : i64 +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i64, i64, i64, i32 // CHECK: } -// CHECK: tt.make_tensor_ptr %{{.*}}, [%[[SHAPE]]], [%{{.*}}], [%{{.*}}] {order = array} : > +// CHECK: tt.make_tensor_ptr %{{.*}}, [%[[DESC]]#1], [%[[DESC]]#2], [%[[DESC]]#3] {order = array} : > module { - tt.func public @if_same_nested_result_not_decoupled(%base: !tt.ptr, %cond0: i1, %cond1: i1) -> !tt.ptr> { + tt.func public @if_same_nested_result_fully_decoupled(%base: !tt.ptr, %cond0: i1, %cond1: i1) -> !tt.ptr> { %c0_i32 = arith.constant 0 : i32 %c1_i32 = arith.constant 1 : i32 %c2_i32 = arith.constant 2 : i32 @@ -729,19 +735,20 @@ module { } } -// CHECK-LABEL: tt.func public @if_same_nested_result_not_decoupled -// CHECK: %[[INNER_OFF:.*]] = scf.if %{{.*}} -> (i32) { -// CHECK: scf.yield %{{.*}} : i32 +// CHECK-LABEL: tt.func public @if_same_nested_result_fully_decoupled +// CHECK: %[[INNER_DESC:.*]]:4 = scf.if %{{.*}} -> (i64, i64, i64, i32) { +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i64, i64, i64, i32 // CHECK: } else { -// CHECK: scf.yield %{{.*}} : i32 +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i64, i64, i64, i32 // CHECK: } -// CHECK: %[[INNER_PTR:.*]] = tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%[[INNER_OFF]]] {order = array} : > -// CHECK: %[[OUTER:.*]] = scf.if %{{.*}} -> (!tt.ptr>) { -// CHECK: scf.yield %[[INNER_PTR]] : !tt.ptr> +// CHECK: %[[OUTER:.*]]:4 = scf.if %{{.*}} -> (i64, i64, i64, i32) { +// CHECK: scf.yield %[[INNER_DESC]]#0, %[[INNER_DESC]]#1, %[[INNER_DESC]]#2, %[[INNER_DESC]]#3 : i64, i64, i64, i32 // CHECK: } else { -// CHECK: scf.yield %[[INNER_PTR]] : !tt.ptr> +// CHECK: scf.yield %[[INNER_DESC]]#0, %[[INNER_DESC]]#1, %[[INNER_DESC]]#2, %[[INNER_DESC]]#3 : i64, i64, i64, i32 // CHECK: } -// CHECK: tt.return %[[OUTER]] : !tt.ptr> +// CHECK: %[[OUTER_BASE:.*]] = tt.int_to_ptr %[[OUTER]]#0 +// CHECK: %[[OUTER_PTR:.*]] = tt.make_tensor_ptr %[[OUTER_BASE]], [%[[OUTER]]#1], [%[[OUTER]]#2], [%[[OUTER]]#3] +// CHECK: tt.return %[[OUTER_PTR]] : !tt.ptr> // ----- @@ -770,13 +777,13 @@ module { } // CHECK-LABEL: tt.func public @if_mixed_block_and_tensor_ptr -// CHECK: %[[RESULT:.*]]:2 = scf.if %{{.*}} -> (i32, tensor<4xi32>) { -// CHECK: scf.yield %{{.*}}, %{{.*}} : i32, tensor<4xi32> +// CHECK: %[[RESULT:.*]]:5 = scf.if %{{.*}} -> (i64, i64, i64, i32, tensor<4xi32>) { +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i64, i64, i64, i32, tensor<4xi32> // CHECK: } else { -// CHECK: scf.yield %{{.*}}, %{{.*}} : i32, tensor<4xi32> +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i64, i64, i64, i32, tensor<4xi32> // CHECK: } -// CHECK-DAG: %[[BLOCK:.*]] = tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%[[RESULT]]#0] {order = array} : > -// CHECK-DAG: %[[TENSOR:.*]] = tt.addptr %{{.*}}, %[[RESULT]]#1 : tensor<4x!tt.ptr>, tensor<4xi32> +// CHECK-DAG: %[[BLOCK:.*]] = tt.make_tensor_ptr %{{.*}}, [%[[RESULT]]#1], [%[[RESULT]]#2], [%[[RESULT]]#3] {order = array} : > +// CHECK-DAG: %[[TENSOR:.*]] = tt.addptr %{{.*}}, %[[RESULT]]#4 : tensor<4x!tt.ptr>, tensor<4xi32> // CHECK: tt.return %[[BLOCK]], %[[TENSOR]] : !tt.ptr>, tensor<4x!tt.ptr> // ----- @@ -805,11 +812,11 @@ module { // CHECK-LABEL: tt.func public @if_without_else_rewrites_nested_block_ptr // CHECK: scf.if %{{.*}} { -// CHECK: %[[FOR:.*]] = scf.for -// CHECK-SAME: -> (i32) { -// CHECK: scf.yield %{{.*}} : i32 +// CHECK: %[[FOR:.*]]:4 = scf.for +// CHECK-SAME: -> (i64, i64, i64, i32) { +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i64, i64, i64, i32 // CHECK: } -// CHECK: %[[PTR:.*]] = tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%[[FOR]]] {order = array} : > +// CHECK: %[[PTR:.*]] = tt.make_tensor_ptr %{{.*}}, [%[[FOR]]#1], [%[[FOR]]#2], [%[[FOR]]#3] {order = array} : > // CHECK: tt.load %[[PTR]] : !tt.ptr> // CHECK: } diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/scalar_pointer_select.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/scalar_pointer_select.mlir new file mode 100644 index 0000000000..b34fc8cb60 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/scalar_pointer_select.mlir @@ -0,0 +1,207 @@ +// RUN: triton-opt --triton-to-linalg %s -verify-each | FileCheck %s + +module attributes {hacc.target = #hacc.target<"Ascend910B2">} { + func.func private @consume_offset_view(memref<1xf32, strided<[1], offset: 1>>) + func.func private @next_offset_view(memref<1xf32, strided<[1], offset: 1>>) -> memref<1xf32, strided<[1], offset: 1>> + + tt.func public @scalar_pointer_select(%lhs: !tt.ptr, %rhs: !tt.ptr, %condition: i1) -> i64 { + %selected = arith.select %condition, %lhs, %rhs : !tt.ptr + %address = tt.ptr_to_int %selected : !tt.ptr -> i64 + tt.return %address : i64 + } + + tt.func public @scalar_pointer_roundtrip(%address: i64) -> i64 { + %pointer = tt.int_to_ptr %address : i64 -> !tt.ptr + %roundtrip = tt.ptr_to_int %pointer : !tt.ptr -> i64 + tt.return %roundtrip : i64 + } + + func.func @pointer_cast_static_offset(%address: i64) -> f32 { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %pointer = hivm.hir.pointer_cast(%address) [%c1] : memref + %view = memref.reinterpret_cast %pointer to offset: [1], sizes: [1], strides: [1] + : memref to memref<1xf32, strided<[1], offset: 1>> + %value = memref.load %view[%c0] : memref<1xf32, strided<[1], offset: 1>> + return %value : f32 + } + + func.func @pointer_cast_dynamic_offset_subviews(%address: i64, %offset: index, %size: index) -> f32 { + %c0 = arith.constant 0 : index + %c8 = arith.constant 8 : index + %pointer = hivm.hir.pointer_cast(%address) [%c8] : memref + %view = memref.reinterpret_cast %pointer to offset: [%offset], sizes: [1, 8], strides: [8, 1] + : memref to memref<1x8xf32, strided<[8, 1], offset: ?>> + %rank_reduced = memref.subview %view[0, 0] [1, %size] [1, 1] + : memref<1x8xf32, strided<[8, 1], offset: ?>> to memref> + %subview = memref.subview %rank_reduced[0] [%size] [1] + : memref> to memref> + %value = memref.load %subview[%c0] : memref> + return %value : f32 + } + + func.func @pointer_cast_offset_return(%address: i64) -> memref<1xf32, strided<[1], offset: 1>> { + %c1 = arith.constant 1 : index + %pointer = hivm.hir.pointer_cast(%address) [%c1] : memref + %view = memref.reinterpret_cast %pointer to offset: [1], sizes: [1], strides: [1] + : memref to memref<1xf32, strided<[1], offset: 1>> + return %view : memref<1xf32, strided<[1], offset: 1>> + } + + func.func @pointer_cast_dynamic_offset_return(%address: i64, %offset: index) -> memref<1xf32, strided<[1], offset: ?>> { + %c1 = arith.constant 1 : index + %pointer = hivm.hir.pointer_cast(%address) [%c1] : memref + %view = memref.reinterpret_cast %pointer to offset: [%offset], sizes: [1], strides: [1] + : memref to memref<1xf32, strided<[1], offset: ?>> + return %view : memref<1xf32, strided<[1], offset: ?>> + } + + func.func @pointer_cast_offset_call(%address: i64) { + %c1 = arith.constant 1 : index + %pointer = hivm.hir.pointer_cast(%address) [%c1] : memref + %view = memref.reinterpret_cast %pointer to offset: [1], sizes: [1], strides: [1] + : memref to memref<1xf32, strided<[1], offset: 1>> + call @consume_offset_view(%view) : (memref<1xf32, strided<[1], offset: 1>>) -> () + return + } + + func.func @pointer_cast_offset_if(%address: i64, %condition: i1, %alternate: memref<1xf32, strided<[1], offset: 1>>) -> memref<1xf32, strided<[1], offset: 1>> { + %c1 = arith.constant 1 : index + %pointer = hivm.hir.pointer_cast(%address) [%c1] : memref + %view = memref.reinterpret_cast %pointer to offset: [1], sizes: [1], strides: [1] + : memref to memref<1xf32, strided<[1], offset: 1>> + %result = scf.if %condition -> (memref<1xf32, strided<[1], offset: 1>>) { + %then = call @next_offset_view(%view) : (memref<1xf32, strided<[1], offset: 1>>) -> memref<1xf32, strided<[1], offset: 1>> + scf.yield %then : memref<1xf32, strided<[1], offset: 1>> + } else { + %else = call @next_offset_view(%alternate) : (memref<1xf32, strided<[1], offset: 1>>) -> memref<1xf32, strided<[1], offset: 1>> + scf.yield %else : memref<1xf32, strided<[1], offset: 1>> + } + return %result : memref<1xf32, strided<[1], offset: 1>> + } + + func.func @pointer_cast_offset_loop(%address: i64, %upper: index) -> memref<1xf32, strided<[1], offset: 1>> { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %pointer = hivm.hir.pointer_cast(%address) [%c1] : memref + %view = memref.reinterpret_cast %pointer to offset: [1], sizes: [1], strides: [1] + : memref to memref<1xf32, strided<[1], offset: 1>> + %result = scf.for %iv = %c0 to %upper step %c1 iter_args(%current = %view) -> (memref<1xf32, strided<[1], offset: 1>>) { + %next = call @next_offset_view(%current) : (memref<1xf32, strided<[1], offset: 1>>) -> memref<1xf32, strided<[1], offset: 1>> + scf.yield %next : memref<1xf32, strided<[1], offset: 1>> + } + return %result : memref<1xf32, strided<[1], offset: 1>> + } + + func.func @pointer_cast_offset_while(%address: i64, %condition: i1) -> memref<1xf32, strided<[1], offset: 1>> { + %c1 = arith.constant 1 : index + %pointer = hivm.hir.pointer_cast(%address) [%c1] : memref + %view = memref.reinterpret_cast %pointer to offset: [1], sizes: [1], strides: [1] + : memref to memref<1xf32, strided<[1], offset: 1>> + %result = scf.while (%current = %view) : (memref<1xf32, strided<[1], offset: 1>>) -> (memref<1xf32, strided<[1], offset: 1>>) { + scf.condition(%condition) %current : memref<1xf32, strided<[1], offset: 1>> + } do { + ^bb0(%current: memref<1xf32, strided<[1], offset: 1>>): + %next = call @next_offset_view(%current) : (memref<1xf32, strided<[1], offset: 1>>) -> memref<1xf32, strided<[1], offset: 1>> + scf.yield %next : memref<1xf32, strided<[1], offset: 1>> + } + return %result : memref<1xf32, strided<[1], offset: 1>> + } +} + +// CHECK-LABEL: func.func @scalar_pointer_select( +// CHECK-SAME: %[[LHS:[^ ,]+]]: memref, %[[RHS:[^ ,]+]]: memref +// CHECK: %[[LHS_INDEX:.*]] = memref.extract_aligned_pointer_as_index %[[LHS]] +// CHECK: %[[LHS_ADDRESS:.*]] = arith.index_cast %[[LHS_INDEX]] : index to i64 +// CHECK: %[[RHS_INDEX:.*]] = memref.extract_aligned_pointer_as_index %[[RHS]] +// CHECK: %[[RHS_ADDRESS:.*]] = arith.index_cast %[[RHS_INDEX]] : index to i64 +// CHECK: %[[SELECTED:.*]] = arith.select %{{.*}}, %[[LHS_ADDRESS]], %[[RHS_ADDRESS]] : i64 +// CHECK-NOT: arith.select {{.*}} : memref +// CHECK-NOT: hivm.hir.pointer_cast +// CHECK-NOT: memref.extract_aligned_pointer_as_index +// CHECK: return %[[SELECTED]] : i64 + +// CHECK-LABEL: func.func @scalar_pointer_roundtrip( +// CHECK-SAME: %[[ADDRESS:[^ ,]+]]: i64 +// CHECK-NOT: hivm.hir.pointer_cast +// CHECK-NOT: memref.extract_aligned_pointer_as_index +// CHECK: return %[[ADDRESS]] : i64 + +// CHECK-LABEL: func.func @pointer_cast_static_offset( +// CHECK-SAME: %[[ADDRESS:[^ ,]+]]: i64 +// CHECK: %[[OFFSET_I64:.*]] = arith.index_cast %[[OFFSET:.*]] : index to i64 +// CHECK: %[[BYTE_WIDTH:.*]] = arith.constant 4 : i64 +// CHECK: %[[BYTE_OFFSET:.*]] = arith.muli %[[OFFSET_I64]], %[[BYTE_WIDTH]] : i64 +// CHECK: %[[REAL_ADDRESS:.*]] = arith.addi %[[ADDRESS]], %[[BYTE_OFFSET]] : i64 +// CHECK: %[[REBASED_POINTER:.*]] = hivm.hir.pointer_cast(%[[REAL_ADDRESS]]) [%[[SIZE:.*]]] : memref +// CHECK: %[[VIEW:.*]] = memref.reinterpret_cast %[[REBASED_POINTER]] to offset: [0], sizes: [1], strides: [1] +// CHECK-SAME: to memref<1xf32, strided<[1]>> +// CHECK: memref.load %[[VIEW]][%{{.*}}] : memref<1xf32, strided<[1]>> + +// CHECK-LABEL: func.func @pointer_cast_dynamic_offset_subviews( +// CHECK-SAME: %[[ADDRESS:[^ ,]+]]: i64, %[[OFFSET:[^ ,]+]]: index, %[[DYNAMIC_SIZE:[^ ,]+]]: index +// CHECK: %[[OFFSET_I64:.*]] = arith.index_cast %[[OFFSET]] : index to i64 +// CHECK: %[[BYTE_WIDTH:.*]] = arith.constant 4 : i64 +// CHECK: %[[BYTE_OFFSET:.*]] = arith.muli %[[OFFSET_I64]], %[[BYTE_WIDTH]] : i64 +// CHECK: %[[REAL_ADDRESS:.*]] = arith.addi %[[ADDRESS]], %[[BYTE_OFFSET]] : i64 +// CHECK: %[[REBASED_POINTER:.*]] = hivm.hir.pointer_cast(%[[REAL_ADDRESS]]) [%[[SIZE:.*]]] : memref +// CHECK: %[[VIEW:.*]] = memref.reinterpret_cast %[[REBASED_POINTER]] to offset: [0], sizes: [1, 8], strides: [8, 1] +// CHECK-SAME: to memref<1x8xf32, strided<[8, 1]>> +// CHECK: %[[RANK_REDUCED:.*]] = memref.subview %[[VIEW]][0, 0] [1, %[[DYNAMIC_SIZE]]] [1, 1] +// CHECK-SAME: to memref> +// CHECK: %[[SUBVIEW:.*]] = memref.subview %[[RANK_REDUCED]][0] [%[[DYNAMIC_SIZE]]] [1] +// CHECK-SAME: to memref> +// CHECK: memref.load %[[SUBVIEW]][%{{.*}}] : memref> + +// CHECK-LABEL: func.func @pointer_cast_offset_return( +// CHECK-SAME: %[[ADDRESS:[^ ,]+]]: i64 +// CHECK: %[[CAPACITY:.*]] = arith.addi %[[VIEW_SIZE:.*]], %[[VIEW_OFFSET:.*]] : index +// CHECK: %[[POINTER:.*]] = hivm.hir.pointer_cast(%[[ADDRESS]]) [%[[CAPACITY]]] : memref +// CHECK: %[[VIEW:.*]] = memref.reinterpret_cast %[[POINTER]] to offset: [1], sizes: [1], strides: [1] +// CHECK-SAME: to memref<1xf32, strided<[1], offset: 1>> +// CHECK: return %[[VIEW]] : memref<1xf32, strided<[1], offset: 1>> + +// CHECK-LABEL: func.func @pointer_cast_dynamic_offset_return( +// CHECK-SAME: %[[ADDRESS:[^ ,]+]]: i64, %[[OFFSET:[^ ,]+]]: index +// CHECK: %[[ZERO:.*]] = arith.constant 0 : index +// CHECK: %[[LEADING_EXTENT:.*]] = arith.maxsi %[[OFFSET]], %[[ZERO]] : index +// CHECK: %[[CAPACITY:.*]] = arith.addi %[[VIEW_SIZE:.*]], %[[LEADING_EXTENT]] : index +// CHECK: %[[POINTER:.*]] = hivm.hir.pointer_cast(%[[ADDRESS]]) [%[[CAPACITY]]] : memref +// CHECK: %[[VIEW:.*]] = memref.reinterpret_cast %[[POINTER]] to offset: [%[[OFFSET]]], sizes: [1], strides: [1] +// CHECK-SAME: to memref<1xf32, strided<[1], offset: ?>> +// CHECK: return %[[VIEW]] : memref<1xf32, strided<[1], offset: ?>> + +// CHECK-LABEL: func.func @pointer_cast_offset_call( +// CHECK-SAME: %[[ADDRESS:[^ ,]+]]: i64 +// CHECK: %[[POINTER:.*]] = hivm.hir.pointer_cast(%[[ADDRESS]]) +// CHECK: %[[VIEW:.*]] = memref.reinterpret_cast %[[POINTER]] to offset: [1], sizes: [1], strides: [1] +// CHECK-SAME: to memref<1xf32, strided<[1], offset: 1>> +// CHECK: call @consume_offset_view(%[[VIEW]]) + +// CHECK-LABEL: func.func @pointer_cast_offset_if( +// CHECK-SAME: %[[ADDRESS:[^ ,]+]]: i64 +// CHECK: %[[POINTER:.*]] = hivm.hir.pointer_cast(%[[ADDRESS]]) +// CHECK: %[[VIEW:.*]] = memref.reinterpret_cast %[[POINTER]] to offset: [1], sizes: [1], strides: [1] +// CHECK-SAME: to memref<1xf32, strided<[1], offset: 1>> +// CHECK: scf.if +// CHECK-SAME: -> (memref<1xf32, strided<[1], offset: 1>>) +// CHECK: call @next_offset_view(%[[VIEW]]) +// CHECK: scf.yield {{.*}} : memref<1xf32, strided<[1], offset: 1>> + +// CHECK-LABEL: func.func @pointer_cast_offset_loop( +// CHECK-SAME: %[[ADDRESS:[^ ,]+]]: i64 +// CHECK: %[[POINTER:.*]] = hivm.hir.pointer_cast(%[[ADDRESS]]) +// CHECK: %[[VIEW:.*]] = memref.reinterpret_cast %[[POINTER]] to offset: [1], sizes: [1], strides: [1] +// CHECK-SAME: to memref<1xf32, strided<[1], offset: 1>> +// CHECK: scf.for +// CHECK-SAME: memref<1xf32, strided<[1], offset: 1>> + +// CHECK-LABEL: func.func @pointer_cast_offset_while( +// CHECK-SAME: %[[ADDRESS:[^ ,]+]]: i64 +// CHECK: %[[POINTER:.*]] = hivm.hir.pointer_cast(%[[ADDRESS]]) +// CHECK: %[[VIEW:.*]] = memref.reinterpret_cast %[[POINTER]] to offset: [1], sizes: [1], strides: [1] +// CHECK-SAME: to memref<1xf32, strided<[1], offset: 1>> +// CHECK: scf.while +// CHECK-SAME: memref<1xf32, strided<[1], offset: 1>> +// CHECK: scf.condition +// CHECK-SAME: memref<1xf32, strided<[1], offset: 1>> diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/tensor_pointer_descriptor_loop.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/tensor_pointer_descriptor_loop.mlir new file mode 100644 index 0000000000..ed8efe2011 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/tensor_pointer_descriptor_loop.mlir @@ -0,0 +1,28 @@ +// RUN: triton-opt --triton-control-flow-opt --triton-to-linalg %s -verify-each | FileCheck %s + +module attributes {hacc.target = #hacc.target<"Ascend910B2">} { + tt.func public @tensor_pointer_descriptor_loop( + %base: !tt.ptr, %output: !tt.ptr, %upper: index) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %delta = arith.constant dense<1> : tensor<4xi32> + %range = tt.make_range {end = 4 : i32, start = 0 : i32} : tensor<4xi32> + %base_tensor = tt.splat %base : !tt.ptr -> tensor<4x!tt.ptr> + %initial = tt.addptr %base_tensor, %range : tensor<4x!tt.ptr>, tensor<4xi32> + %final = scf.for %iv = %c0 to %upper step %c1 iter_args(%ptr = %initial) -> (tensor<4x!tt.ptr>) { + %next = tt.addptr %ptr, %delta : tensor<4x!tt.ptr>, tensor<4xi32> + scf.yield %next : tensor<4x!tt.ptr> + } + %value = tt.load %final : tensor<4x!tt.ptr> + %output_tensor = tt.splat %output : !tt.ptr -> tensor<4x!tt.ptr> + %output_ptr = tt.addptr %output_tensor, %range : tensor<4x!tt.ptr>, tensor<4xi32> + tt.store %output_ptr, %value : tensor<4x!tt.ptr> + tt.return + } +} + +// CHECK-LABEL: func.func @tensor_pointer_descriptor_loop +// CHECK: scf.for +// CHECK-SAME: tensor<4xi32> +// CHECK-NOT: tensor<4x!tt.ptr +// CHECK-NOT: unrealized_conversion_cast diff --git a/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/bubbleupoperation.mlir b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/bubbleupoperation.mlir index 626f670d1d..4d3228436d 100644 --- a/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/bubbleupoperation.mlir +++ b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/bubbleupoperation.mlir @@ -69,6 +69,18 @@ tt.func @test_addptr_extract_bubbleup(%a: tensor<128x!tt.ptr>, %b: tensor<1 tt.return %1 : !tt.ptr } +// CHECK-LABEL: tt.func @test_select_pointer_extract_bubbleup +// CHECK: %[[CONDITION:.*]] = tensor.extract %{{.*}}[%{{.*}}] {DiscreteMemAccess} : tensor<4xi1> +// CHECK: %[[TRUE_VALUE:.*]] = tensor.extract %{{.*}}[%{{.*}}] {DiscreteMemAccess} : tensor<4x!tt.ptr> +// CHECK: %[[FALSE_VALUE:.*]] = tensor.extract %{{.*}}[%{{.*}}] {DiscreteMemAccess} : tensor<4x!tt.ptr> +// CHECK: arith.select %[[CONDITION]], %[[TRUE_VALUE]], %[[FALSE_VALUE]] : !tt.ptr +tt.func @test_select_pointer_extract_bubbleup( + %condition: tensor<4xi1>, %true_value: tensor<4x!tt.ptr>, + %false_value: tensor<4x!tt.ptr>, %i: index) -> !tt.ptr { + %selected = arith.select %condition, %true_value, %false_value : tensor<4xi1>, tensor<4x!tt.ptr> + %pointer = tensor.extract %selected[%i] : tensor<4x!tt.ptr> + tt.return %pointer : !tt.ptr +} // CHECK-LABEL: tt.func @test_ceil_extract_bubbleup tt.func @test_ceil_extract_bubbleup(%a: tensor<128xf32>, %i: index, %c: f32) -> f32 { diff --git a/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/tensor_pointer_select.mlir b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/tensor_pointer_select.mlir new file mode 100644 index 0000000000..6ae53301db --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/tensor_pointer_select.mlir @@ -0,0 +1,27 @@ +// RUN: triton-opt %s --triton-to-unstructure | FileCheck %s + +tt.func public @opaque_tensor_pointer_select( + %lhs: !tt.ptr, %rhs: !tt.ptr) -> tensor<4xf32> { + %zero = arith.constant dense<0> : tensor<4xi32> + %one = arith.constant dense<1> : tensor<4xi32> + %range = tt.make_range {end = 4 : i32, start = 0 : i32} : tensor<4xi32> + %lhs_tensor = tt.splat %lhs : !tt.ptr -> tensor<4x!tt.ptr> + %rhs_tensor = tt.splat %rhs : !tt.ptr -> tensor<4x!tt.ptr> + %lhs_ptrs = tt.addptr %lhs_tensor, %range : tensor<4x!tt.ptr>, tensor<4xi32> + %rhs_ptrs = tt.addptr %rhs_tensor, %range : tensor<4x!tt.ptr>, tensor<4xi32> + %bits = arith.andi %range, %one : tensor<4xi32> + %condition = arith.cmpi eq, %bits, %zero : tensor<4xi32> + %selected = arith.select %condition, %lhs_ptrs, %rhs_ptrs : tensor<4xi1>, tensor<4x!tt.ptr> + %advanced = tt.addptr %selected, %one : tensor<4x!tt.ptr>, tensor<4xi32> + %loaded = tt.load %advanced : tensor<4x!tt.ptr> + tt.return %loaded : tensor<4xf32> +} + +// CHECK-LABEL: tt.func public @opaque_tensor_pointer_select +// CHECK: %[[SELECTED:.*]] = arith.select {{.*}} : tensor<4xi1>, tensor<4x!tt.ptr> +// CHECK: scf.for +// CHECK: %[[LANE_PTR:.*]] = tensor.extract %[[SELECTED]]{{\[}}%{{.*}}] {DiscreteMemAccess} : tensor<4x!tt.ptr> +// CHECK: %[[ACCESS_PTR:.*]] = tt.addptr %[[LANE_PTR]], +// CHECK-SAME: : !tt.ptr, i64 +// CHECK: tt.load %[[ACCESS_PTR]] {DiscreteMemAccess} : !tt.ptr +// CHECK: tt.return diff --git a/third_party/ascend/unittest/pytest_ut/test_control_flow_pointer_boundary.py b/third_party/ascend/unittest/pytest_ut/test_control_flow_pointer_boundary.py new file mode 100644 index 0000000000..75e74c26d2 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_control_flow_pointer_boundary.py @@ -0,0 +1,247 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import torch +import torch_npu # noqa: F401 +import triton +import triton.language as tl + +BLOCK = 32 +N = 128 + + +@triton.jit +def block_ptr_if_descriptor_kernel(a, b, out, choose_ptr, descriptor_ptr, BLOCK: tl.constexpr): + choose = tl.load(choose_ptr) + size_a = tl.load(descriptor_ptr) + size_b = tl.load(descriptor_ptr + 1) + stride_a = tl.load(descriptor_ptr + 2) + stride_b = tl.load(descriptor_ptr + 3) + offset_a = tl.load(descriptor_ptr + 4) + offset_b = tl.load(descriptor_ptr + 5) + if choose != 0: + pointer = tl.make_block_ptr( + base=a + 3, + shape=(size_a, ), + strides=(stride_a, ), + offsets=(offset_a, ), + block_shape=(BLOCK, ), + order=(0, ), + ) + else: + pointer = tl.make_block_ptr( + base=b + 11, + shape=(size_b, ), + strides=(stride_b, ), + offsets=(offset_b, ), + block_shape=(BLOCK, ), + order=(0, ), + ) + value = tl.load(pointer, boundary_check=(0, ), padding_option="zero") + tl.store(out + tl.arange(0, BLOCK), value) + + +@triton.jit +def block_ptr_for_descriptor_kernel(a, b, out, scalar_out, steps_ptr, descriptor_ptr, BLOCK: tl.constexpr): + steps = tl.load(steps_ptr) + size_a = tl.load(descriptor_ptr) + size_b = tl.load(descriptor_ptr + 1) + stride_a = tl.load(descriptor_ptr + 2) + stride_b = tl.load(descriptor_ptr + 3) + pointer = tl.make_block_ptr( + base=a, + shape=(size_a, ), + strides=(stride_a, ), + offsets=(1, ), + block_shape=(BLOCK, ), + order=(0, ), + ) + ordinary_result = 17 + for i in tl.range(0, steps): + if (i & 1) == 0: + pointer = tl.advance(pointer, (2, )) + else: + pointer = tl.make_block_ptr( + base=b, + shape=(size_b, ), + strides=(stride_b, ), + offsets=(i + 3, ), + block_shape=(BLOCK, ), + order=(0, ), + ) + ordinary_result = ordinary_result + i + 1 + value = tl.load(pointer, boundary_check=(0, ), padding_option="zero") + tl.store(out + tl.arange(0, BLOCK), value) + tl.store(scalar_out, ordinary_result) + + +@triton.jit +def block_ptr_while_descriptor_kernel(a, b, out, steps_ptr, switch_at_ptr, n: tl.constexpr, BLOCK: tl.constexpr): + steps = tl.load(steps_ptr) + switch_at = tl.load(switch_at_ptr) + pointer = tl.make_block_ptr( + base=a, + shape=(n, ), + strides=(1, ), + offsets=(0, ), + block_shape=(BLOCK, ), + order=(0, ), + ) + i = 0 + while i < steps: + if i == switch_at: + pointer = tl.make_block_ptr( + base=b, + shape=(n - 3, ), + strides=(1, ), + offsets=(1, ), + block_shape=(BLOCK, ), + order=(0, ), + ) + else: + pointer = tl.advance(pointer, (2, )) + i += 1 + value = tl.load(pointer, boundary_check=(0, ), padding_option="zero") + tl.store(out + tl.arange(0, BLOCK), value) + + +@triton.jit +def scalar_base_tensor_ptr_loop_kernel(x, out, steps_ptr, n: tl.constexpr, BLOCK: tl.constexpr): + steps = tl.load(steps_ptr) + lane = tl.arange(0, BLOCK) + pointers = x + 3 + lane + for _ in tl.range(0, steps): + pointers = pointers + 2 + index = 3 + lane + 2 * steps + value = tl.load(pointers, mask=index < n, other=0.0) + tl.store(out + lane, value) + + +@triton.jit +def opaque_tensor_ptr_loop_kernel(a, b, out, steps_ptr, n: tl.constexpr, BLOCK: tl.constexpr): + steps = tl.load(steps_ptr) + lane = tl.arange(0, BLOCK) + pointers = tl.where((lane & 1) == 0, a + lane, b + lane) + for _ in tl.range(0, steps): + pointers = pointers + 1 + index = lane + steps + value = tl.load(pointers, mask=index < n, other=0.0) + tl.store(out + lane, value) + + +def _inputs(): + a_cpu = torch.arange(2048, dtype=torch.float32) + b_cpu = 100000.0 + torch.arange(2048, dtype=torch.float32) * 3.0 + return a_cpu, b_cpu, a_cpu.npu(), b_cpu.npu() + + +def _device_i32(value): + return torch.tensor([value], dtype=torch.int32, device="npu") + + +def _slice(source, base_offset, logical_size, stride, offset): + expected = torch.zeros(BLOCK, dtype=source.dtype) + for lane in range(BLOCK): + logical_index = offset + lane + if 0 <= logical_index < logical_size: + expected[lane] = source[base_offset + logical_index * stride] + return expected + + +def _assert_output(actual, expected): + torch.npu.synchronize() + torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("choose", [0, 1]) +def test_block_ptr_if_carries_complete_dynamic_descriptor(choose): + a_cpu, b_cpu, a, b = _inputs() + descriptor = torch.tensor([40, 29, 2, 3, 3, 5], dtype=torch.int64, device="npu") + choose_ptr = _device_i32(choose) + out = torch.empty(BLOCK, dtype=torch.float32, device="npu") + + block_ptr_if_descriptor_kernel[(1, )](a, b, out, choose_ptr, descriptor, BLOCK=BLOCK) + + if choose: + expected = _slice(a_cpu, 3, 40, 2, 3) + else: + expected = _slice(b_cpu, 11, 29, 3, 5) + _assert_output(out, expected) + + +@pytest.mark.parametrize("steps", [0, 2, 5]) +def test_block_ptr_for_carries_changing_descriptor_and_ordinary_result(steps): + a_cpu, b_cpu, a, b = _inputs() + descriptor = torch.tensor([60, 55, 2, 3], dtype=torch.int64, device="npu") + out = torch.empty(BLOCK, dtype=torch.float32, device="npu") + scalar_out = torch.empty(1, dtype=torch.int32, device="npu") + + block_ptr_for_descriptor_kernel[(1, )](a, b, out, scalar_out, _device_i32(steps), descriptor, BLOCK=BLOCK) + + source, logical_size, stride, offset = a_cpu, 60, 2, 1 + for i in range(steps): + if (i & 1) == 0: + offset += 2 + else: + source, logical_size, stride, offset = b_cpu, 55, 3, i + 3 + _assert_output(out, _slice(source, 0, logical_size, stride, offset)) + assert scalar_out.cpu().item() == 17 + steps * (steps + 1) // 2 + + +@pytest.mark.parametrize("steps,switch_at", [(0, -1), (4, -1), (4, 2)]) +def test_block_ptr_while_carries_descriptor(steps, switch_at): + a_cpu, b_cpu, a, b = _inputs() + out = torch.empty(BLOCK, dtype=torch.float32, device="npu") + + block_ptr_while_descriptor_kernel[(1, )](a, b, out, _device_i32(steps), _device_i32(switch_at), n=N, BLOCK=BLOCK) + + source, logical_size, offset = a_cpu, N, 0 + for i in range(steps): + if i == switch_at: + source, logical_size, offset = b_cpu, N - 3, 1 + else: + offset += 2 + _assert_output(out, _slice(source, 0, logical_size, 1, offset)) + + +@pytest.mark.parametrize("steps", [0, 2, 4]) +def test_scalar_base_tensor_pointer_loop(steps): + a_cpu, _, a, _ = _inputs() + out = torch.empty(BLOCK, dtype=torch.float32, device="npu") + + scalar_base_tensor_ptr_loop_kernel[(1, )](a, out, _device_i32(steps), n=N, BLOCK=BLOCK) + + _assert_output(out, _slice(a_cpu, 0, N, 1, 3 + 2 * steps)) + + +@pytest.mark.parametrize("steps", [0, 2, 4]) +def test_opaque_tensor_pointer_loop(steps): + a_cpu, b_cpu, a, b = _inputs() + out = torch.empty(BLOCK, dtype=torch.float32, device="npu") + + opaque_tensor_ptr_loop_kernel[(1, )](a, b, out, _device_i32(steps), n=N, BLOCK=BLOCK) + + expected = torch.zeros(BLOCK, dtype=torch.float32) + for lane in range(BLOCK): + index = lane + steps + if index < N: + expected[lane] = a_cpu[index] if (lane & 1) == 0 else b_cpu[index] + _assert_output(out, expected)