[FIX] Support AOT cross-compilation with COMPILE_ONLY cache save - #382
Closed
coderfeli wants to merge 3 commits into
Closed
[FIX] Support AOT cross-compilation with COMPILE_ONLY cache save#382coderfeli wants to merge 3 commits into
coderfeli wants to merge 3 commits into
Conversation
Two fixes:
1. primitive.py: atom_set_value() was calling fly.atom_set_value(atom, field, value)
but the ODS-generated binding expects (result, atom, field, value). Also, the
field string must be wrapped in StringAttr for the MLIR op to accept it.
This broke any kernel using copy_atom.set_value("soffset", ...) — including
the fused RoPE+KV cache kernel introduced in #377.
2. rocdl/__init__.py: cluster_workgroup_id_{x,y,z} are not present in the
current LLVM ROCDL dialect, causing NameError on import. Use globals().get()
to make them optional, matching the pattern used for other optional ops.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Move COMPILE_ONLY return after cache save so pkl is persisted - Remove address space check in DLTensorAdaptor to allow CPU tensor tracing for cross-arch compilation (e.g. ARCH=gfx950 on gfx942 host) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
COMPILE_ONLYreturn aftercache_manager.set()so pkl is actually persisted to disk. PreviouslyCOMPILE_ONLY=1returned before saving, making AOT pre-compilation useless.DLTensorAdaptor.buildMemRefDesc()to allow CPU tensor tracing. This enablesCOMPILE_ONLY=1 ARCH=gfx950on a gfx942 host (or any machine with LLVM AMDGPU backend).Usage
Test plan
COMPILE_ONLY=1+ CPU tensors → pkl produced (232KB)hipModuleLaunchKernelwill fault)🤖 Generated with Claude Code