feat(compile): surface AMDGPU kernel resource-usage info during compile#796
Closed
fsx950223 wants to merge 3 commits into
Closed
feat(compile): surface AMDGPU kernel resource-usage info during compile#796fsx950223 wants to merge 3 commits into
fsx950223 wants to merge 3 commits into
Conversation
Parse LLVM's "; Kernel info:" AsmPrinter comment block (NumVgprs, NumSgprs, Occupancy, LDS/scratch usage) out of the ISA text FLYDSL_DUMP_IR already dumps, and print it during compilation. Requires the companion ROCDL AsmVerbose patch in llvm-project for the comment block to be emitted; the parser degrades gracefully to an empty dict otherwise. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds a small utility to parse AMDGPU “; Kernel info:” resource-usage comments from generated ISA, and wires it into the existing ISA dump path so compile-time dumps can surface register/LDS/occupancy information.
Changes:
- Add
flydsl.utils.kernel_info.parse_kernel_info()to extract key/value pairs from LLVM’s AMDGPU AsmPrinter “Kernel info” comment block. - Print parsed kernel resource info from
jit_function._dump_isa()whenFLYDSL_DUMP_IR=1produces ISA text. - Add unit tests for the parser and occupancy helper.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
python/flydsl/utils/kernel_info.py |
New parser for AMDGPU “Kernel info” block (currently missing get_occupancy implementation and has slightly misleading doc wording). |
python/flydsl/compiler/jit_function.py |
Hooks the parser into the ISA-dump path and prints parsed info when available. |
tests/unit/test_kernel_info.py |
Adds unit tests for parsing and occupancy extraction (currently imports get_occupancy which isn’t defined). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+15
to
+19
| This is only emitted when the target machine is built with ``AsmVerbose`` | ||
| on. FlyDSL's ROCDL serializer (``mlir/lib/Target/LLVM/ROCDL/Target.cpp``, | ||
| ``SerializeGPUModuleBase::getTargetOptions()``) enables it, so | ||
| ``jit_function._dump_isa``'s ``.s`` output always carries this block. | ||
|
|
Comment on lines
+43
to
+50
| info: Dict[str, str] = {} | ||
| for line in isa_text[block_start:block_end].splitlines()[1:]: | ||
| line = line.lstrip(";").strip() | ||
| if ":" not in line: | ||
| continue | ||
| key, _, value = line.partition(":") | ||
| info[key.strip()] = value.strip() | ||
| return info |
The companion LLVM patch (ROCm/llvm-project#3177) now makes ROCDL's AsmVerbose an opt-in via `-asm-verbose` in gpu-module-to-binary's `opts` argument rather than always-on, so pass it explicitly from _dump_isa. Also restores flydsl.utils.kernel_info.get_occupancy(), which was dropped from the previous commit by an editing mishap despite still being imported by tests/unit/test_kernel_info.py. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Contributor
Author
|
Pushed a follow-up commit:
|
No callers need the Occupancy field on its own; parse_kernel_info() already exposes it as part of the full dict. Also corrects a stale docstring claim that AsmVerbose is unconditionally enabled -- it's opt-in via -asm-verbose as of the companion LLVM patch. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Contributor
Author
|
Superseded — cherry-picked onto #613 (see comment there) since both touch |
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
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
flydsl.utils.kernel_info.parse_kernel_info()/get_occupancy(), which parse LLVM's AMDGPU AsmPrinter; Kernel info:comment block (NumVgprs,NumSgprs,Occupancy,LDSByteSize,ScratchSize, ...) out of ISA text, similar to how Triton exposeskernel.n_regs/occupancy via itsllc-based codegen.jit_function._dump_isa(the existingFLYDSL_DUMP_IR=1ISA-dump path) so compiling any kernel now prints[flydsl.compile] kernel info: {...}alongside the.sdump.Dependency
This requires the companion LLVM/MLIR patch ROCm/llvm-project#3177, which enables
AsmVerbosefor the ROCDLgpu-module-to-binaryserialization path. Without it, FlyDSL's codegen only emits the.amdgpu_metadataYAML note (noOccupancyfield), andparse_kernel_inforeturns an empty dict.Test plan
pytest tests/unit/test_kernel_info.py -v-- 3/3 pass (pure parser tests,l0_backend_agnostic, no GPU/build required).scripts/check_python_style.sh --include-localclean.pa_decode_ps_kernelviatests/kernels/test_pa.py'srun_pa_decode_ps_testwithFLYDSL_DUMP_IR=1against a toolchain built with the companion LLVM patch; confirmed[flydsl.compile] kernel info: {..., 'Occupancy': '3', 'NumVgprs': '130', ...}is printed and matches the real register allocation.🤖 Generated with Claude Code