fix(printer): print atomic kwarg as pl.AtomicType enum, not raw int

[Hzfengsy]

Summary

• The Python printer emitted pl.tile.store(..., atomic=1) for stores carrying the atomic-add combine mode, but the DSL signature is atomic: AtomicType. Pyright flagged the printed source as Literal[1] incompatible with AtomicType.
• Restore the enum form in the printer (pl.AtomicType.Add / pl.AtomicType.None_), mirroring the existing set_pipe / wait_pipe → pl.PipeType.<Name> handling. Call::kwargs_ storage stays int (consistent with the binding layer that casts all enums → int).
• Round-trip preserved: pl.AtomicType is already exposed in pypto.language, so the parser resolves pl.AtomicType.Add via normal Python attribute lookup. Covers tile.store, tensor.assemble, and pld.tensor.put (all share the printer path).
• Also fix two pre-existing misc-include-cleaner violations in src/ir/op/distributed/put.cpp (missing direct includes for <cstddef> / pypto/core/logging.h) that were surfaced by adding <string> / pypto/core/error.h to comm.h.

Before / After

tile.store with atomic-add combine mode in a dumped IR file:

# Before
out__tile = pl.tile.store(t__tile, [0, 0], out__ssa_v0, atomic=1)
#                                                       ^^^^^^^^ Literal[1] is not AtomicType (Pyright)

# After
out__tile = pl.tile.store(t__tile, [0, 0], out__ssa_v0, atomic=pl.AtomicType.Add)

Test plan

• tests/ut/ir/printing/test_python_printer.py::test_python_print_atomic_kwarg_uses_enum_form (new) — pins both Add and None_ print forms
• tests/ut/ir/parser/test_put_op.py::test_put_round_trips_through_printer_and_parser — existing round-trip via parse → print
• tests/ut/jit/test_split_k.py::test_split_k_matmul_* — end-to-end split-K with atomic-add codegen
• Full sweep: tests/ut/ir/, tests/ut/language/, tests/ut/codegen/test_pto_codegen_ops.py — 4347 passed, 28 skipped
• clang-tidy clean

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: c085c04d-25f9-4c14-9222-c4a4c33f4396

📥 Commits

Reviewing files that changed from the base of the PR and between e7e41f5 and eeeb292.

📒 Files selected for processing (4)

• include/pypto/ir/comm.h
• src/ir/op/distributed/put.cpp
• src/ir/transforms/python_printer.cpp
• tests/ut/ir/printing/test_python_printer.py

---

📝 Walkthrough

Walkthrough

This PR adds round-trippable printing of the atomic enum kwarg in tile store operations by introducing an AtomicTypeToString converter and integrating it into the Python IR printer to emit pl.AtomicType.<Name> instead of raw integers.

Changes

AtomicType Enum Printing

Layer / File(s)	Summary
AtomicType to string conversion helper include/pypto/ir/comm.h	Adds <string> and pypto/core/error.h dependencies, then defines AtomicTypeToString(AtomicType) as an inline switch-based converter mapping kNone → "None_" and kAdd → "Add", with pypto::TypeError for unknown values.
Python printer atomic kwarg rendering src/ir/transforms/python_printer.cpp	Includes pypto/ir/comm.h and extends the Call kwargs printer to detect int-typed atomic arguments, converting them to pl.AtomicType.<name> strings via AtomicTypeToString instead of emitting raw integers.
Test atomic kwarg enum form tests/ut/ir/printing/test_python_printer.py	Adds test_python_print_atomic_kwarg_uses_enum_form to verify that tile store calls with integer-encoded atomic values print as pl.AtomicType.Add and pl.AtomicType.None_, and do not contain raw numeric literals.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• hw-native-sys/pypto#295: Both PRs modify src/ir/transforms/python_printer.cpp's Call kwargs printing to render enum-valued kwargs as pl.<EnumType>.<Name> strings instead of raw integers (main: atomic via AtomicTypeToString, retrieved: set_pipe/wait_pipe via PipeTypeToString).

Suggested reviewers

• lyfne123

Poem

🐰 A rabbit hops through enum lands,
Where integers meet string commands,
AtomicType now speaks with grace,
In pl.AtomicType form and place,
Print-round-trip perfection spans! ✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 57.14% 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 accurately and concisely describes the main change: converting the printer to emit AtomicType enum form instead of raw integers for the atomic kwarg.
Description check	✅ Passed	The description thoroughly explains the problem (Pyright flagging Literal[1] incompatible with AtomicType), the solution (printing enum form mirroring PipeType handling), and comprehensive test coverage.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

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 and usage tips.}

[Copilot]

Pull request overview

Fixes the Python printer emitting atomic=1 (raw int) for ops carrying the atomic kwarg, which violated the DSL's atomic: AtomicType signature and tripped Pyright. The printer now emits pl.AtomicType.<Name>, mirroring the existing PipeType handling.

Changes:

• Add AtomicTypeToString helper in include/pypto/ir/comm.h and dispatch on the atomic kwarg in IRPythonPrinter::VisitExpr_ for Calls.
• Fix incidental misc-include-cleaner violations in src/ir/op/distributed/put.cpp exposed by adding <string>/pypto/core/error.h to comm.h.
• New unit test pinning both Add and None_ enum-form printing.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.

File	Description
include/pypto/ir/comm.h	Adds AtomicTypeToString helper; pulls in <string> and pypto/core/error.h.
src/ir/transforms/python_printer.cpp	Prints atomic kwarg as pl.AtomicType.<Name> instead of raw int.
src/ir/op/distributed/put.cpp	Adds direct includes for <cstddef> and pypto/core/logging.h.
tests/ut/ir/printing/test_python_printer.py	New test pinning the enum-form printing for Add and None_.

[gemini-code-assist]

Code Review

This pull request adds support for printing the atomic keyword argument using its Python enum form (pl.AtomicType.<Name>) instead of a raw integer in the Python printer, ensuring type correctness and round-trip compatibility. It also includes corresponding unit tests. The review feedback suggests keeping the core C++ IR header comm.h clean by localizing the Python-specific string conversion logic for AtomicType directly inside python_printer.cpp, which aligns with existing patterns for other enums and avoids header bloat.

[gemini-code-assist]

Separation of Concerns & Header Bloat

Including <string> and "pypto/core/error.h" in a core IR header like comm.h introduces unnecessary compilation coupling and increases header bloat.

Furthermore, defining Python-specific printing helpers (such as AtomicTypeToString which handles the Python-specific "None_" suffix) inside a core C++ IR header couples the core IR definitions with Python-specific printer details.

We should keep comm.h clean and handle this string conversion locally within python_printer.cpp, matching how other printer-specific conversions (like CastModeToString and SplitModeToPythonString) are handled.

[gemini-code-assist]

Separation of Concerns & Header Bloat

This Python-specific helper function should be removed from the core C++ IR header to keep comm.h decoupled from Python printing details. We can inline this conversion directly inside python_printer.cpp where it is used.

[gemini-code-assist]

Localize Python-Specific Printing Logic

Instead of calling a helper from the core comm.h header, we can inline the AtomicType string conversion directly here. This keeps all Python-specific printing logic (including the "None_" suffix workaround) local to the printer, matching the established patterns for other enums in this file. For the default case in the enum switch, use INTERNAL_UNREACHABLE to align with project standards.

      } else if (key == "atomic") {
        // Stored as int (the DSL casts AtomicType -> int before stashing on
        // kwargs_; nb::isinstance<AtomicType> in bindings does the same). The
        // public DSL signature is atomic: AtomicType, so restore the enum
        // form on print to keep the output type-correct for static checkers
        // and round-trippable through the parser (pl.AtomicType is exposed).
        std::string atomic_str;
        switch (static_cast<AtomicType>(int_val)) {
          case AtomicType::kNone:
            atomic_str = "None_";
            break;
          case AtomicType::kAdd:
            atomic_str = "Add";
            break;
          default:
            INTERNAL_UNREACHABLE;
        }
        stream_ << prefix_ << ".AtomicType." << atomic_str;

References

1. For partial switch statements over enums, use a default case with INTERNAL_UNREACHABLE for clarity.