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: |
zhangstevenunity
left a comment
zhangstevenunity
left a comment
There was a problem hiding this comment.
Reviewed the source-backed @pto.jit path end-to-end: the loader (resolve/parse/select/ABI-verify), specialization-cache identity, native-build flag overrides, the docs-as-test files harness, and the vpto case-discovery rewrite. The feature is well-built and the Python test coverage is thorough. No blocking issues; two items below.
Note on the existing automated reviews:
- The Codex bot run failed on INSUFFICIENT_BALANCE and produced no findings.
- All three Gemini findings are false positives:
_walk_ops/_select_entry"raw Operation has no .operation" (HIGH x2): incorrect.mlir.ir.Operationexposes.operation(returns self), soop.operation.nameand the recursive_walk_ops(op.operation)are safe for both Operation and OpView. Additionallymake_context()does not allow unregistered dialects, soModule.parserejects unknown ops before_walk_opsruns; and the ABI path already depends on the typed FuncOp view (entry.type.inputs), which the passing count/type/void diagnostics tests confirm is present.- tempfile NameError (MEDIUM):
test_docs_as_test.pyalready importstempfile(line 20) andcontextmanager(line 10).
My own findings are posted inline (1x P2 validation reproducibility, 1x P3 diagnostic gap).
| [[ -f "${requested_dir}/kernel.pto" ]] || | ||
| die "case ${CASE_NAME} must provide kernel.pto" | ||
| local requested_path="${CASES_ROOT}/${CASE_NAME}" | ||
| [[ -e "${requested_path}" ]] || die "unknown case: ${CASE_NAME}" |
There was a problem hiding this comment.
P2: the PR's documented SIM validation command no longer resolves against the final layout.
The vmadd case was flattened from a directory (micro-op/a5-extra/vmadd/) to a single file (micro-op/a5-extra/vmadd.py), but the Validation section still runs:
CASE_NAME='micro-op/a5-extra/vmadd' ... run_host_vpto_validation.sh
With the new resolution logic, CASE_NAME is matched by the exact -e "${CASES_ROOT}/${CASE_NAME}" test on this line. micro-op/a5-extra/vmadd no longer exists (it is vmadd.py), so this dies with "unknown case: micro-op/a5-extra/vmadd". The correct name is now micro-op/a5-extra/vmadd.py.
Please either (a) update the recorded command to CASE_NAME='micro-op/a5-extra/vmadd.py' and confirm the SIM run passes on the final file-based case, or (b) accept the extension-less basename (append .py when the bare path is absent). As written, the only device/SIM validation for this feature cannot be reproduced from the PR's own command, which suggests the final flattened layout may not have been re-run.
| content_digest = hashlib.sha256(mlir_text.encode("utf-8")).hexdigest() | ||
| ctx = make_context() | ||
| with ctx, Location.unknown(): | ||
| module = Module.parse(mlir_text) |
There was a problem hiding this comment.
P3: malformed PTO source bypasses the friendly jit_source_* diagnostics.
The loader wraps file-not-found (jit_source_file_error), entry selection (jit_source_entry_error), and ABI mismatch (jit_source_abi_error) in named diagnostics that quote source= and the resolved path. But a syntactically invalid source file makes Module.parse(mlir_text) raise a bare mlir.ir.MLIRError that names neither the source= argument nor the resolved path -- which is a very likely first failure when hand-editing a .pto file. Consider catching the parse failure here and re-raising through a jit_source_* helper. Same spot: the module.operation.verify() result on line 90 is discarded, so a parse-clean-but-invalid module is handed to the native build without a friendly message.
A3 板测完成（有跳过）
|
4 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