[Enh] Add CompiledFunction.max_blocks_per_cu()

lib/Runtime/ROCm/FlyRocmRuntimeWrappers.cpp

@@ -45,6 +45,20 @@ extern "C" hipFunction_t mgpuModuleGetFunction(hipModule_t module, const char *n
   return function;
 }
 
+// Max active blocks of kernel `name` per CU; loads a temporary module, launches nothing.
+extern "C" int32_t mgpuModuleOccupancyMaxActiveBlocks(void *data, const char *name,
+                                                      int32_t blockSize, int32_t smem) {
+  hipModule_t module = mgpuModuleLoad(data, /*gpuBlobSize=*/0);
+  if (!module)
+    return 0;
+  int blocks = 0;
+  if (hipFunction_t function = mgpuModuleGetFunction(module, name))
+    HIP_REPORT_IF_ERROR(
+        hipModuleOccupancyMaxActiveBlocksPerMultiprocessor(&blocks, function, blockSize, smem));
+  mgpuModuleUnload(module);
+  return blocks;
+}
+
 extern "C" void mgpuLaunchKernel(hipFunction_t function, intptr_t gridX, intptr_t gridY,
                                  intptr_t gridZ, intptr_t blockX, intptr_t blockY, intptr_t blockZ,
                                  int32_t smem, hipStream_t stream, void **params, void **extra,

python/flydsl/compiler/jit_executor.py

@@ -386,6 +386,41 @@ def dump(self, compiled: bool = True):
                 print("=" * 60)
                 print(self._source_ir)
 
+    def max_blocks_per_cu(self) -> int:
+        from .._mlir.dialects import gpu
+        from .jit_function import _create_mlir_context
+
+        def constant(value):
+            if value is None:
+                return 0
+            if getattr(value.owner, "name", None) != "llvm.mlir.constant":
+                raise ValueError("max_blocks_per_cu() needs constant block and smem sizes")
+            return ir.IntegerAttr(value.owner.attributes["value"]).value
+
+        binaries, launches = {}, set()
+
+        def visit(op):
+            if op.name == "gpu.binary":
+                # The JIT loads object 0 (selectObject in FlyLLVMTranslation.cpp).
+                objects = ir.ArrayAttr(op.attributes["objects"])
+                binaries[ir.StringAttr(op.attributes["sym_name"]).value] = gpu.ObjectAttr(objects[0]).object
+            elif op.name == "gpu.launch_func":
+                launch = op.opview
+                block = constant(launch.blockSizeX) * constant(launch.blockSizeY) * constant(launch.blockSizeZ)
+                launches.add((*ir.SymbolRefAttr(launch.kernel).value, block, constant(launch.dynamicSharedMemorySize)))
+            return ir.WalkResult.ADVANCE
+
+        with _create_mlir_context():
+            ir.Module.parse(self._ir_text).operation.walk(visit)
+        if len(launches) != 1:
+            raise ValueError(f"max_blocks_per_cu() needs exactly one kernel launch, found {len(launches)}")
+        module_name, kernel_name, block, smem = launches.pop()
+        lib = ctypes.CDLL(_resolve_runtime_libs()[0])
+        blocks = lib.mgpuModuleOccupancyMaxActiveBlocks(binaries[module_name], kernel_name.encode(), block, smem)
+        if blocks <= 0:
+            raise RuntimeError("max_blocks_per_cu(): HIP occupancy query failed")
+        return blocks
+
     @property
     def ir(self) -> str:
         return self._ir_text

python/flydsl/compiler/jit_function.py

@@ -1723,6 +1723,10 @@ def __call__(self, *args):
         assert state is not None
         return state(args)
 
+    def max_blocks_per_cu(self) -> int:
+        """Max resident blocks of this kernel per CU on the current device, without launching it."""
+        return self._keepalive.max_blocks_per_cu()
+
     def export_to_c(self, file_path: str, file_name: str, function_prefix: str = "") -> None:
         """Export this specialization as a standalone ``.o`` and C header.
 

tests/unit/test_max_blocks_per_cu.py

@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 FlyDSL Project Contributors
+
+import pytest
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+
+pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower]
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+if torch is None or not torch.cuda.is_available():
+    pytest.skip("CUDA/ROCm not available", allow_module_level=True)
+
+
+@flyc.kernel
+def _kernel(x: fx.Tensor):
+    pass
+
+
+@flyc.jit
+def _launch(x: fx.Tensor, smem: fx.Constexpr[int], stream: fx.Stream = fx.Stream(None)):
+    _kernel(x).launch(grid=(1, 1, 1), block=(64, 1, 1), smem=smem, stream=stream)
+
+
+def test_max_blocks_per_cu():
+    # CPU tensors are compile placeholders, so nothing is launched.
+    free = flyc.compile(_launch, torch.zeros(64), 0).max_blocks_per_cu()
+    lds_bound = flyc.compile(_launch, torch.zeros(64), 48 * 1024).max_blocks_per_cu()
+    assert 1 <= lds_bound < free