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

python/flydsl/compiler/jit_function.py

@@ -24,6 +24,7 @@
 from ..expr.meta import tracing_context
 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 .diagnostics import (
@@ -651,8 +652,9 @@ def _dump_isa(*, dump_dir: Path, ctx: ir.Context, asm: str, verify: bool, stage_
         di_pass = (
             "ensure-debug-info-scope-on-llvm-func{emission-kind=LineTablesOnly}," if env.debug.enable_debug_info else ""
         )
+        cmd_opts = "-asm-verbose" + (" -g" if env.debug.enable_debug_info else "")
         pm = PassManager.parse(
-            f'builtin.module({di_pass}gpu-module-to-binary{{format=isa opts="{"-g" if env.debug.enable_debug_info else ""}" section= toolkit=}})',
+            f'builtin.module({di_pass}gpu-module-to-binary{{format=isa opts="{cmd_opts}" section= toolkit=}})',
             context=ctx,
         )
         pm.enable_verifier(bool(verify))
@@ -664,6 +666,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,51 @@
+# 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()``) supports enabling it via
+``-asm-verbose`` in ``gpu-module-to-binary``'s ``opts`` argument, which
+``jit_function._dump_isa`` passes so its ``.s`` output 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,49 @@
+# 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 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") == {}