feat(codegen): Add kernel wrapper generation for PTO backend#272
feat(codegen): Add kernel wrapper generation for PTO backend#272Hzfengsy merged 1 commit intohw-native-sys:mainfrom
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThe pull request refactors the PTO backend code generation and compilation pipeline by extracting PTOAS handling from compile.py into a dedicated pto_codegen.py module, adding orchestration code generation support to C++ bindings, updating the pass manager with new optimization passes, and providing comprehensive test coverage for the new implementation. Changes
Sequence DiagramsequenceDiagram
participant App as Application
participant Compile as compile.py
participant PtoCg as pto_codegen.py
participant PTO as PTOCodegen
participant Ptoas as PTOAS Tool
participant Orch as orchestration<br/>(shared codegen)
participant Output as Output Files
App->>Compile: compile with PTO backend
Compile->>PtoCg: generate(program, output_dir)
PtoCg->>PTO: for each InCore function
PTO->>PTO: create .pto representation
PtoCg->>Ptoas: invoke with .pto
Ptoas->>Ptoas: compile to C++
Ptoas-->>PtoCg: return .cpp output
PtoCg->>PtoCg: _preprocess_ptoas_output()
PtoCg->>PtoCg: _generate_arg_unpacking()
PtoCg->>PtoCg: _generate_kernel_wrapper()
PtoCg->>Orch: generate_orchestration(program, func)
Orch-->>PtoCg: orchestration code + metadata
PtoCg->>Output: _write_files() all artifacts
Output-->>Compile: files dict
Compile->>Output: persist to ptoas/, kernels/aiv/
Output-->>App: build output
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Comment |
Summary of ChangesHello @zhangqi-chen, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the PyPTO IR compiler's PTO backend by introducing a comprehensive kernel wrapper generation mechanism and refactoring its code generation logic into a dedicated module. These changes streamline the compilation process for PTOAS, ensure compatibility with the CCE/orchestration calling convention, and provide clearer separation of concerns. Additionally, new C++ bindings for orchestration functions and an enriched PTOAS optimization strategy contribute to a more robust and flexible compilation framework. Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a significant and well-structured refactoring for the PTO backend codegen. Extracting the logic into a new pto_codegen.py module is a good design choice that improves modularity. The new kernel wrapper generation is a crucial feature that bridges the gap between ptoas output and the CCE calling convention. The added C++ bindings and corresponding tests are also valuable additions.
I've identified one critical issue regarding a hardcoded path that could lead to runtime failures, along with a couple of medium-severity suggestions for improving the robustness of the new module. Overall, this is a solid feature addition, and addressing these points will enhance its reliability.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
examples/ir_parser/vector_example_dag.py (1)
136-142: Minor: Inconsistent step numbers in warning messages.The warning messages use
[5]but appear under different logical steps:
- Line 137 is checking the kernel directory (Step 4), so should use
[4]- Line 142 is checking the orchestration file (Step 5), so
[5]is correct📝 Suggested fix
kernel_dir = os.path.join(output_dir, "kernels") if not os.path.isdir(kernel_dir): - print(f"\n[5] Warning: {kernel_dir} not found") + print(f"\n[4] Warning: {kernel_dir} not found") # Step 5: Show orchestration code orch_file = os.path.join(output_dir, "orchestration", "orch_vector.cpp")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/ir_parser/vector_example_dag.py` around lines 136 - 142, The warning for the kernel directory uses the wrong step number: update the print for the kernel_dir existence check to use "[4]" instead of "[5]" so it matches the logical Step 4 (locate the kernel_dir variable in the block that checks os.path.isdir(kernel_dir)); leave the orch_file check and its "[5]" warning unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@python/pypto/ir/pto_codegen.py`:
- Around line 279-281: The wrapper emission always writes to "kernels/aiv/..."
in _generate_kernel_wrapper handling (assignment to kernel_rel and
result_files), but _generate_config_file uses func_name_to_core_type to resolve
source paths (which may be "aic"), causing path mismatches; change the
kernel_rel computation to pick the core type for this function from
func_name_to_core_type (e.g., core = func_name_to_core_type.get(func.name,
<sane-default>)) and join os.path.join("kernels", core, f"{func.name}.cpp")
before assigning result_files[kernel_rel] so emitted wrapper paths match the
config lookup used by _generate_config_file.
---
Nitpick comments:
In `@examples/ir_parser/vector_example_dag.py`:
- Around line 136-142: The warning for the kernel directory uses the wrong step
number: update the print for the kernel_dir existence check to use "[4]" instead
of "[5]" so it matches the logical Step 4 (locate the kernel_dir variable in the
block that checks os.path.isdir(kernel_dir)); leave the orch_file check and its
"[5]" warning unchanged.
ℹ️ Review info
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (9)
docs/dev/codegen/00-pto_codegen.mdexamples/ir_parser/vector_example_dag.pypython/bindings/modules/codegen.cpppython/pypto/ir/compile.pypython/pypto/ir/pass_manager.pypython/pypto/ir/pto_codegen.pypython/pypto/pypto_core/codegen.pyitests/ut/codegen/test_pto_codegen.pytests/ut/ir/transforms/test_pass_manager.py
Extract PTO codegen logic into dedicated pto_codegen module with kernel wrapper generation, config file generation, and PTOAS output preprocessing. Expose OrchestrationResult, GenerateOrchestration, and InferFunctionCoreType through C++ bindings and type stubs. Add PTOAS strategy to pass manager with ConvertToSSA, FlattenCallExpr, and RunVerifier passes.
2 participants
Extract PTO codegen logic into dedicated pto_codegen module with kernel wrapper generation, config file generation, and PTOAS output preprocessing. Expose OrchestrationResult, GenerateOrchestration, and InferFunctionCoreType through C++ bindings and type stubs. Add PTOAS strategy to pass manager with ConvertToSSA, FlattenCallExpr, and RunVerifier passes.