feat(ptodsl): support source-backed jit kernels#892
Conversation
Codex Review该评论由 review 机器人自动更新。
SummaryReview failed at stage Findings未生成结构化 findings，因为 review 过程提前失败。 Log Tail |
![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 support for loading existing hand-written PTO files in PTODSL using the source parameter in the @pto.jit decorator. This allows users to maintain a Python signature as the host ABI declaration while loading the kernel implementation directly from a PTO IR file instead of tracing the Python body. The changes include comprehensive diagnostics, a new SourceModuleLoader for parsing and ABI verification, integration with the compilation and native build pipelines, and extensive test coverage. The review feedback highlights critical robustness improvements, specifically addressing potential AttributeError crashes in _walk_ops and _select_entry when handling raw MLIR operations, as well as a potential NameError in the test suite due to a missing tempfile import.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| def _walk_ops(root_op): | ||
| for region in root_op.regions: | ||
| for block in region.blocks: | ||
| for op in block.operations: | ||
| yield op | ||
| yield from _walk_ops(op.operation) |
There was a problem hiding this comment.
The recursive helper _walk_ops assumes that every operation yielded by block.operations is an OpView subclass (which has an .operation attribute). However, if the source file contains unregistered operations or operations from dialects that are not loaded, block.operations will yield raw mlir.ir.Operation objects. Calling op.operation on a raw Operation will raise an AttributeError and crash the loader.\n\nWe should make _walk_ops robust by checking if the object has an operation attribute before accessing it, or by passing the raw operation to the recursive call.
| def _walk_ops(root_op): | |
| for region in root_op.regions: | |
| for block in region.blocks: | |
| for op in block.operations: | |
| yield op | |
| yield from _walk_ops(op.operation) | |
| def _walk_ops(root_op): | |
| op = root_op.operation if hasattr(root_op, "operation") else root_op | |
| for region in op.regions: | |
| for block in region.blocks: | |
| for nested_op in block.operations: | |
| yield nested_op | |
| yield from _walk_ops(nested_op) |
| def _select_entry(self, module: Module, resolved_path: Path): | ||
| matches = [] | ||
| for op in _walk_ops(module.operation): | ||
| if op.operation.name != "func.func": | ||
| continue | ||
| if getattr(op, "is_external", False): | ||
| continue | ||
| if _symbol_name(op) == self._module_spec.function_name: | ||
| matches.append(op) |
There was a problem hiding this comment.
Similarly to _walk_ops, _select_entry accesses op.operation.name directly. If op is a raw mlir.ir.Operation (e.g., an unregistered operation), this will raise an AttributeError. We should resolve the raw operation safely before checking its name.
| def _select_entry(self, module: Module, resolved_path: Path): | |
| matches = [] | |
| for op in _walk_ops(module.operation): | |
| if op.operation.name != "func.func": | |
| continue | |
| if getattr(op, "is_external", False): | |
| continue | |
| if _symbol_name(op) == self._module_spec.function_name: | |
| matches.append(op) | |
| def _select_entry(self, module: Module, resolved_path: Path): | |
| matches = [] | |
| for op in _walk_ops(module.operation): | |
| raw_op = op.operation if hasattr(op, "operation") else op | |
| if raw_op.name != "func.func": | |
| continue | |
| if getattr(op, "is_external", False): | |
| continue | |
| if _symbol_name(op) == self._module_spec.function_name: | |
| matches.append(op) |
| yield None | ||
| return | ||
|
|
||
| with tempfile.TemporaryDirectory() as temp_dir: |
There was a problem hiding this comment.
If tempfile is not imported in this file, calling tempfile.TemporaryDirectory() will raise a NameError. You can import tempfile inline to make this context manager self-contained and robust.
| with tempfile.TemporaryDirectory() as temp_dir: | |
| import tempfile | |
| with tempfile.TemporaryDirectory() as temp_dir: |
2 participants
Summary
This PR lets PTODSL launch a hand-written PTO source file through the existing
@pto.jitruntime path by addingsource=...support to launchable entrykernels.
It keeps the Python function signature as the host ABI declaration, skips
tracing the Python body when
sourceis provided, and validates that theselected PTO entry matches the declared ABI before native build and launch.
What Changed
@pto.jit(source=...)support for launchable entry kernelsfunc.funcby JIT namePython
KernelSignatureentry=Falsepto.const_expr.compile(...)constexpr bindings--pto-backend=<backend>mode="explicit"->--pto-level=level3command coverage
test/vpto/cases/micro-op/a5-extra/vmaddinto a PTODSLsource-backed runtime case using the existing
kernel.ptoValidation
python3 ptodsl/tests/test_jit_compile.pypython3 ptodsl/tests/test_jit_diagnostics.pyPTOAS_ENV_SKIP_SMOKE_TEST=1 WORK_SPACE=/tmp/pto-vpto-ptodsl-vmadd-pr CASE_NAME='micro-op/a5-extra/vmadd' DEVICE=SIM test/vpto/scripts/run_host_vpto_validation.sh