From b4f9ee56f9ed80027e77bd0796a071f618510cfe Mon Sep 17 00:00:00 2001 From: zhaojingkai Date: Thu, 13 Aug 2026 09:50:03 +0800 Subject: [PATCH] feat(control-flow): decouple pointer descriptors across SCF Carry complete block-pointer descriptors across supported for, while, and if boundaries as pointer-free SSA components, and rebuild pointers only at their use sites. Keep tensor-pointer bases outside loop signatures when they are invariant while carrying complete lane offsets. Teach TritonToLinalg and TritonToUnstructure to lower integer pointer carriers, scalar and opaque pointer joins, descriptor loops, lane offsets, and rebased memref layouts. Preserve exact descriptor producer slots, retain legacy conversion for mixed pointer boundaries, and keep externally typed memref boundaries layout-compatible. Remove the module-wide addptr-base restriction so unrelated and local make_tensor_ptr operations retain their previous behavior. Add one end-to-end pytest covering dynamic if/for/while block-pointer descriptors, changing bases, ordinary loop results, scalar-base tensor pointers, and opaque lane-wise tensor pointers. --- .../ControlFlowAnalysis.h | 1 - .../TritonControlFlowOpt/ControlFlowRewrite.h | 28 +- .../include/TritonToLinalg/BlockPtrAnalysis.h | 20 +- .../TritonToLinalg/TritonOpConverter.h | 79 ++- .../TritonToUnstructure/BubbleUpOperation.h | 5 + .../BlockPtrDecompose.cpp | 147 +++-- .../ControlFlowAnalysis.cpp | 6 +- .../ControlFlowRewrite.cpp | 142 +++-- .../TensorPtrDecompose.cpp | 150 ++--- .../lib/TritonToLinalg/BlockPtrAnalysis.cpp | 226 +++++--- .../lib/TritonToLinalg/TritonOpConverter.cpp | 336 ++++++++++- .../lib/TritonToLinalg/TritonToLinalgPass.cpp | 532 +++++++++++++++--- .../TritonToUnstructure/BubbleUpOperation.cpp | 16 + .../TritonToUnstructure/OffsetAnalysis.cpp | 65 +++ .../TritonToUnstructure/ReplaceArguments.cpp | 11 +- .../UnstructureConversionPass.cpp | 31 +- .../block_ptr_addptr_base.mlir | 95 ++++ ...lid.mlir => block_ptr_different_base.mlir} | 16 +- .../pointer_descriptor_boundary_marker.mlir | 74 +++ ..._descriptor_boundary_mixed_downstream.mlir | 38 ++ .../scf_pointer_decouple.mlir | 181 +++--- .../TritonToLinalg/scalar_pointer_select.mlir | 207 +++++++ .../tensor_pointer_descriptor_loop.mlir | 28 + .../bubbleupoperation.mlir | 12 + .../tensor_pointer_select.mlir | 27 + .../test_control_flow_pointer_boundary.py | 247 ++++++++ 26 files changed, 2277 insertions(+), 443 deletions(-) create mode 100644 third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/block_ptr_addptr_base.mlir rename third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/{block_ptr_different_base_invalid.mlir => block_ptr_different_base.mlir} (50%) create mode 100644 third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/pointer_descriptor_boundary_marker.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/pointer_descriptor_boundary_mixed_downstream.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToLinalg/scalar_pointer_select.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToLinalg/tensor_pointer_descriptor_loop.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToUnstructure/tensor_pointer_select.mlir create mode 100644 third_party/ascend/unittest/pytest_ut/test_control_flow_pointer_boundary.py 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)