diff --git a/include/pypto/codegen/pto/pto_codegen.h b/include/pypto/codegen/pto/pto_codegen.h index 361c96ce20..b1b4bda049 100644 --- a/include/pypto/codegen/pto/pto_codegen.h +++ b/include/pypto/codegen/pto/pto_codegen.h @@ -370,6 +370,8 @@ class PTOCodegen : public CodegenBase { void SetCurrentResultBuf(const std::string& buf); void RegisterTileBufType(const std::string& ssa_name, const std::string& type_string); std::string GetSSATileBufType(const std::string& ssa_name) const; + void AliasTileVarToExistingBuf(const ir::VarPtr& var, const std::string& ssa_name, + const std::string& type_string = ""); struct SubviewMaterializationInfo { std::string source_ssa; std::string source_type; diff --git a/include/pypto/ir/transforms/utils/cross_core_pipe.h b/include/pypto/ir/transforms/utils/cross_core_pipe.h index 7f42ed30e2..df09eec420 100644 --- a/include/pypto/ir/transforms/utils/cross_core_pipe.h +++ b/include/pypto/ir/transforms/utils/cross_core_pipe.h @@ -80,7 +80,7 @@ CallPtr CreateImportPeerBuffer(const std::string& buffer_name, const std::string // (otherwise PTOAS derives it from `dir_mask`). CallPtr CreateInitializePipe(core_affinity::CoreSide side, int dir_mask, int slot_size_bytes, const ExprPtr& c2v_consumer_buf, const ExprPtr& v2c_consumer_buf, - std::optional slot_num, const Span& span); + std::optional slot_num, int pipe_id, const Span& span); void CollectCrossCorePipeMetadata(const std::vector& stmts, CrossCorePipeMetadata& metadata); CrossCorePipeMetadata CollectDominatingPipeSetupMetadata(const std::vector& stmts); diff --git a/src/backend/common/pto_ops_common.cpp b/src/backend/common/pto_ops_common.cpp index 306dd2cdb0..629411fde9 100644 --- a/src/backend/common/pto_ops_common.cpp +++ b/src/backend/common/pto_ops_common.cpp @@ -2221,8 +2221,10 @@ static void EmitLogicalTpushValidShapeRestore(const CallPtr& op, codegen::PTOCod static std::string FormatFrontendPipeAttrs(const CallPtr& op, int split) { std::ostringstream oss; oss << "{"; - if (op->HasKwarg("id")) { - const int id = op->GetKwarg("id", 0); + const bool has_explicit_id = op->HasKwarg("id"); + const int id = has_explicit_id ? op->GetKwarg("id", 0) : (split != 0 ? split : 0); + // Keep an explicit id=0 distinguishable from the omitted-id default. + if (id != 0 || has_explicit_id) { CHECK(id >= 0) << "Frontend pipe 'id' attribute must be non-negative, got " << id; oss << "id = " << id << ", "; } @@ -3506,7 +3508,22 @@ void RegisterPTOOps(Backend& backend, const std::unordered_set& exc if (src_space.has_value() && dst_space.has_value() && *src_space == *dst_space) { auto src_offset = As((*src_tile->memref_)->byte_offset_); auto dst_offset = As((*dst_tile->memref_)->byte_offset_); - if (src_offset && dst_offset && src_offset->value_ == dst_offset->value_) { + const bool same_dtype = src_tile->dtype_ == dst_tile->dtype_; + const bool same_tile_view = ir::tile_view_semantics::GetEffectiveTileView(*src_tile) == + ir::tile_view_semantics::GetEffectiveTileView(*dst_tile); + bool same_shape = src_tile->shape_.size() == dst_tile->shape_.size(); + for (size_t i = 0; same_shape && i < src_tile->shape_.size(); ++i) { + same_shape = IsSameDimExpr(src_tile->shape_[i], dst_tile->shape_[i]); + } + if (*src_space == ir::MemorySpace::Acc && same_dtype && same_shape) { + // PTOAS has no Acc->Acc tmov form. Treat same-shaped Acc moves as + // SSA aliases; these are inserted as layout/placement repairs before + // consumers such as tile.store, and copying would be illegal anyway. + codegen.SetCurrentResultBuf(codegen.GetExprAsCode(op->args_[0])); + return std::string(""); + } + if (src_offset && dst_offset && src_offset->value_ == dst_offset->value_ && same_dtype && + same_tile_view) { // Alias the destination to the source SSA value so downstream // references use the source's defined buffer, not the destination's // alloc_tile (which would be unwritten after eliding the tmov). @@ -3620,9 +3637,14 @@ void RegisterPTOOps(Backend& backend, const std::unordered_set& exc reg("tile.load", [](const ir::CallPtr& op, codegen::CodegenBase& codegen) { return MakeTileLoadCodegenPTO(op, codegen); }); - reg("tile.store", [](const ir::CallPtr& op, codegen::CodegenBase& codegen) { - return MakeTileStoreCodegenPTO(op, codegen); - }); + if (exclude_ops.count("tile.store") == 0) { + auto reg_entry = backend.RegisterOp("tile.store"); + reg_entry + .f_codegen([](const ir::CallPtr& op, codegen::CodegenBase& codegen) { + return MakeTileStoreCodegenPTO(op, codegen); + }) + .set_input_layout(0, ir::TileLayout::row_major); + } // Distributed N6 ops — cross-rank tile load + per-rank signal notify/wait + // synchronous bulk get/put. See MakeRemoteLoadCodegenPTO / // MakeNotifyCodegenPTO / MakeWaitCodegenPTO / MakeGetCodegenPTO / diff --git a/src/codegen/pto/pto_codegen.cpp b/src/codegen/pto/pto_codegen.cpp index 1baf45b024..52c46cc66b 100644 --- a/src/codegen/pto/pto_codegen.cpp +++ b/src/codegen/pto/pto_codegen.cpp @@ -1255,6 +1255,16 @@ void PTOCodegen::RegisterTileBufType(const std::string& ssa_name, const std::str fs_.ssa_to_tile_buf_type[ssa_name] = type_string; } +void PTOCodegen::AliasTileVarToExistingBuf(const ir::VarPtr& var, const std::string& ssa_name, + const std::string& type_string) { + INTERNAL_CHECK(var != nullptr) << "Internal error: cannot alias null tile var"; + INTERNAL_CHECK(!ssa_name.empty()) << "Internal error: cannot alias tile var to empty SSA"; + BindVarToMlir(var, ssa_name); + if (!type_string.empty()) { + RegisterTileBufType(ssa_name, type_string); + } +} + std::string PTOCodegen::GetSSATileBufType(const std::string& ssa_name) const { auto it = fs_.ssa_to_tile_buf_type.find(ssa_name); return it != fs_.ssa_to_tile_buf_type.end() ? it->second : std::string{}; @@ -1383,6 +1393,13 @@ void PTOCodegen::VisitStmt_(const AssignStmtPtr& op) { const bool alias_scatter_result_to_input = ShouldAliasScatterResultToInput(op); const bool alias_array_update_to_input = ShouldAliasArrayUpdateResultToInput(op); + if (As(op->value_)) { + auto key = GetVarKey(op->var_); + if (fs_.var_to_mlir.find(key) != fs_.var_to_mlir.end()) { + return; + } + } + if (auto tile_type = ir::GetTileTypeWithMemRef(op->var_->GetType())) { if (!is_set_validshape && !alias_scatter_result_to_input) { EmitAllocTileForVar(op->var_, tile_type); diff --git a/src/ir/op/tile_ops/memory.cpp b/src/ir/op/tile_ops/memory.cpp index 48ca8c09ae..0d4a546c2a 100644 --- a/src/ir/op/tile_ops/memory.cpp +++ b/src/ir/op/tile_ops/memory.cpp @@ -398,6 +398,12 @@ TypePtr DeduceTileCreateTileType(const std::vector& args, std::optional creation_space = std::nullopt; if (flat_layout) { creation_space = MemorySpace::Mat; + } else if (target_memory_opt.has_value() && *target_memory_opt == MemorySpace::Left) { + tile_view.blayout = TileLayout::col_major; + tile_view.slayout = TileLayout::row_major; + } else if (target_memory_opt.has_value() && *target_memory_opt == MemorySpace::Right) { + tile_view.blayout = TileLayout::row_major; + tile_view.slayout = TileLayout::col_major; } else if (target_memory_opt.has_value() && *target_memory_opt == MemorySpace::Acc) { tile_view.blayout = TileLayout::col_major; tile_view.slayout = TileLayout::row_major; diff --git a/src/ir/transforms/convert_tensor_to_tile_ops_pass.cpp b/src/ir/transforms/convert_tensor_to_tile_ops_pass.cpp index 1b43751a4d..7d1a1bf459 100644 --- a/src/ir/transforms/convert_tensor_to_tile_ops_pass.cpp +++ b/src/ir/transforms/convert_tensor_to_tile_ops_pass.cpp @@ -348,6 +348,17 @@ class TypePropagatingMutator : public IRMutator { return std::make_shared(op->name_hint_, new_init->GetType(), new_init, op->span_); } + ExprPtr VisitExpr_(const TupleGetItemExprPtr& op) override { + auto tuple = VisitExpr(op->tuple_); + if (auto make_tuple = As(tuple)) { + if (op->index_ >= 0 && static_cast(op->index_) < make_tuple->elements_.size()) { + return VisitExpr(make_tuple->elements_[static_cast(op->index_)]); + } + } + if (tuple.get() == op->tuple_.get()) return op; + return std::make_shared(tuple, op->index_, op->span_); + } + /// Override ForStmt to update return_vars types to match iter_arg types. StmtPtr VisitStmt_(const ForStmtPtr& op) override { auto result = IRMutator::VisitStmt_(op); @@ -576,6 +587,11 @@ class TensorToTileMutator : public TypePropagatingMutator { // remapped vars that the result expression references. auto new_result = VisitExpr(conv_result.result); + if (As(new_result)) { + var_remap_[op->var_.get()] = new_result; + return SeqStmts::Flatten(std::move(stmts), op->span_); + } + auto tile_name = MakeTileValueName(op->var_->name_hint_); auto tile_var = std::make_shared(tile_name, new_result->GetType(), op->var_->span_); stmts.push_back(std::make_shared(tile_var, new_result, op->span_)); diff --git a/src/ir/transforms/convert_to_ssa_pass.cpp b/src/ir/transforms/convert_to_ssa_pass.cpp index 74dd81d7a4..eac76268a0 100644 --- a/src/ir/transforms/convert_to_ssa_pass.cpp +++ b/src/ir/transforms/convert_to_ssa_pass.cpp @@ -819,6 +819,47 @@ class SSAConverter { auto cond = SubstExpr(op->condition_); auto before = cur_; + // Frontend constants such as ``if True:`` can wrap manual-dependency + // producers that are consumed later in the parent scope. If we keep the + // artificial IfStmt until SSA conversion, those producer variables appear + // branch-local and the verifier rejects their later uses. Fold constants + // here, including explicit return_vars: the kept branch's trailing yield is + // stripped and each return_var is rebound to its yielded SSA value. + if (auto const_cond = StaticBoolValue(cond)) { + if (!*const_cond && !op->else_body_.has_value()) { + cur_ = before; + return std::make_shared(std::vector{}, op->span_); + } + auto kept_body = ConvertStmt(*const_cond ? op->then_body_ : *op->else_body_); + if (op->return_vars_.empty()) { + return kept_body; + } + auto stripped = StripTrailingYield(kept_body, op->return_vars_.size()); + std::vector extra_assigns; + for (size_t i = 0; i < op->return_vars_.size(); ++i) { + auto rv_key = op->return_vars_[i].get(); + auto yielded = stripped.yielded_values[i]; + if (auto yielded_var = As(yielded)) { + cur_[rv_key] = yielded_var; + continue; + } + int v = NextVersion(rv_key); + auto nrv = std::make_shared(BuildAutoNamedVersion(rv_key->name_hint_, "rv", v), + op->return_vars_[i]->GetType(), op->return_vars_[i]->span_); + extra_assigns.push_back(std::make_shared(nrv, yielded, op->span_)); + cur_[rv_key] = nrv; + } + if (extra_assigns.empty()) return stripped.body; + std::vector stmts; + if (auto seq = As(stripped.body)) { + stmts = seq->stmts_; + } else if (stripped.body) { + stmts.push_back(stripped.body); + } + stmts.insert(stmts.end(), extra_assigns.begin(), extra_assigns.end()); + return SeqStmts::Flatten(std::move(stmts), op->span_); + } + // Convert then branch auto new_then = ConvertStmt(op->then_body_); auto then_ver = cur_; @@ -1224,6 +1265,55 @@ class SSAConverter { return SeqStmts::Flatten({s, yield}, span); } + struct StrippedYield { + StmtPtr body; + std::vector yielded_values; + }; + + static StmtPtr EmptyBody(const Span& span) { + return std::make_shared(std::vector{}, span); + } + + static std::optional StaticBoolValue(const ExprPtr& e) { + if (auto b = As(e)) return b->value_; + if (auto i = As(e); i && i->dtype() == DataType::BOOL) return i->value_ != 0; + return std::nullopt; + } + + static StrippedYield StripTrailingYield(const StmtPtr& s, size_t return_var_count) { + AssertNoMidBodyYield(s); + if (auto scope = As(s)) { + auto copy = MutableCopy(scope); + auto stripped = StripTrailingYield(scope->body_, return_var_count); + copy->body_ = stripped.body ? stripped.body : EmptyBody(scope->body_->span_); + return {copy, std::move(stripped.yielded_values)}; + } + if (auto scope = As(s)) { + auto copy = MutableCopy(scope); + auto stripped = StripTrailingYield(scope->body_, return_var_count); + copy->body_ = stripped.body ? stripped.body : EmptyBody(scope->body_->span_); + return {copy, std::move(stripped.yielded_values)}; + } + if (auto seq = As(s)) { + INTERNAL_CHECK_SPAN(!seq->stmts_.empty(), seq->span_) + << "ConvertToSSA: IfStmt with return_vars must end with YieldStmt"; + auto stmts = seq->stmts_; + auto stripped = StripTrailingYield(stmts.back(), return_var_count); + if (stripped.body) { + stmts.back() = stripped.body; + } else { + stmts.pop_back(); + } + return {SeqStmts::Flatten(std::move(stmts), seq->span_), std::move(stripped.yielded_values)}; + } + auto yield = As(s); + INTERNAL_CHECK_SPAN(yield, s->span_) << "ConvertToSSA: IfStmt with return_vars must end with YieldStmt"; + INTERNAL_CHECK_SPAN(yield->value_.size() == return_var_count, yield->span_) + << "ConvertToSSA: yielded value count " << yield->value_.size() + << " does not match return_vars count " << return_var_count; + return {nullptr, yield->value_}; + } + // ── State ────────────────────────────────────────────────────────── std::unordered_map cur_; // var pointer → latest version diff --git a/src/ir/transforms/expand_mixed_kernel_pass.cpp b/src/ir/transforms/expand_mixed_kernel_pass.cpp index 8debee2a17..0ac72ab78b 100644 --- a/src/ir/transforms/expand_mixed_kernel_pass.cpp +++ b/src/ir/transforms/expand_mixed_kernel_pass.cpp @@ -345,7 +345,11 @@ void CollectCVBoundaryMoves(const std::vector& stmts, // ============================================================================ std::vector> MakeSplitKwargs(int split = 0) { - return {{"split", std::any(split)}}; + std::vector> kwargs{{"split", std::any(split)}}; + if (split != 0) { + kwargs.emplace_back("id", std::any(split)); + } + return kwargs; } CallPtr CreateTpush(const std::string& op_name, const ExprPtr& tile, const Span& span, int split = 0) { @@ -393,6 +397,54 @@ bool NeedsPostTpopMove(CoreSide side, const TileType& dest_type) { return dest_type.memory_space_.value() != GetBoundaryTpopMemory(side); } +bool SameDimExpr(const ExprPtr& lhs, const ExprPtr& rhs) { + if (lhs.get() == rhs.get()) return true; + auto lhs_const = std::dynamic_pointer_cast(lhs); + auto rhs_const = std::dynamic_pointer_cast(rhs); + return lhs_const && rhs_const && lhs_const->value_ == rhs_const->value_; +} + +bool SameTileShape(const TileType& lhs, const TileType& rhs) { + if (lhs.shape_.size() != rhs.shape_.size()) return false; + for (size_t i = 0; i < lhs.shape_.size(); ++i) { + if (!SameDimExpr(lhs.shape_[i], rhs.shape_[i])) return false; + } + return true; +} + +bool IsSameSpaceTpopMoveAlias(const AssignStmtPtr& assign, const TpopDefs& tpop_defs, + const std::unordered_map& tpop_var_remap, + VarPtr* canonical_tpop) { + if (!assign) return false; + auto call = std::dynamic_pointer_cast(assign->value_); + if (!IsOp(call, "tile.move") || call->args_.empty()) return false; + auto source_var = AsVarLike(call->args_[0]); + if (!source_var) return false; + + VarPtr source_after_remap = source_var; + if (auto remap_it = tpop_var_remap.find(source_var.get()); + remap_it != tpop_var_remap.end() && remap_it->second) { + source_after_remap = remap_it->second; + } else if (tpop_defs.count(source_var.get()) == 0) { + return false; + } + + auto src_type = std::dynamic_pointer_cast(source_after_remap->GetType()); + auto dst_type = std::dynamic_pointer_cast(assign->var_->GetType()); + if (!src_type || !dst_type) return false; + auto src_space = src_type->GetMemorySpace(); + auto dst_space = dst_type->GetMemorySpace(); + if (!src_space.has_value() || !dst_space.has_value() || src_space.value() != dst_space.value()) return false; + if (src_type->dtype_ != dst_type->dtype_) return false; + if (!SameTileShape(*src_type, *dst_type)) return false; + if (tile_view_semantics::GetEffectiveTileView(*src_type) != tile_view_semantics::GetEffectiveTileView(*dst_type)) { + return false; + } + + if (canonical_tpop) *canonical_tpop = source_after_remap; + return true; +} + std::string BuildBoundaryTpopName(CoreSide side, const std::string& dest_name) { return dest_name + ((side == CoreSide::AIC) ? "_mat" : "_vec"); } @@ -670,7 +722,8 @@ std::vector BuildCoreBody(CoreSide side, const std::vector& st std::unordered_map& tpop_var_remap, std::unordered_set& superseded_tpop_vars, const std::map& gm_sync_pushes, - const std::map>& gm_sync_pops) { + const std::map>& gm_sync_pops, + const TpopDefs& tpop_defs) { const auto* handler = PassContext::Current()->GetBackendHandler(); // AIC keeps CUBE, skips VECTOR; AIV keeps VECTOR, skips CUBE CoreAffinity keep_affinity = (side == CoreSide::AIC) ? CoreAffinity::CUBE : CoreAffinity::VECTOR; @@ -765,7 +818,7 @@ std::vector BuildCoreBody(CoreSide side, const std::vector& st view_ms = shape_tt->memory_space_.value(); // NOLINT(bugprone-unchecked-optional-access) needs_post_move = NeedsPostTpopMove(side, *shape_tt); } - auto tpop_type = BuildBoundaryTpopType(side, shape_source); + auto tpop_type = needs_post_move ? BuildBoundaryTpopType(side, shape_source) : shape_source; // Consumer-side transfer view. For op-driven boundaries the cross-core // data lands in a FRESH transfer tile of this side's memory (Vec on // AIV, Mat on AIC) — exactly like a plain move-boundary destination — @@ -827,6 +880,14 @@ std::vector BuildCoreBody(CoreSide side, const std::vector& st if (superseded_tpop_vars.count(assign->var_.get()) > 0) continue; } + if (auto assign = std::dynamic_pointer_cast(stmt)) { + VarPtr canonical_tpop; + if (IsSameSpaceTpopMoveAlias(assign, tpop_defs, tpop_var_remap, &canonical_tpop)) { + tpop_var_remap[assign->var_.get()] = canonical_tpop; + continue; + } + } + // GM cross-lane sync (issue #1433): on the consumer lane, emit a fence tpop // just before the keyed stmt. That stmt is the load itself when producer and // consumer share a body, or the loop/branch enclosing the load (so the fence @@ -856,19 +917,19 @@ std::vector BuildCoreBody(CoreSide side, const std::vector& st // Recurse into compound statements, building pruned copies if (auto for_stmt = std::dynamic_pointer_cast(stmt)) { auto new_body = BuildCoreBody(side, FlattenBody(for_stmt->body_), stmt_map, boundary_moves, - tpop_var_remap, superseded_tpop_vars, gm_sync_pushes, gm_sync_pops); + tpop_var_remap, superseded_tpop_vars, gm_sync_pushes, gm_sync_pops, tpop_defs); auto new_for = MutableCopy(for_stmt); new_for->body_ = MakeBody(new_body, for_stmt->span_); result.push_back(new_for); } else if (auto if_stmt = std::dynamic_pointer_cast(stmt)) { auto new_then = BuildCoreBody(side, FlattenBody(if_stmt->then_body_), stmt_map, boundary_moves, - tpop_var_remap, superseded_tpop_vars, gm_sync_pushes, gm_sync_pops); + tpop_var_remap, superseded_tpop_vars, gm_sync_pushes, gm_sync_pops, tpop_defs); std::optional new_else; const auto& else_body = if_stmt->else_body_; if (else_body.has_value()) { auto new_else_stmts = BuildCoreBody(side, FlattenBody(*else_body), stmt_map, boundary_moves, tpop_var_remap, - superseded_tpop_vars, gm_sync_pushes, gm_sync_pops); + superseded_tpop_vars, gm_sync_pushes, gm_sync_pops, tpop_defs); new_else = MakeBody(new_else_stmts, if_stmt->span_); } auto new_if = MutableCopy(if_stmt); @@ -877,7 +938,7 @@ std::vector BuildCoreBody(CoreSide side, const std::vector& st result.push_back(new_if); } else if (auto while_stmt = std::dynamic_pointer_cast(stmt)) { auto new_body = BuildCoreBody(side, FlattenBody(while_stmt->body_), stmt_map, boundary_moves, - tpop_var_remap, superseded_tpop_vars, gm_sync_pushes, gm_sync_pops); + tpop_var_remap, superseded_tpop_vars, gm_sync_pushes, gm_sync_pops, tpop_defs); auto new_while = MutableCopy(while_stmt); new_while->body_ = MakeBody(new_body, while_stmt->span_); result.push_back(new_while); @@ -1145,7 +1206,7 @@ ExpandedKernel ExpandMixedFunction(const FunctionPtr& func, bool create_group = // Build AIC body (recursive — handles MIXED compound stmts) std::unordered_map aic_tpop_remap; auto aic_stmts = BuildCoreBody(CoreSide::AIC, stmts, stmt_map, boundary_moves, aic_tpop_remap, - superseded_tpop_vars, gm_sync_pushes, gm_sync_pops); + superseded_tpop_vars, gm_sync_pushes, gm_sync_pops, tpop_defs); // Remove ReturnStmt from AIC (AIC doesn't return values) std::vector aic_stmts_no_return; @@ -1173,7 +1234,7 @@ ExpandedKernel ExpandMixedFunction(const FunctionPtr& func, bool create_group = // Build AIV body (recursive — handles MIXED compound stmts) std::unordered_map aiv_tpop_remap; auto aiv_stmts = BuildCoreBody(CoreSide::AIV, stmts, stmt_map, boundary_moves, aiv_tpop_remap, - superseded_tpop_vars, gm_sync_pushes, gm_sync_pops); + superseded_tpop_vars, gm_sync_pushes, gm_sync_pops, tpop_defs); auto aiv_final = FinalizeTpopTfrees(FinalizeSplitCoreBody(aiv_stmts, original_def_map, remap_keys(aiv_tpop_remap)), CoreSide::AIV, aiv_tpop_remap); diff --git a/src/ir/transforms/resolve_backend_op_layouts_pass.cpp b/src/ir/transforms/resolve_backend_op_layouts_pass.cpp index bdd4fe170f..ef840300a8 100644 --- a/src/ir/transforms/resolve_backend_op_layouts_pass.cpp +++ b/src/ir/transforms/resolve_backend_op_layouts_pass.cpp @@ -76,6 +76,19 @@ bool RequiresRowMajor(const std::optional& required_layout) { bool IsRowMajor(const TileTypePtr& tile_type) { return GetTileLayout(tile_type) == TileLayout::row_major; } +bool ShouldRepairRowMajorInput(const CallPtr& call, size_t arg_index, const TileTypePtr& tile_type) { + if (!tile_type || IsRowMajor(tile_type)) { + return false; + } + // TSTORE accepts the tile's logical valid shape. Rewriting a [N, 1] column + // vector to [1, N] would change the store footprint, so keep vector stores + // in their original shape and only repair matrix col-major stores. + if (call->op_->name_ == "tile.store" && arg_index == 0 && IsColumnVector(tile_type)) { + return false; + } + return true; +} + ExprPtr MakeShapeTuple(const std::vector& dims, const Span& span) { return std::make_shared(dims, span); } @@ -131,7 +144,7 @@ bool NeedsInputRepair(const CallPtr& call, const backend::BackendTileLayoutSpec& if (!tile_type) { continue; // Non-tile inputs (scalars, shapes) are not subject to layout repair } - if (!IsRowMajor(tile_type)) { + if (ShouldRepairRowMajorInput(call, i, tile_type)) { return true; } } @@ -166,8 +179,11 @@ class BackendLayoutRepairMutator : public IRMutator { return IRMutator::VisitStmt_(op); } - INTERNAL_CHECK_SPAN(result_tile_type, op->span_) - << "ResolveBackendOpLayouts expects constrained op assignment targets to be TileType"; + const bool needs_output_repair = NeedsOutputRepair(result_tile_type, *layout_spec); + if (needs_output_repair) { + INTERNAL_CHECK_SPAN(result_tile_type, op->span_) + << "ResolveBackendOpLayouts expects output-repaired op assignment targets to be TileType"; + } std::vector rewritten; std::vector new_args = call->args_; @@ -179,7 +195,7 @@ class BackendLayoutRepairMutator : public IRMutator { } auto tile_type = As(call->args_[i]->GetType()); - if (!tile_type || IsRowMajor(tile_type)) { + if (!ShouldRepairRowMajorInput(call, i, tile_type)) { continue; } @@ -206,7 +222,7 @@ class BackendLayoutRepairMutator : public IRMutator { INTERNAL_CHECK_SPAN(repaired_call, call->span_) << "ResolveBackendOpLayouts: repaired consumer must remain a Call"; - if (NeedsOutputRepair(result_tile_type, *layout_spec)) { + if (needs_output_repair) { auto row_major_var = std::make_shared(NextTempName(op->var_->name_hint_, {"row_major"}), repaired_call->GetType(), call->span_); rewritten.push_back(std::make_shared(row_major_var, repaired_call, op->span_)); @@ -251,7 +267,7 @@ class BackendLayoutRepairMutator : public IRMutator { continue; } auto tile_type = As(call->args_[i]->GetType()); - if (!tile_type || IsRowMajor(tile_type)) { + if (!ShouldRepairRowMajorInput(call, i, tile_type)) { continue; } auto repair_var_name = NextTempName("layout_fix", {"row_major", "arg" + std::to_string(i)}); diff --git a/src/ir/transforms/utils/cross_core_pipe.cpp b/src/ir/transforms/utils/cross_core_pipe.cpp index 291da19a6c..df82db140a 100644 --- a/src/ir/transforms/utils/cross_core_pipe.cpp +++ b/src/ir/transforms/utils/cross_core_pipe.cpp @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -148,6 +149,15 @@ std::string BuildPipeBufferName(const std::string& func_name, core_affinity::Pip ((direction == core_affinity::PipeDirection::C2V) ? "_c2v_slot_buffer" : "_v2c_slot_buffer"); } +std::string BuildPipeBufferNameForId(const std::string& func_name, core_affinity::PipeDirection direction, + int pipe_id) { + if (pipe_id == 0) { + return BuildPipeBufferName(func_name, direction); + } + return func_name + "_pipe" + std::to_string(pipe_id) + + ((direction == core_affinity::PipeDirection::C2V) ? "_c2v_slot_buffer" : "_v2c_slot_buffer"); +} + CallPtr CreateSystemOpCall(const std::string& op_name, const std::vector>& kwargs, const Span& span) { return CreateSystemOpCall(op_name, {}, kwargs, span); @@ -176,7 +186,7 @@ CallPtr CreateImportPeerBuffer(const std::string& buffer_name, const std::string CallPtr CreateInitializePipe(core_affinity::CoreSide side, int dir_mask, int slot_size_bytes, const ExprPtr& c2v_consumer_buf, const ExprPtr& v2c_consumer_buf, - std::optional slot_num, const Span& span) { + std::optional slot_num, int pipe_id, const Span& span) { INTERNAL_CHECK_SPAN(slot_size_bytes >= 0 && slot_size_bytes <= std::numeric_limits::max(), span) << "Cross-core slot_size out of range: " << slot_size_bytes; std::vector> kwargs = {{"dir_mask", std::any(dir_mask)}, @@ -186,10 +196,54 @@ CallPtr CreateInitializePipe(core_affinity::CoreSide side, int dir_mask, int slo << "Cross-core slot_num override must be positive: " << slot_num.value(); kwargs.emplace_back("slot_num", std::any(slot_num.value())); } + if (pipe_id != 0) { + kwargs.emplace_back("id", std::any(pipe_id)); + } const std::string op_name = core_side_ops::InitializePipeOp(side); return CreateSystemOpCall(op_name, {c2v_consumer_buf, v2c_consumer_buf}, kwargs, span); } +void CollectCrossCorePipeMetadataById(const std::vector& stmts, + std::map& metadata_by_id) { + for (const auto& stmt : stmts) { + auto assign = std::dynamic_pointer_cast(stmt); + auto eval = std::dynamic_pointer_cast(stmt); + CallPtr call; + if (assign) { + call = std::dynamic_pointer_cast(assign->value_); + } else if (eval) { + call = std::dynamic_pointer_cast(eval->expr_); + } + auto op = call ? std::dynamic_pointer_cast(call->op_) : nullptr; + if (op) { + const int split = call->GetKwarg("split", 0); + const int pipe_id = call->HasKwarg("id") ? call->GetKwarg("id", 0) : (split != 0 ? split : 0); + auto& metadata = metadata_by_id[pipe_id]; + if (IsOp(op, "tile.tpush_to_aiv") && call->args_.size() == 1) { + RecordTileSlotSize(metadata.c2v, call->args_[0]->GetType()); + } else if (IsOp(op, "tile.tpush_to_aic") && call->args_.size() == 1) { + RecordTileSlotSize(metadata.v2c, call->args_[0]->GetType()); + } else if (IsOp(op, "tile.tpop_from_aiv") && assign) { + RecordTileSlotSize(metadata.v2c, assign->var_->GetType()); + } else if (IsOp(op, "tile.tpop_from_aic") && assign) { + RecordTileSlotSize(metadata.c2v, assign->var_->GetType()); + } + } + + if (auto for_stmt = std::dynamic_pointer_cast(stmt)) { + CollectCrossCorePipeMetadataById(FlattenBody(for_stmt->body_), metadata_by_id); + } else if (auto if_stmt = std::dynamic_pointer_cast(stmt)) { + CollectCrossCorePipeMetadataById(FlattenBody(if_stmt->then_body_), metadata_by_id); + const auto& else_body = if_stmt->else_body_; + if (else_body) { + CollectCrossCorePipeMetadataById(FlattenBody(*else_body), metadata_by_id); + } + } else if (auto while_stmt = std::dynamic_pointer_cast(stmt)) { + CollectCrossCorePipeMetadataById(FlattenBody(while_stmt->body_), metadata_by_id); + } + } +} + void CollectCrossCorePipeMetadata(const std::vector& stmts, CrossCorePipeMetadata& metadata) { for (const auto& stmt : stmts) { auto assign = std::dynamic_pointer_cast(stmt); @@ -281,76 +335,83 @@ AutomaticPipeSetup BuildAutomaticPipeSetup(const std::string& func_name, const s return {}; } - const int dir_mask = BuildDirMask(combined); - auto common_slot_size = GetCommonSlotSizeBytes(combined); - if (dir_mask == 0 || !common_slot_size.has_value()) { - return {}; - } - - // Ring depth: pl.split(mode, slot_num=N) override, else the PTOAS-matching - // default (8 unidirectional / 4 bidirectional). The reserved buffer and the - // emitted initialize_pipe slot_num attribute both use this value, so PTOAS - // and the auto-reserved buffer stay consistent on a3 (local footprint = - // slot_num when local_slot_num is omitted) and a5 (footprint = slot_num). if (slot_num_override.has_value()) { INTERNAL_CHECK_SPAN(slot_num_override.value() > 0, span) << "Cross-core slot_num override must be positive: " << slot_num_override.value(); } - const int effective_slot_num = slot_num_override.value_or(GetSlotNumForDirMask(dir_mask)); - // Bound-check the slot size before multiplying so an oversized inferred size - // can't overflow the int64 buffer_size computation. - const int64_t slot_size_i64 = common_slot_size.value(); - INTERNAL_CHECK_SPAN(slot_size_i64 >= 0 && slot_size_i64 <= std::numeric_limits::max(), span) - << "Cross-core slot_size out of range: " << slot_size_i64; - const int slot_size_bytes = static_cast(slot_size_i64); - const int64_t buffer_size = slot_size_i64 * effective_slot_num; - AutomaticPipeSetup setup; - std::shared_ptr aic_v2c_reserve_var; - std::shared_ptr aic_c2v_import_var; - std::shared_ptr aiv_c2v_reserve_var; - std::shared_ptr aiv_v2c_import_var; + std::map metadata_by_id; + CollectCrossCorePipeMetadataById(aic_stmts, metadata_by_id); + CollectCrossCorePipeMetadataById(aiv_stmts, metadata_by_id); + + AutomaticPipeSetup setup; auto zero_i32 = [&]() { return std::make_shared(0, DataType::INT32, span); }; auto var_as_expr = [](const std::shared_ptr& v) -> ExprPtr { return std::static_pointer_cast(v); }; - if (dir_mask & core_affinity::kDirMaskV2C) { - const auto v2c_name = BuildPipeBufferName(func_name, core_affinity::PipeDirection::V2C); - auto v2c_reserve = CreateReserveBuffer(v2c_name, buffer_size, span); - aic_v2c_reserve_var = std::make_shared(v2c_name, v2c_reserve->GetType(), span); - setup.aic_stmts.push_back(std::make_shared(aic_v2c_reserve_var, v2c_reserve, span)); - auto v2c_import = CreateImportPeerBuffer(v2c_name, aic_name, span); - aiv_v2c_import_var = std::make_shared(v2c_name + "_import", v2c_import->GetType(), span); - setup.aiv_stmts.push_back(std::make_shared(aiv_v2c_import_var, v2c_import, span)); - } + for (const auto& [pipe_id, pipe_metadata] : metadata_by_id) { + const int dir_mask = BuildDirMask(pipe_metadata); + auto common_slot_size = GetCommonSlotSizeBytes(pipe_metadata); + if (dir_mask == 0 || !common_slot_size.has_value()) { + continue; + } - if (dir_mask & core_affinity::kDirMaskC2V) { - const auto c2v_name = BuildPipeBufferName(func_name, core_affinity::PipeDirection::C2V); - auto c2v_reserve = CreateReserveBuffer(c2v_name, buffer_size, span); - aiv_c2v_reserve_var = std::make_shared(c2v_name, c2v_reserve->GetType(), span); - setup.aiv_stmts.push_back(std::make_shared(aiv_c2v_reserve_var, c2v_reserve, span)); - auto c2v_import = CreateImportPeerBuffer(c2v_name, aiv_name, span); - aic_c2v_import_var = std::make_shared(c2v_name + "_import", c2v_import->GetType(), span); - setup.aic_stmts.push_back(std::make_shared(aic_c2v_import_var, c2v_import, span)); - } + // Ring depth: pl.split(mode, slot_num=N) override, else the PTOAS-matching + // default (8 unidirectional / 4 bidirectional). Split and non-split transfer + // boundaries use distinct logical pipe ids so PTOAS never sees split=0 and + // split=1/2 traffic on the same pipe. + const int effective_slot_num = slot_num_override.value_or(GetSlotNumForDirMask(dir_mask)); + // Bound-check the slot size before multiplying so an oversized inferred size + // can't overflow the int64 buffer_size computation. + const int64_t slot_size_i64 = common_slot_size.value(); + INTERNAL_CHECK_SPAN(slot_size_i64 >= 0 && slot_size_i64 <= std::numeric_limits::max(), span) + << "Cross-core slot_size out of range: " << slot_size_i64; + const int slot_size_bytes = static_cast(slot_size_i64); + const int64_t buffer_size = slot_size_i64 * effective_slot_num; + + std::shared_ptr aic_v2c_reserve_var; + std::shared_ptr aic_c2v_import_var; + std::shared_ptr aiv_c2v_reserve_var; + std::shared_ptr aiv_v2c_import_var; + + if (dir_mask & core_affinity::kDirMaskV2C) { + const auto v2c_name = BuildPipeBufferNameForId(func_name, core_affinity::PipeDirection::V2C, pipe_id); + auto v2c_reserve = CreateReserveBuffer(v2c_name, buffer_size, span); + aic_v2c_reserve_var = std::make_shared(v2c_name, v2c_reserve->GetType(), span); + setup.aic_stmts.push_back(std::make_shared(aic_v2c_reserve_var, v2c_reserve, span)); + auto v2c_import = CreateImportPeerBuffer(v2c_name, aic_name, span); + aiv_v2c_import_var = std::make_shared(v2c_name + "_import", v2c_import->GetType(), span); + setup.aiv_stmts.push_back(std::make_shared(aiv_v2c_import_var, v2c_import, span)); + } - // AIC: c2v operand = import on Cube; v2c operand = reserve on Cube (matches PTO codegen order). - const ExprPtr aic_c2v_arg = aic_c2v_import_var ? var_as_expr(aic_c2v_import_var) : ExprPtr(zero_i32()); - const ExprPtr aic_v2c_arg = aic_v2c_reserve_var ? var_as_expr(aic_v2c_reserve_var) : ExprPtr(zero_i32()); - // AIV: c2v operand = reserve on Vector; v2c operand = import on Vector. - const ExprPtr aiv_c2v_arg = aiv_c2v_reserve_var ? var_as_expr(aiv_c2v_reserve_var) : ExprPtr(zero_i32()); - const ExprPtr aiv_v2c_arg = aiv_v2c_import_var ? var_as_expr(aiv_v2c_import_var) : ExprPtr(zero_i32()); - - setup.aic_stmts.push_back( - std::make_shared(CreateInitializePipe(core_affinity::CoreSide::AIC, dir_mask, slot_size_bytes, - aic_c2v_arg, aic_v2c_arg, slot_num_override, span), - span)); - setup.aiv_stmts.push_back( - std::make_shared(CreateInitializePipe(core_affinity::CoreSide::AIV, dir_mask, slot_size_bytes, - aiv_c2v_arg, aiv_v2c_arg, slot_num_override, span), - span)); + if (dir_mask & core_affinity::kDirMaskC2V) { + const auto c2v_name = BuildPipeBufferNameForId(func_name, core_affinity::PipeDirection::C2V, pipe_id); + auto c2v_reserve = CreateReserveBuffer(c2v_name, buffer_size, span); + aiv_c2v_reserve_var = std::make_shared(c2v_name, c2v_reserve->GetType(), span); + setup.aiv_stmts.push_back(std::make_shared(aiv_c2v_reserve_var, c2v_reserve, span)); + auto c2v_import = CreateImportPeerBuffer(c2v_name, aiv_name, span); + aic_c2v_import_var = std::make_shared(c2v_name + "_import", c2v_import->GetType(), span); + setup.aic_stmts.push_back(std::make_shared(aic_c2v_import_var, c2v_import, span)); + } + + // AIC: c2v operand = import on Cube; v2c operand = reserve on Cube (matches PTO codegen order). + const ExprPtr aic_c2v_arg = aic_c2v_import_var ? var_as_expr(aic_c2v_import_var) : ExprPtr(zero_i32()); + const ExprPtr aic_v2c_arg = aic_v2c_reserve_var ? var_as_expr(aic_v2c_reserve_var) : ExprPtr(zero_i32()); + // AIV: c2v operand = reserve on Vector; v2c operand = import on Vector. + const ExprPtr aiv_c2v_arg = aiv_c2v_reserve_var ? var_as_expr(aiv_c2v_reserve_var) : ExprPtr(zero_i32()); + const ExprPtr aiv_v2c_arg = aiv_v2c_import_var ? var_as_expr(aiv_v2c_import_var) : ExprPtr(zero_i32()); + + setup.aic_stmts.push_back( + std::make_shared(CreateInitializePipe(core_affinity::CoreSide::AIC, dir_mask, slot_size_bytes, + aic_c2v_arg, aic_v2c_arg, slot_num_override, pipe_id, span), + span)); + setup.aiv_stmts.push_back( + std::make_shared(CreateInitializePipe(core_affinity::CoreSide::AIV, dir_mask, slot_size_bytes, + aiv_c2v_arg, aiv_v2c_arg, slot_num_override, pipe_id, span), + span)); + } return setup; } diff --git a/src/ir/transforms/utils/split_axis_utils.cpp b/src/ir/transforms/utils/split_axis_utils.cpp index da63366d00..2a7757cd33 100644 --- a/src/ir/transforms/utils/split_axis_utils.cpp +++ b/src/ir/transforms/utils/split_axis_utils.cpp @@ -161,10 +161,14 @@ bool ValidOperandNeedsLocalize(const ExprPtr& valid_dim, const ExprPtr& original CallPtr RebuildCallWithSplit(const CallPtr& call, int split_int) { std::vector> new_kwargs; bool has_split = false; + bool has_id = false; for (const auto& [key, val] : call->kwargs_) { if (key == "split") { new_kwargs.emplace_back("split", std::any(split_int)); has_split = true; + } else if (key == "id" && split_int != 0) { + new_kwargs.emplace_back("id", std::any(split_int)); + has_id = true; } else { new_kwargs.emplace_back(key, val); } @@ -172,6 +176,9 @@ CallPtr RebuildCallWithSplit(const CallPtr& call, int split_int) { if (!has_split) { new_kwargs.emplace_back("split", std::any(split_int)); } + if (split_int != 0 && !has_id) { + new_kwargs.emplace_back("id", std::any(split_int)); + } return std::make_shared(call->op_, call->args_, std::move(new_kwargs), call->GetType(), call->span_); } @@ -221,10 +228,14 @@ CallPtr RebuildTpopWithHalvedShape(const CallPtr& call, int split_int, int split std::vector> new_kwargs; bool has_split = false; + bool has_id = false; for (const auto& [key, val] : call->kwargs_) { if (key == "split") { new_kwargs.emplace_back("split", std::any(split_int)); has_split = true; + } else if (key == "id" && split_int != 0) { + new_kwargs.emplace_back("id", std::any(split_int)); + has_id = true; } else { new_kwargs.emplace_back(key, val); } @@ -232,6 +243,9 @@ CallPtr RebuildTpopWithHalvedShape(const CallPtr& call, int split_int, int split if (!has_split) { new_kwargs.emplace_back("split", std::any(split_int)); } + if (split_int != 0 && !has_id) { + new_kwargs.emplace_back("id", std::any(split_int)); + } return std::make_shared(call->op_, call->args_, std::move(new_kwargs), new_result_type, call->span_); } diff --git a/tests/ut/ir/transforms/test_resolve_backend_op_layouts_pass.py b/tests/ut/ir/transforms/test_resolve_backend_op_layouts_pass.py index 9f8f4897e0..72066b0bd2 100644 --- a/tests/ut/ir/transforms/test_resolve_backend_op_layouts_pass.py +++ b/tests/ut/ir/transforms/test_resolve_backend_op_layouts_pass.py @@ -308,7 +308,13 @@ def repro( blayout=pl.TileLayout.col_major, slayout=pl.TileLayout.row_major, ) - stored: pl.Tensor[[16, 256], pl.FP32] = pl.store(result, [0, 0], out) + stored_row_major_arg0: pl.Tile[[16, 256], pl.FP32, pl.MemorySpace.Vec] = pl.tile.move( + result, + target_memory=pl.MemorySpace.Vec, + blayout=pl.TileLayout.row_major, + slayout=pl.TileLayout.none_box, + ) + stored: pl.Tensor[[16, 256], pl.FP32] = pl.store(stored_row_major_arg0, [0, 0], out) return stored After = _run_pass(Before) @@ -403,6 +409,66 @@ def repro( After = _run_pass(Before) ir.assert_structural_equal(After, Expected) + def test_tile_store_repairs_col_major_input(self): + """`tile.store` must materialize a row-major Vec tile before PTO TSTORE.""" + + @pl.program + class Before: + @pl.function(type=pl.FunctionType.InCore) + def repro( + self, + out: pl.Out[pl.Tensor[[16, 256], pl.FP32]], + ) -> pl.Tensor[[16, 256], pl.FP32]: + src: pl.Tile[[16, 256], pl.FP32] = pl.tile.create( + [16, 256], dtype=pl.FP32, target_memory=pl.MemorySpace.Vec + ) + col_major: pl.Tile[ + [16, 256], + pl.FP32, + pl.MemorySpace.Vec, + pl.TileView(blayout=pl.TileLayout.col_major, slayout=pl.TileLayout.row_major), + ] = pl.tile.move( + src, + target_memory=pl.MemorySpace.Vec, + blayout=pl.TileLayout.col_major, + slayout=pl.TileLayout.row_major, + ) + stored: pl.Tensor[[16, 256], pl.FP32] = pl.store(col_major, [0, 0], out) + return stored + + @pl.program + class Expected: + @pl.function(type=pl.FunctionType.InCore) + def repro( + self, + out: pl.Out[pl.Tensor[[16, 256], pl.FP32]], + ) -> pl.Tensor[[16, 256], pl.FP32]: + src: pl.Tile[[16, 256], pl.FP32, pl.MemorySpace.Vec] = pl.tile.create( + [16, 256], dtype=pl.FP32, target_memory=pl.MemorySpace.Vec + ) + col_major: pl.Tile[ + [16, 256], + pl.FP32, + pl.MemorySpace.Vec, + pl.TileView(blayout=pl.TileLayout.col_major, slayout=pl.TileLayout.row_major), + ] = pl.tile.move( + src, + target_memory=pl.MemorySpace.Vec, + blayout=pl.TileLayout.col_major, + slayout=pl.TileLayout.row_major, + ) + stored_row_major_arg0: pl.Tile[[16, 256], pl.FP32, pl.MemorySpace.Vec] = pl.tile.move( + col_major, + target_memory=pl.MemorySpace.Vec, + blayout=pl.TileLayout.row_major, + slayout=pl.TileLayout.none_box, + ) + stored: pl.Tensor[[16, 256], pl.FP32] = pl.store(stored_row_major_arg0, [0, 0], out) + return stored + + After = _run_pass(Before) + ir.assert_structural_equal(After, Expected) + def test_unconfigured_backend_is_left_unchanged(self): """With no backend configured the pass is a no-op (`RewriteFunction` guard).