feat(compile): surface AMDGPU kernel resource-usage info during compile

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>

Author: fsx950223

python/flydsl/compiler/jit_function.py

@@ -24,6 +24,7 @@
 from .._mlir.passmanager import PassManager
 from ..expr.typing import Constexpr, Stream
 from ..utils import env, log
+from ..utils.kernel_info import parse_kernel_info
 from .ast_rewriter import ASTRewriter
 from .backends import compile_backend_name, get_backend
 from .jit_argument import convert_to_jit_arguments, is_type_param_annotation, resolve_signature
@@ -698,6 +699,11 @@ def _dump_isa(*, dump_dir: Path, ctx: ir.Context, asm: str, verify: bool, stage_
         dump_dir.mkdir(parents=True, exist_ok=True)
         out = dump_dir / f"{stage_name}.s"
         out.write_text(isa_text, encoding="utf-8")
+
+        kernel_info = parse_kernel_info(isa_text)
+        if kernel_info:
+            print(f"[flydsl.compile] kernel info: {kernel_info}")
+
         return out
     except Exception as exc:
         log().debug(f"[dump_isa] failed: {exc}")

python/flydsl/utils/kernel_info.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 FlyDSL Project Contributors
+"""Parse the LLVM AMDGPU AsmPrinter "Kernel info" comment block.
+
+When LLVM lowers a kernel to AMDGPU assembly, it emits an exact,
+compiler-computed resource-usage summary as a comment block, e.g.::
+
+    ; Kernel info:
+    ; NumVgprs: 6
+    ; NumAgprs: 0
+    ; Occupancy: 8
+    ; LDSByteSize: 0 bytes/workgroup (compile time only)
+    ...
+
+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.
+
+This module parses that block into a plain dict, so callers don't have to
+re-derive occupancy/register usage from device properties themselves.
+"""
+
+from typing import Dict
+
+KERNEL_INFO_HEADER = "; Kernel info:"
+
+
+def parse_kernel_info(isa_text: str) -> Dict[str, str]:
+    """Parse the ``; Kernel info:`` comment block out of AMDGPU ISA text.
+
+    Returns an empty dict if the block is not present (e.g. non-AMDGPU ISA,
+    or a source that didn't come from asm-verbose codegen).
+    """
+    block_start = isa_text.find(KERNEL_INFO_HEADER)
+    if block_start == -1:
+        return {}
+
+    block_end = isa_text.find("\n\t", block_start)
+    if block_end == -1:
+        block_end = len(isa_text)
+
+    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

tests/unit/test_kernel_info.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 FlyDSL Project Contributors
+
+"""Unit tests for flydsl.utils.kernel_info's AMDGPU ISA "Kernel info" parser."""
+
+import pytest
+
+from flydsl.utils.kernel_info import get_occupancy, parse_kernel_info
+
+pytestmark = [pytest.mark.l0_backend_agnostic]
+
+SAMPLE_ISA = """\
+	.text
+	.globl	add_kernel
+; Kernel info:
+; codeLenInByte = 72
+; TotalNumSgprs: 14
+; NumVgprs: 6
+; NumAgprs: 0
+; TotalNumVgprs: 6
+; ScratchSize: 0
+; MemoryBound: 1
+; FloatMode: 240
+; IeeeMode: 1
+; LDSByteSize: 0 bytes/workgroup (compile time only)
+; SGPRBlocks: 1
+; VGPRBlocks: 0
+; NumSGPRsForWavesPerEU: 14
+; NumVGPRsForWavesPerEU: 6
+; AccumOffset: 8
+; Occupancy: 8
+; WaveLimiterHint : 1
+; COMPUTE_PGM_RSRC2:SCRATCH_EN: 0
+; COMPUTE_PGM_RSRC2:USER_SGPR: 2
+	.section	.rodata
+"""
+
+
+def test_parse_kernel_info_extracts_all_fields():
+    info = parse_kernel_info(SAMPLE_ISA)
+    assert info["NumVgprs"] == "6"
+    assert info["NumAgprs"] == "0"
+    assert info["Occupancy"] == "8"
+    assert info["LDSByteSize"] == "0 bytes/workgroup (compile time only)"
+    assert info["TotalNumSgprs"] == "14"
+
+
+def test_parse_kernel_info_missing_block_returns_empty():
+    assert parse_kernel_info("no kernel info here") == {}
+
+
+def test_get_occupancy():
+    assert get_occupancy(SAMPLE_ISA) == 8
+    assert get_occupancy("nothing") is None