feat(codegen): add use_ptoas_multi_buffer switch for ptoas multi-buffer

CMakeLists.txt

@@ -153,6 +153,7 @@ set(PYPTO_SOURCES
     src/ir/transforms/canonicalize_tile_slice_pass.cpp
     src/ir/transforms/fold_no_op_reshape_pass.cpp
     src/ir/transforms/stamp_tfree_split_pass.cpp
+    src/ir/transforms/convert_to_ptoas_multi_buffer_pass.cpp
     src/ir/transforms/materialize_runtime_scopes_pass.cpp
     src/ir/transforms/fuse_create_assemble_to_slice_pass.cpp
     src/ir/transforms/memory_reuse_pass.cpp

docs/en/dev/passes/00-pass_manager.md

@@ -388,6 +388,7 @@ The PTO-oriented tile stage shared by `Default` and `DebugTileOptimization` is:
 11. [`StampTfreeSplit`](22-stamp_tfree_split.md) (copies each cross-core tpop's split/pipe-id onto its matching tfree op)
 12. `NormalizeReturnOrder`
 13. [`SkewCrossCorePipeline`](24-skew_cross_core_pipeline.md) (cross-core cube/vector software-pipeline skew; runs immediately before LowerPipelineLoops)
+    - `ConvertToPtoasMultiBuffer` — **optional, gated by `PassContext.use_ptoas_multi_buffer`** (no-op by default; no dedicated doc). Runs immediately after `SkewCrossCorePipeline`. When enabled, converts same-core `pl.pipeline` ping-pong loops to ptoas multi-buffer slots: tags vec/mat tile-producing `Call`s with `kMultiBufferSlotsAttr = N` and downgrades `pipeline_stages` to 1 so `LowerPipelineLoops` skips body replication. Codegen then emits `pto.alloc_multi_tile` / `pto.multi_tile_get[iv % N]` (with `AllocateMemoryAddr` reserving N contiguous slots and `MemoryReuse` keeping the tile exclusive), delegating slot rotation + cross-iteration sync to ptoas instead of pypto's body-replication ping-pong.
 14. [`LowerPipelineLoops`](25-lower_pipeline_loops.md)
 15. [`CanonicalizeIOOrder`](26-canonicalize_io_order.md)
 16. [`MaterializeTensorStrides`](27-materialize_tensor_strides.md) — wired into the default pipeline starting from RFC #1300 P6

docs/zh-cn/dev/passes/00-pass_manager.md

@@ -388,6 +388,7 @@ with passes.PassContext([passes.VerificationInstrument(passes.VerificationMode.A
 11. [`StampTfreeSplit`](22-stamp_tfree_split.md)（把每个跨核 tpop 的 split/pipe-id 复制到与之配对的 tfree 算子上）
 12. `NormalizeReturnOrder`
 13. [`SkewCrossCorePipeline`](24-skew_cross_core_pipeline.md)（cube/vector 跨核软流水 skew；紧接在 LowerPipelineLoops 之前运行）
+    - `ConvertToPtoasMultiBuffer` —— **可选，由 `PassContext.use_ptoas_multi_buffer` 门控**（默认 no-op；无独立文档）。紧接在 `SkewCrossCorePipeline` 之后运行。开启后把同核 `pl.pipeline` ping-pong 循环转换为 ptoas 多缓冲槽位：给 vec/mat 产生 tile 的 `Call` 打上 `kMultiBufferSlotsAttr = N` 并把 `pipeline_stages` 降到 1，使 `LowerPipelineLoops` 跳过循环体复制。codegen 随后发射 `pto.alloc_multi_tile` / `pto.multi_tile_get[iv % N]`（`AllocateMemoryAddr` 预留 N 个连续槽位、`MemoryReuse` 保持该 tile 独占），把槽位轮转与跨迭代同步交给 ptoas，而非 pypto 的循环体复制 ping-pong。
 14. [`LowerPipelineLoops`](25-lower_pipeline_loops.md)
 15. [`CanonicalizeIOOrder`](26-canonicalize_io_order.md)
 16. [`MaterializeTensorStrides`](27-materialize_tensor_strides.md) —— 自 RFC #1300 P6 起接入默认 pipeline

include/pypto/codegen/pto/pto_codegen.h

@@ -337,6 +337,20 @@ class PTOCodegen : public CodegenBase {
    */
   void EmitAllocTileForVar(const ir::VarPtr& tile_var, const std::shared_ptr<const ir::TileType>& tile_type);
 
+  /**
+   * @brief Hoist `pto.alloc_multi_tile` for ptoas multi-buffer tiles in a loop.
+   *
+   * Pre-scans `body` for tile-producing `Call`s tagged with
+   * `kMultiBufferSlotsAttr` (set by `ConvertToPtoasMultiBuffer`), stopping at
+   * nested `ForStmt` boundaries so each rotating tile is hoisted at its own
+   * enclosing loop. For each, emits one `pto.alloc_multi_tile addr=<base>
+   * count=N` (using the tile's MemRef byte offset as the level3 base) BEFORE the
+   * `scf.for` line and records the slot in `fs_.multi_buffer_tiles`. Returns the
+   * number of tiles hoisted (0 ⇒ ordinary loop, no multi-buffer state set).
+   * `EmitAllocTileForVar` then emits `pto.multi_tile_get %mb[iv % N]` per use.
+   */
+  int PrescanAndHoistMultiBuffer(const ir::StmtPtr& body, const std::string& loop_var_ssa);
+
   /**
    * @brief Resolve the DPS element vars of a tuple-returning op call
    *
@@ -746,6 +760,21 @@ class PTOCodegen : public CodegenBase {
     /// vars share one handle (PTOAS in-place aliasing).
     std::set<std::string> emitted_tile_alloc_names;
 
+    /// ptoas multi-buffer lowering (kMultiBufferSlotsAttr). Populated per
+    /// pipeline loop by the pre-scan in `VisitStmt_(ForStmtPtr)`: the hoisted
+    /// `pto.alloc_multi_tile` SSA + types for each rotating tile var, keyed by
+    /// tile Var. `EmitAllocTileForVar` reads this to emit a `pto.multi_tile_get`
+    /// slot pick (indexed by `mb_loop_var_ssa % N`) instead of `pto.alloc_tile`.
+    struct MultiBufferSlot {
+      std::string mb_ssa;         ///< hoisted alloc_multi_tile result SSA
+      std::string mtb_type_str;   ///< `!pto.multi_tile_buf<...>` type string
+      std::string tile_type_str;  ///< per-slot `!pto.tile_buf<...>` type string
+      int slots = 0;              ///< N
+    };
+    std::map<const ir::Var*, MultiBufferSlot> multi_buffer_tiles;
+    std::string mb_loop_var_ssa;               ///< current pipeline loop iv SSA (rotation index source)
+    std::map<int, std::string> mb_remui_by_n;  ///< N → cached `iv % N` index SSA for the current loop
+
     ir::FunctionPtr current_function;
     ir::VarPtr current_result_var;
     std::string current_result_buf;
@@ -806,6 +835,10 @@ class PTOCodegen : public CodegenBase {
       memref_identity_to_mlir.clear();
       emitted_tile_alloc_names.clear();
 
+      multi_buffer_tiles.clear();
+      mb_loop_var_ssa.clear();
+      mb_remui_by_n.clear();
+
       current_function.reset();
       current_result_var.reset();
       current_result_buf.clear();

include/pypto/ir/transforms/pass_context.h

@@ -262,12 +262,17 @@ class PassContext {
    * @param memory_planner Who plans on-chip buffer memory (default: PyPTO).
    *        PtoAS makes the pipeline skip the pypto allocation passes so the
    *        ptoas PlanMemory pass owns allocation instead.
+   * @param use_ptoas_multi_buffer When true, `ConvertToPtoasMultiBuffer`
+   *        rewrites `pl.pipeline` double/multi-buffered loops to emit ptoas
+   *        `pto.alloc_multi_tile` / `pto.multi_tile_get` slot rotation instead
+   *        of pypto's own body-replication ping-pong (default: false).
    */
   explicit PassContext(std::vector<PassInstrumentPtr> instruments,
                        VerificationLevel verification_level = VerificationLevel::Basic,
                        DiagnosticPhase diagnostic_phase = DiagnosticPhase::PrePipeline,
                        DiagnosticCheckSet disabled_diagnostics = {DiagnosticCheck::UnusedControlFlowResult},
-                       MemoryPlanner memory_planner = MemoryPlanner::PyPTO);
+                       MemoryPlanner memory_planner = MemoryPlanner::PyPTO,
+                       bool use_ptoas_multi_buffer = false);
 
   /**
    * @brief Push this context onto the thread-local stack
@@ -315,6 +320,16 @@ class PassContext {
    */
   [[nodiscard]] const DiagnosticCheckSet& GetDisabledDiagnostics() const;
 
+  /**
+   * @brief Whether ptoas multi-buffer lowering is enabled for this context.
+   *
+   * When true, `ConvertToPtoasMultiBuffer` converts `pl.pipeline`
+   * double/multi-buffered loops into ptoas `pto.alloc_multi_tile` /
+   * `pto.multi_tile_get` slot rotation (delegating ping-pong + cross-iteration
+   * sync to ptoas) instead of pypto's body-replication ping-pong.
+   */
+  [[nodiscard]] bool UsePtoasMultiBuffer() const;
+
   /**
    * @brief Get the instruments registered on this context
    */
@@ -350,6 +365,7 @@ class PassContext {
   DiagnosticPhase diagnostic_phase_;
   DiagnosticCheckSet disabled_diagnostics_;
   MemoryPlanner memory_planner_;
+  bool use_ptoas_multi_buffer_;
   PassContext* previous_;
 
   static thread_local PassContext* current_;

include/pypto/ir/transforms/pass_properties.h

@@ -333,6 +333,19 @@ inline const PassProperties kFoldNoOpReshapeProperties{
 
 inline const PassProperties kStampTfreeSplitProperties{.required = {IRProperty::SplitIncoreOrch}};
 
+// -- Convert pl.pipeline ping-pong to ptoas multi-buffer ---------------------
+//
+// Gated by PassContext::UsePtoasMultiBuffer(); a no-op otherwise. Runs right
+// before LowerPipelineLoops (so only same-core Pipeline loops remain — cross-
+// core ones were demoted by SkewCrossCorePipeline). Tags vec/mat tile-producing
+// Calls with kMultiBufferSlotsAttr and downgrades the loop's pipeline_stages to
+// 1 (suppressing body replication). Preserves all pipeline-stage properties.
+inline const PassProperties kConvertToPtoasMultiBufferProperties{
+    .required = {IRProperty::SSAForm, IRProperty::SplitIncoreOrch, IRProperty::IncoreTileOps,
+                 IRProperty::TileOps2D, IRProperty::TileMemoryInferred, IRProperty::NormalizedStmtStructure},
+    .produced = {IRProperty::SSAForm, IRProperty::SplitIncoreOrch, IRProperty::IncoreTileOps,
+                 IRProperty::TileOps2D, IRProperty::TileMemoryInferred, IRProperty::NormalizedStmtStructure}};
+
 }  // namespace pass
 }  // namespace ir
 }  // namespace pypto

include/pypto/ir/transforms/passes.h

@@ -776,6 +776,24 @@ Pass MaterializeRuntimeScopes();
  */
 Pass StampTfreeSplit();
 
+/**
+ * @brief Convert `pl.pipeline` ping-pong loops to ptoas multi-buffer slots
+ *
+ * Gated by `PassContext::UsePtoasMultiBuffer()` — a no-op when the switch is
+ * off, so the default pipeline is unchanged. When on, for each same-core
+ * `ForKind::Pipeline` loop with `pipeline_stages > 1` it:
+ *   - tags every vec/mat tile-producing `Call` in the loop body with
+ *     `kMultiBufferSlotsAttr = N` (the slot count), and
+ *   - downgrades the loop's `pipeline_stages` to 1 so `LowerPipelineLoops`
+ *     skips body replication.
+ *
+ * Codegen then emits one hoisted `pto.alloc_multi_tile addr=<base> count=N`
+ * per tagged tile plus `pto.multi_tile_get %mb[iv % N]` at each use, delegating
+ * slot rotation and cross-iteration synchronization to ptoas instead of pypto's
+ * own body-replication ping-pong. Runs immediately before `LowerPipelineLoops`.
+ */
+Pass ConvertToPtoasMultiBuffer();
+
 /**
  * @brief Verify properties on a program and throw on errors
  *

include/pypto/ir/transforms/utils/attrs.h

@@ -144,6 +144,24 @@ inline bool PipelineMembershipsConflict(const std::vector<std::pair<int32_t, int
   return false;
 }
 
+/// Attribute key on a tile-producing ``Call`` marking that the tile it defines
+/// should be lowered to a **ptoas multi-buffer slot** rather than an ordinary
+/// per-iteration ``pto.alloc_tile``. The value is the slot count ``N`` (``int``,
+/// ``>= 2``) — the enclosing ``pl.pipeline`` loop's ``pipeline_stages``.
+///
+/// Set by ``ConvertToPtoasMultiBuffer`` (gated by
+/// ``PassContext::UsePtoasMultiBuffer()``), which also downgrades the enclosing
+/// loop's ``pipeline_stages`` to 1 so ``LowerPipelineLoops`` skips body
+/// replication — the ping-pong is delegated to ptoas instead. Consumed by
+/// ``PTOCodegen``: it hoists one ``pto.alloc_multi_tile addr=<base> count=N``
+/// before the loop and emits ``pto.multi_tile_get %mb[iv % N]`` at each use,
+/// leaving cross-iteration slot rotation + synchronization to ptoas.
+///
+/// Only tiles in ``vec`` / ``mat`` local memory are tagged (ptoas multi-buffer
+/// supports only those spaces); L0 matmul spaces (Left/Right/Acc/Bias) are
+/// matmul-managed and left as ordinary allocs.
+inline constexpr const char* kMultiBufferSlotsAttr = "multi_buffer_slots";
+
 /// Return a copy of `attrs` with any entry matching `key` removed. The order of
 /// the remaining entries is preserved.
 inline std::vector<std::pair<std::string, std::any>> StripAttr(

python/bindings/modules/passes.cpp

@@ -277,13 +277,14 @@ void BindPass(nb::module_& m) {
                           "verification and the diagnostic channel (warnings + performance\n"
                           "hints) for PassPipeline.")
       .def(nb::init<std::vector<PassInstrumentPtr>, VerificationLevel, DiagnosticPhase, DiagnosticCheckSet,
-                    MemoryPlanner>(),
+                    MemoryPlanner, bool>(),
            nb::arg("instruments"), nb::arg("verification_level") = VerificationLevel::Basic,
            nb::arg("diagnostic_phase") = DiagnosticPhase::PrePipeline,
            nb::arg("disabled_diagnostics") = DiagnosticCheckSet{DiagnosticCheck::UnusedControlFlowResult},
-           nb::arg("memory_planner") = MemoryPlanner::PyPTO,
+           nb::arg("memory_planner") = MemoryPlanner::PyPTO, nb::arg("use_ptoas_multi_buffer") = false,
            "Create a PassContext with instruments, verification level, diagnostic phase gate, "
-           "optional disabled diagnostic checks, and memory planner selection")
+           "optional disabled diagnostic checks, memory planner selection, and optional ptoas "
+           "multi-buffer lowering")
       .def("__enter__",
            [](PassContext& self) -> PassContext& {
              self.EnterContext();
@@ -296,6 +297,8 @@ void BindPass(nb::module_& m) {
            "Get the diagnostic phase gate for this context")
       .def("get_disabled_diagnostics", &PassContext::GetDisabledDiagnostics,
            "Get the diagnostic checks suppressed by this context")
+      .def("use_ptoas_multi_buffer", &PassContext::UsePtoasMultiBuffer,
+           "Whether ptoas multi-buffer lowering is enabled for this context")
       .def("get_instruments", &PassContext::GetInstruments, "Get the instruments registered on this context")
       .def("get_memory_planner", &PassContext::GetMemoryPlanner,
            "Get the memory planner selection for this context")
@@ -527,6 +530,11 @@ void BindPass(nb::module_& m) {
   passes.def("stamp_tfree_split", &pass::StampTfreeSplit,
              "Copy each cross-core tpop's split/pipe-id onto its matching tfree op so codegen\n"
              "reads them from the op directly. Covers mixed-kernel and explicit AIC/AIV tfrees.");
+  passes.def("convert_to_ptoas_multi_buffer", &pass::ConvertToPtoasMultiBuffer,
+             "Convert pl.pipeline ping-pong loops to ptoas multi-buffer slots (gated by\n"
+             "PassContext.use_ptoas_multi_buffer). Tags vec/mat tile Calls with the slot count\n"
+             "and downgrades pipeline_stages to 1 so codegen emits pto.alloc_multi_tile /\n"
+             "pto.multi_tile_get and ptoas owns slot rotation + sync. No-op when the switch is off.");
   passes.def("materialize_runtime_scopes", &pass::MaterializeRuntimeScopes,
              "Materialize implicit orchestration scopes as explicit RuntimeScopeStmt nodes.\n\n"
              "For every Orchestration function, inserts AUTO RuntimeScopeStmt (manual_=false)\n"

python/pypto/ir/compile.py

@@ -53,7 +53,7 @@ def _backend_type_for_platform(platform: str | None, fallback: BackendType) -> B
     raise ValueError(f"Invalid platform {platform!r}. Expected 'a2a3sim', 'a2a3', 'a5sim', or 'a5'.")
 
 
-def compile(  # noqa: PLR0913
+def compile(  # noqa: PLR0913, PLR0912
     program: _ir_core.Program,
     output_dir: str | None = None,
     strategy: OptimizationStrategy = OptimizationStrategy.Default,
@@ -69,6 +69,7 @@ def compile(  # noqa: PLR0913
     distributed_config: Any = None,
     block_dim: int | None = None,
     analyze_auto_scopes_for_deps: bool = False,
+    use_ptoas_multi_buffer: bool | None = None,
 ) -> "CompiledProgram | DistributedCompiledProgram":
     """Compile a Program through passes and codegen.
 
@@ -133,6 +134,12 @@ def compile(  # noqa: PLR0913
             behavior unless explicitly enabled. User-written manual scopes do
             not get compiler deps, but covered read-only inputs may still be
             rewritten to NoDep.
+        use_ptoas_multi_buffer: When True, ``ConvertToPtoasMultiBuffer`` lowers
+            ``pl.pipeline`` double/multi-buffered loops to ptoas
+            ``pto.alloc_multi_tile`` / ``pto.multi_tile_get`` slot rotation
+            (ptoas owns ping-pong + cross-iteration sync) instead of pypto's
+            body-replication ping-pong. None uses the default (off, or the
+            ``PYPTO_PTOAS_MULTI_BUFFER`` env var).
 
     Returns:
         A :class:`CompiledProgram` that wraps the output directory and can
@@ -176,6 +183,11 @@ def compile(  # noqa: PLR0913
             "compile() was called with memory_planner while a PassContext is already active. "
             "Set the memory planner on the existing PassContext instead."
         )
+    if use_ptoas_multi_buffer is not None and outer is not None:
+        raise RuntimeError(
+            "compile() was called with use_ptoas_multi_buffer while a PassContext is already active. "
+            "Set use_ptoas_multi_buffer on the existing PassContext instead."
+        )
 
     # --- Compile profiling ---------------------------------------------------
     prof = get_active_profiler()
@@ -209,7 +221,19 @@ def compile(  # noqa: PLR0913
         dphase = diagnostic_phase if diagnostic_phase is not None else _passes.get_default_diagnostic_phase()
         disabled = disabled_diagnostics if disabled_diagnostics is not None else default_disabled
         mplan = memory_planner if memory_planner is not None else _passes.MemoryPlanner.PYPTO
-    ctx = _passes.PassContext(instruments, vlevel, dphase, disabled, mplan)
+
+    # Resolve ptoas multi-buffer lowering. The explicit ``use_ptoas_multi_buffer``
+    # argument (from compile() / RunConfig) is the reliable source of truth; the
+    # PYPTO_PTOAS_MULTI_BUFFER env var is a global fallback that applies whether
+    # or not an outer PassContext is active. An outer context can only force it
+    # *on* (its default is off), so it never suppresses the env var.
+    env_mbuf = os.environ.get("PYPTO_PTOAS_MULTI_BUFFER", "0").strip().lower() in ("1", "true", "yes", "on")
+    outer_mbuf = outer.use_ptoas_multi_buffer() if outer is not None else False
+    if use_ptoas_multi_buffer is not None:
+        mbuf = use_ptoas_multi_buffer
+    else:
+        mbuf = outer_mbuf or env_mbuf
+    ctx = _passes.PassContext(instruments, vlevel, dphase, disabled, mplan, mbuf)
 
     if mplan == _passes.MemoryPlanner.PTOAS:
         logger.warning(

python/pypto/ir/pass_manager.py

@@ -162,6 +162,13 @@ def _register_passes(cls):
             ("StampTfreeSplit", lambda: passes.stamp_tfree_split()),
             ("NormalizeReturnOrder", lambda: passes.normalize_return_order()),
             ("SkewCrossCorePipeline", lambda: passes.skew_cross_core_pipeline()),
+            # Gated by PassContext.use_ptoas_multi_buffer (no-op otherwise). Runs
+            # right before LowerPipelineLoops so only same-core Pipeline loops
+            # remain (cross-core ones were demoted by SkewCrossCorePipeline). Tags
+            # vec/mat tile Calls with the slot count and downgrades pipeline_stages
+            # to 1 so codegen emits ptoas multi-buffer ops instead of pypto
+            # replicating the body for ping-pong.
+            ("ConvertToPtoasMultiBuffer", lambda: passes.convert_to_ptoas_multi_buffer()),
             ("LowerPipelineLoops", lambda: passes.lower_pipeline_loops()),
             ("CanonicalizeIOOrder", lambda: passes.canonicalize_io_order()),
             # MaterializeTensorStrides fills empty stride slots on every
@@ -407,8 +414,17 @@ def after_pass(_pass_obj: passes.Pass, program: core_ir.Program) -> None:
             inner_phase = passes.DiagnosticPhase.PRE_PIPELINE
         else:
             inner_phase = outer_phase
-
-        with passes.PassContext([*outer_instruments, *extra_instruments], level, inner_phase, disabled):
+        # Preserve ALL outer-context pass config here — this reconstructs the
+        # context the pipeline runs under, so any PassContext setting the passes
+        # read (e.g. memory_planner; use_ptoas_multi_buffer gating
+        # ConvertToPtoasMultiBuffer) must be carried over, or dump mode silently
+        # changes pass behavior.
+        mplan = ctx.get_memory_planner() if ctx else passes.MemoryPlanner.PYPTO
+        mbuf = ctx.use_ptoas_multi_buffer() if ctx else False
+
+        with passes.PassContext(
+            [*outer_instruments, *extra_instruments], level, inner_phase, disabled, mplan, mbuf
+        ):
             try:
                 return self._pipeline.run(input_ir)
             finally: