fix(auto-tile): Add M/N tiling for loop-carried matmul_acc

[tonibohnlein]

Summary

Fixes #2232 by adding real M/N tiling for the canonical loop-carried split-K matmul/matmul_acc form. The rewrite moves the output grid outside the K reduction, so each legal L0C sub-tile completes every source K block and is stored before the next sub-tile begins; the oversized full accumulator is never materialized.

The change also:

• uses one backend-aware physical L0C footprint contract in the chooser and allocators, including padded INT32 accumulator rows;
• pads boundary Mat/L0 operand boxes while preserving their logical valid_shape, then propagates that physical/valid distinction through matmul, matmul_acc, and stores;
• represents PTO compact mode in TileView and automatically marks partial Left/Right tile.extract results as compact, so TEXTRACT copies only valid INT8 tail data;
• coalesces padded tile.set_validshape accumulator initializers with the in-place reduction buffer and rejects irreconcilable distinct Acc-to-Acc copies before codegen; and
• retains PH-AT-006 for non-canonical, caller-owned oversized matmul_acc forms that still cannot be safely retiled.

Validation

• cmake --build build --parallel 2
• Focused host integration suite: 1,470 passed after rebasing onto current main (one untouched current-main parser assertion that expects NameError but receives ParserTypeError on local Python 3.14 was excluded)
• Issue-specific host coverage: physical L0C accounting, full M/N grid reconstruction, nested K retiling, symbolic and padded boundaries, allocation/reuse, compact-mode serialization, and PTO codegen
• Ascend 910B2 validation before the final rebase: 57/57 AutoTile system-test nodes passed, including exact INT32 results for [16, 1152] and the padded M/N boundary case [272, 144] under both PyPTO and PTOAS memory planners
• The formerly failing secondary-K boundary row passed repeatedly under both planners and emitted physical Right/Acc N tails of 32 with logical valid N of 16, with no Acc-to-Acc move or full [32, 1152] accumulator

The device run used the feature-equivalent pre-rebase tip d2526fc6; the final rebase preserved the feature patches, integrated the new singular valid_shape API and fractal-byte documentation from main, and was rebuilt and host-revalidated at the PR head.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 99c972f3-2b28-400f-b398-3c024dad0ed2

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

The PR adds TileView.compact support across the IR, Python API, serialization, parsing, and PTO code generation. It also adds backend-aware physical L0C sizing and canonical split-K M/N tiling with updated memory reuse behavior and regression tests.

Changes

Tile metadata and representation

Layer / File(s)	Summary
CompactMode and TileView contract include/pypto/ir/type.h, src/ir/type.cpp, python/bindings/modules/ir.cpp, python/pypto/ir/*, python/pypto/language/*	Adds CompactMode::{null, normal} and exposes it through TileView constructors, properties, Python bindings, and language exports.
Compact propagation and persistence src/ir/op/tile_ops/*, src/ir/serialization/*, src/codegen/pto/*, src/backend/common/*	Preserves compact metadata through transforms, serialization, subview validation, tile-buffer formatting, and same-address move handling.

Physical L0C allocation and tiling

Layer / File(s)	Summary
Backend alignment and footprint contract include/pypto/backend/*, include/pypto/ir/transforms/utils/l0c_footprint.h, include/pypto/ir/transforms/utils/l0_tile_chooser.h, src/ir/transforms/utils/l0_tile_chooser.cpp	Adds backend-specific L0C M-row alignment, boxed layout alignment, and overflow-safe physical footprint calculations.
Physical allocation and memory planning src/ir/transforms/init_memref.cpp, src/ir/transforms/auto_tile_matmul_l0_pass.cpp	Uses physical aligned sizes for MemRef allocation, slices, accumulator footprints, capacity checks, and double-buffer planning.

AutoTileMatmulL0 split-K tiling

Layer / File(s)	Summary
Physical and valid matmul shapes src/ir/op/tile_ops/matmul.cpp, src/ir/op/type_inference.cpp	Separates physical extents from logical valid shapes and validates physical operand compatibility.
Canonical split-K output rewriting src/ir/transforms/auto_tile_matmul_l0_pass.cpp, include/pypto/ir/transforms/passes.h	Recognizes supported split-K chains, clones complete reductions per M/N output tile, narrows loads, tracks valid windows, and preserves store attributes.

Validation and control-flow safety

Layer / File(s)	Summary
Accumulator memory-reuse safety src/ir/transforms/memory_reuse_pass.cpp, tests/ut/ir/transforms/test_memory_reuse.py	Allows shared-storage tile.set_validshape retargeting and rejects divergent Acc control-flow values instead of generating Acc-to-Acc moves.
Regression coverage tests/ut/backend/*, tests/ut/ir/*, tests/ut/language/*, tests/st/runtime/ops/test_auto_tile_matmul.py	Covers compact metadata, backend alignment, physical allocation, chooser capacity, split-K execution, serialization, and parser round trips.

Estimated code review effort: 5 (Critical) | ~120 minutes

Possibly related issues

• #2232: The AutoTileMatmulL0 changes directly address split-K M/N tiling, physical accumulator sizing, and preservation of supported K tiling.
• [Performance] Support valid-shape-preserving packed tile layouts for small Sinkhorn reductions #1958: The compact mode and valid-shape handling support packed partial-tile representations.

Possibly related PRs

• hw-native-sys/pypto#2070: Directly overlaps AutoTileMatmulL0 M/N tiling, split-K handling, and L0C capacity tests.
• hw-native-sys/pypto#2100: Directly relates to preserving tile moves when source and destination representations differ.
• hw-native-sys/pypto#2178: Directly relates to AutoTileMatmulL0 pipeline capacity and L0C planning.

Suggested labels: enhancement

Poem

I carry compact tiles through the code,
With padded rows in a measured load.
Split-K blooms across each M/N square,
While safe Acc moves stop in midair.
The rabbit checks each shape and byte—
Then hops through tests by morning light.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 46.83% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the primary change: adding M/N tiling for loop-carried matmul_acc.
Description check	✅ Passed	The description directly explains the M/N tiling fix, supporting changes, linked issue, and validation.
Linked Issues check	✅ Passed	The changes implement M/N tiling, preserve K tiling, add diagnostics, and provide coverage for the linked issue requirements [#2232].
Out of Scope Changes check	✅ Passed	The implementation, API updates, documentation, and tests support the stated matmul_acc tiling and physical-layout objectives.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 3

🧹 Nitpick comments (4)

src/ir/transforms/init_memref.cpp (1)

803-813: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Remove the shadowing ctx re-declaration at line 822.

This change introduces const auto* ctx = PassContext::Current(); at line 803. Line 822 declares a second ctx with the same initializer inside the declared-allocation check. The inner declaration shadows the outer one and adds no behavior. Reuse the outer ctx.

♻️ Proposed cleanup at line 822

   if (!declared_allocs.empty()) {
-    const auto* ctx = PassContext::Current();
     CHECK(ctx == nullptr || ctx->GetMemoryPlanner() != MemoryPlanner::PtoAS)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/ir/transforms/init_memref.cpp` around lines 803 - 813, Remove the inner
ctx declaration in the declared-allocation check and reuse the existing ctx
initialized from PassContext::Current(). Keep the surrounding
DeclaredAllocCollector initialization and backend-handler behavior unchanged.

src/ir/transforms/utils/l0_tile_chooser.cpp (1)

538-542: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Report the physical footprint in the L0c capacity failure message.

The message prints the logical minimum tile min_m x min_n. With box_align_* or l0c_align_m set, the rejected footprint is *min_c_elements, which can be much larger. Include that value so a caller can see why the logical minimum does not fit.

♻️ Proposed message change

   CHECK(min_c_elements && *min_c_elements <= static_cast<uint64_t>(std::numeric_limits<int64_t>::max()) &&
         C0_base >= static_cast<int64_t>(*min_c_elements))
       << "ChooseL0Tile: L0c capacity " << C0_base << " elements is too small to fit the minimum tile ("
-      << cfg.min_m << " x " << cfg.min_n << ")";
+      << cfg.min_m << " x " << cfg.min_n << "), whose physical footprint after box/L0C-row alignment is "
+      << (min_c_elements ? std::to_string(*min_c_elements) : std::string("unrepresentable")) << " elements";

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/ir/transforms/utils/l0_tile_chooser.cpp` around lines 538 - 542, Update
the failure message in the ChooseL0Tile validation to include the physical
minimum L0c footprint from min_c_elements, alongside the existing logical min_m
x min_n dimensions, so aligned configurations report the actual rejected
capacity requirement.

src/ir/transforms/auto_tile_matmul_l0_pass.cpp (2)

1486-1496: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Hoist the loop-invariant out_ty lookup out of the grid loop.

out_ty at line 1490 depends only on match.matmul, so it does not change across grid steps. Compute it once before the nj loop.

♻️ Proposed change

+  auto out_ty = As<TileType>(match.matmul->var_->GetType());
+  INTERNAL_CHECK_SPAN(out_ty, match.matmul->span_)
+      << "Internal error: canonical split-K matmul result is not a TileType";
   int step = 0;
   for (int64_t nj = 0; nj < num_n; ++nj) {
@@
       const std::string suffix = "_mn" + std::to_string(step);
-      auto out_ty = As<TileType>(match.matmul->var_->GetType());
       const auto window = BuildCanonicalOutputWindow(match, m_eff, n_eff, *output_box_alignment);

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/ir/transforms/auto_tile_matmul_l0_pass.cpp` around lines 1486 - 1496,
Move the loop-invariant TileType lookup for out_ty out of the nested grid loops,
computing it once before the nj loop while keeping the existing match.matmul
source and all subsequent uses unchanged.

---

1017-1020: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Reuse PreserveCallAttrs instead of duplicating the attribute rebuild.

Lines 1018-1020 rebuild the deduced Call with the original attributes. PreserveCallAttrs at line 1340 performs the identical rebuild. Move PreserveCallAttrs above DirectGmPlacer and call it here, so the attribute-preservation rule has one definition.

♻️ Proposed refactor

     auto deduced = reg.Create("tile.store", {sub, offs, chain_in}, kwargs_, sp_);
-    auto scall = attrs_.empty() ? deduced
-                                : std::make_shared<Call>(deduced->op_, deduced->args_, deduced->kwargs_,
-                                                         attrs_, deduced->GetType(), deduced->span_);
+    auto scall = attrs_.empty()
+                     ? deduced
+                     : std::make_shared<Call>(deduced->op_, deduced->args_, deduced->kwargs_, attrs_,
+                                              deduced->GetType(), deduced->span_);

Then replace the body with a shared helper once PreserveCallAttrs is declared before this class.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/ir/transforms/auto_tile_matmul_l0_pass.cpp` around lines 1017 - 1020,
Extract or move PreserveCallAttrs before DirectGmPlacer so it is available at
this call site, then replace the inline attrs_.empty() conditional and Call
reconstruction in the tile.store creation path with that helper. Remove the
duplicate attribute-preservation implementation while keeping the resulting Call
attributes unchanged.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@python/pypto/ir/builder.py`:
- Around line 692-693: In the ir.TileView signature, move the compact parameter
after span so existing eighth-position positional calls continue passing ir.Span
to span. Preserve the current compact default and all other parameter behavior.

In `@src/backend/common/pto_ops_elementwise.cpp`:
- Around line 570-572: Update the same_layout compatibility check in the
elementwise operation to also require src_view.pad == dst_view.pad before taking
the no-op aliasing path. Preserve the existing tmov behavior for any layout or
padding mismatch.

In `@src/ir/transforms/memory_reuse_pass.cpp`:
- Around line 602-622: Add a regression test in
tests/ut/ir/transforms/test_memory_reuse.py covering an Acc initializer whose
seed passes through tile.set_validshape with padding. Assert the initializer and
accumulator coalesce to the same Acc allocation, while the resulting logical
valid_shape remains unchanged.

---

Nitpick comments:
In `@src/ir/transforms/auto_tile_matmul_l0_pass.cpp`:
- Around line 1486-1496: Move the loop-invariant TileType lookup for out_ty out
of the nested grid loops, computing it once before the nj loop while keeping the
existing match.matmul source and all subsequent uses unchanged.
- Around line 1017-1020: Extract or move PreserveCallAttrs before DirectGmPlacer
so it is available at this call site, then replace the inline attrs_.empty()
conditional and Call reconstruction in the tile.store creation path with that
helper. Remove the duplicate attribute-preservation implementation while keeping
the resulting Call attributes unchanged.

In `@src/ir/transforms/init_memref.cpp`:
- Around line 803-813: Remove the inner ctx declaration in the
declared-allocation check and reuse the existing ctx initialized from
PassContext::Current(). Keep the surrounding DeclaredAllocCollector
initialization and backend-handler behavior unchanged.

In `@src/ir/transforms/utils/l0_tile_chooser.cpp`:
- Around line 538-542: Update the failure message in the ChooseL0Tile validation
to include the physical minimum L0c footprint from min_c_elements, alongside the
existing logical min_m x min_n dimensions, so aligned configurations report the
actual rejected capacity requirement.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 987a1dd7-4946-4742-8775-6b68222d0107

📥 Commits

Reviewing files that changed from the base of the PR and between 6730e96 and 1c7e416.

📒 Files selected for processing (71)

• docs/en/dev/backend/00-backend_handler.md
• docs/en/dev/codegen/00-pto_codegen.md
• docs/en/dev/ir/02-types.md
• docs/en/dev/ir/04-serialization.md
• docs/en/dev/language/00-python_syntax.md
• docs/en/dev/passes/15-auto_tile_matmul_l0.md
• docs/en/dev/passes/31-memory_reuse.md
• docs/zh/dev/backend/00-backend_handler.md
• docs/zh/dev/codegen/00-pto_codegen.md
• docs/zh/dev/ir/02-types.md
• docs/zh/dev/ir/04-serialization.md
• docs/zh/dev/language/00-python_syntax.md
• docs/zh/dev/passes/15-auto_tile_matmul_l0.md
• docs/zh/dev/passes/31-memory_reuse.md
• include/pypto/backend/910B/backend_910b_handler.h
• include/pypto/backend/common/backend_handler.h
• include/pypto/codegen/pto/pto_type_utils.h
• include/pypto/ir/tile_view_semantics.h
• include/pypto/ir/transforms/passes.h
• include/pypto/ir/transforms/utils/l0_tile_chooser.h
• include/pypto/ir/transforms/utils/l0c_footprint.h
• include/pypto/ir/transforms/utils/memref_utils.h
• include/pypto/ir/transforms/utils/tile_buf_signature.h
• include/pypto/ir/type.h
• include/pypto/ir/type_inference.h
• python/bindings/modules/backend.cpp
• python/bindings/modules/ir.cpp
• python/bindings/modules/passes.cpp
• python/pypto/ir/builder.py
• python/pypto/ir/type.py
• python/pypto/language/__init__.py
• python/pypto/language/parser/ast_parser.py
• python/pypto/language/parser/type_resolver.py
• python/pypto/pypto_core/backend.pyi
• python/pypto/pypto_core/ir.pyi
• python/pypto/pypto_core/passes.pyi
• src/backend/common/pto_ops_datamove.cpp
• src/backend/common/pto_ops_elementwise.cpp
• src/backend/common/pto_ops_shared.cpp
• src/codegen/pto/pto_codegen.cpp
• src/codegen/pto/pto_type_utils.cpp
• src/ir/op/tile_ops/matmul.cpp
• src/ir/op/tile_ops/transform.cpp
• src/ir/op/type_inference.cpp
• src/ir/serialization/deserializer.cpp
• src/ir/serialization/serializer.cpp
• src/ir/transforms/auto_tile_matmul_l0_pass.cpp
• src/ir/transforms/flatten_tile_nd_to_2d/rewrite.cpp
• src/ir/transforms/infer_tile_memory_space_pass.cpp
• src/ir/transforms/init_memref.cpp
• src/ir/transforms/memory_reuse_pass.cpp
• src/ir/transforms/python_printer.cpp
• src/ir/transforms/structural_equal.cpp
• src/ir/transforms/structural_hash.cpp
• src/ir/transforms/utils/l0_tile_chooser.cpp
• src/ir/type.cpp
• tests/st/runtime/ops/test_auto_tile_matmul.py
• tests/ut/backend/test_backend_910b.py
• tests/ut/backend/test_backend_950.py
• tests/ut/codegen/test_pto_codegen_ops.py
• tests/ut/ir/core/test_tile_view_equality.py
• tests/ut/ir/memory/test_memref.py
• tests/ut/ir/operators/test_tile_ops.py
• tests/ut/ir/transforms/test_allocate_memory_addr_pass.py
• tests/ut/ir/transforms/test_auto_tile_matmul_acc_mn.py
• tests/ut/ir/transforms/test_auto_tile_matmul_l0.py
• tests/ut/ir/transforms/test_init_memref.py
• tests/ut/ir/transforms/test_l0_tile_chooser.py
• tests/ut/ir/transforms/test_memory_reuse.py
• tests/ut/ir/transforms/test_serialization.py
• tests/ut/language/parser/test_type_resolver.py