refactor(codegen): use Out/InOut params for orchestration output tensors#620
refactor(codegen): use Out/InOut params for orchestration output tensors#620lyfne123 merged 4 commits intohw-native-sys:mainfrom
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughRefactors orchestration codegen to stop creating local return tensors and instead emit reference aliases for kernel outputs; updates many examples and tests to accept external output handles ( Changes
Sequence Diagram(s)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. 📝 Coding Plan
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. Comment |
Summary of ChangesHello, 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 refactors the orchestration code generation process by shifting from implicit inference of output tensors to explicit declaration using Highlights
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. Footnotes
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/st/runtime/test_ctrl_flow.py (1)
524-532:⚠️ Potential issue | 🟡 MinorUpdate orchestrator functions for consistency or add explanatory comment.
The three test cases (
TestForLoopBreak,TestForLoopContinue, andTestForLoopBreakContinue) usepl.create_tensor()to allocate output tensors in their orchestrator functions, while all other test cases in this file usepl.Out[...]parameters. Additionally, the InCore kernel functions within these same test cases usepl.Out[...], creating an inconsistency within each test class.Either update these orchestrator functions to use
pl.Out[...]for consistency, or add a brief comment explaining why this pattern differs.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/st/runtime/test_ctrl_flow.py` around lines 524 - 532, The orchestrator functions (e.g., orchestrator in TestForLoopBreak/TestForLoopContinue/TestForLoopBreakContinue) are inconsistent with the rest of the file by allocating outputs with pl.create_tensor rather than taking pl.Out[...] parameters while their kernels (kernel_break, kernel_continue, etc.) use pl.Out; update each orchestrator signature to accept the output tensor as a pl.Out[[256, 64], pl.FP32] parameter and wire that through when calling the corresponding kernel (e.g., pass the pl.Out c into kernel_break), or if create_tensor is intentional, add a short explanatory comment above each orchestrator explaining why create_tensor is required for these tests to justify the deviation.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@tests/st/runtime/test_ctrl_flow.py`:
- Around line 524-532: The orchestrator functions (e.g., orchestrator in
TestForLoopBreak/TestForLoopContinue/TestForLoopBreakContinue) are inconsistent
with the rest of the file by allocating outputs with pl.create_tensor rather
than taking pl.Out[...] parameters while their kernels (kernel_break,
kernel_continue, etc.) use pl.Out; update each orchestrator signature to accept
the output tensor as a pl.Out[[256, 64], pl.FP32] parameter and wire that
through when calling the corresponding kernel (e.g., pass the pl.Out c into
kernel_break), or if create_tensor is intentional, add a short explanatory
comment above each orchestrator explaining why create_tensor is required for
these tests to justify the deviation.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 5a2ac031-bf8d-4d67-a831-dc874c631271
📒 Files selected for processing (25)
docs/en/dev/codegen/02-orchestration_codegen.mddocs/zh-cn/dev/codegen/02-orchestration_codegen.mdexamples/ir_parser/batch_paged_attention_example.pyexamples/ir_parser/orchestration_example.pyexamples/ir_parser/paged_attention_example.pyexamples/ir_parser/vector_example_dag.pyexamples/language/beginner/basic_ops.pyexamples/language/beginner/elementwise.pyexamples/language/beginner/hello_world.pyexamples/language/beginner/matmul.pyexamples/language/intermediate/activation.pyexamples/language/intermediate/ffn_activations.pyexamples/language/intermediate/layer_norm.pyexamples/language/intermediate/rms_norm.pyexamples/language/intermediate/softmax.pyexamples/language/intermediate/vector_dag.pyexamples/language/llm_models/llama_7b_mini.pysrc/codegen/orchestration/orchestration_codegen.cpptests/st/codegen/test_batch_paged_attention.pytests/st/codegen/test_paged_attention.pytests/st/runtime/test_ctrl_flow.pytests/st/runtime/test_dynamic_shape.pytests/st/runtime/test_fillpad.pytests/st/runtime/test_matmul.pytests/ut/codegen/test_orchestration_codegen.py
|
Warning Gemini is experiencing higher than usual traffic and was unable to create the review. Please try again in a few hours by commenting |
Replace return-var inference with explicit pl.Out/pl.InOut parameter annotations in orchestration functions. Output tensors are now passed as params rather than allocated via pl.create_tensor() in the body. C++ codegen removes dead code: CountReturnTensors, CountExpectedArgs, GetIntermediateTensorType, return_vars tracking, and return_names_ member. Adds alias generation for InCore call return values that map to Out/InOut args.
…t params Migrate output tensor declarations from local pl.create_tensor() to explicit pl.Out[...] parameters across all examples and tests, aligning with the Out/InOut orchestration codegen refactor.
Remove output_tensors, output_tensor_assigns, tuple_element_map, and call_to_result_var which were populated but never read after the Out/InOut param refactor. Also remove the no-op SetTupleElementMap method and its call site.
2 participants
Summary
Out/InOutfunction parameters instead of inferring them from return statementspl.Out[...]parameter syntax for output tensorsOrchestrationInfoCollectorthat were no longer read after the refactorTesting
Related Issues
fix #583