feat(rocke): expose ASIC revision through typed HIP queries

dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/__init__.py

@@ -33,6 +33,10 @@
                       `HIP_LAUNCH_PARAM_BUFFER_POINTER` arg-buffer
                       lifetime race.
 
+  - ``device_capabilities``: reports named device features from the
+                      active gfx target and ASIC revision. Callers do
+                      not need to compare revision numbers themselves.
+
   - ``packing``     : torch-agnostic kernel-arg packing (`pack_args`,
                       `pack_args_kernelparams`) for the AMDGPU kernarg
                       ABI. Used by both the numpy and torch paths.
@@ -84,6 +88,11 @@
 from __future__ import annotations
 
 from .comgr import ComgrError, ComgrTimings, build_hsaco_from_llvm_ir
+from .device_capabilities import (
+    DeviceCapabilities,
+    DeviceCapability,
+    get_device_capabilities,
+)
 from .hip_module import HipError, Runtime
 from .launcher import (
     DeviceMem,
@@ -109,6 +118,8 @@
 __all__ = [
     "ComgrError",
     "ComgrTimings",
+    "DeviceCapabilities",
+    "DeviceCapability",
     "DeviceMem",
     "HipError",
     "KernelLauncher",
@@ -120,6 +131,7 @@
     "WorkspacePool",
     "build_hsaco_from_llvm_ir",
     "empty_workspace",
+    "get_device_capabilities",
     "launch_torch_kernel",
     "no_fence",
     "pack_args",

dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/device_capabilities.py

@@ -0,0 +1,128 @@
+# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+# SPDX-License-Identifier: MIT
+
+"""Named device capabilities reported from HIP properties.
+
+Keep revision checks in one place so dispatch and launch code do not compare
+``asicRevision`` values directly. Features tied to an unknown revision are reported
+as unsupported. The current table covers gfx1250 and can be extended with other
+targets. HIP still decides whether a launch succeeds.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+
+from .hip_module import _get_device_asic_revision, get_device_arch
+
+
+class DeviceCapability(str, Enum):
+    """Device features recognized by rocKE."""
+
+    WORKGROUP_CLUSTER_LAUNCH = "workgroup_cluster_launch"
+    TDM_MULTICAST = "tdm_multicast"
+    MX_WMMA_FP4_32X16 = "mx_wmma_fp4_32x16"
+    MX_BLOCK16_CONVERSION = "mx_block16_conversion"
+
+
+@dataclass(frozen=True)
+class _CapabilityPolicy:
+    """Capabilities shared by a target and added by each known revision."""
+
+    base: frozenset[DeviceCapability]
+    by_revision: dict[int, frozenset[DeviceCapability]]
+
+
+_CAPABILITY_POLICIES = {
+    "gfx1250": _CapabilityPolicy(
+        base=frozenset({DeviceCapability.WORKGROUP_CLUSTER_LAUNCH}),
+        by_revision={
+            0: frozenset(),
+            1: frozenset(
+                {
+                    DeviceCapability.TDM_MULTICAST,
+                    DeviceCapability.MX_WMMA_FP4_32X16,
+                    DeviceCapability.MX_BLOCK16_CONVERSION,
+                }
+            ),
+        },
+    )
+}
+
+
+@dataclass(frozen=True)
+class DeviceCapabilities:
+    """Capability result for one HIP device.
+
+    ``check`` returns a short explanation when a feature is not supported. Callers
+    do not need to compare ASIC revision numbers themselves.
+    """
+
+    arch: str | None
+    asic_revision: int | None
+    supported: frozenset[DeviceCapability]
+
+    def check(self, capability: DeviceCapability) -> tuple[bool, str]:
+        """Return whether ``capability`` is supported and explain the result."""
+
+        if not isinstance(capability, DeviceCapability):
+            raise TypeError(
+                "capability must be DeviceCapability, "
+                f"got {type(capability).__name__}"
+            )
+        if capability in self.supported:
+            return True, "supported"
+        if self.arch is None:
+            return False, "device architecture is unavailable"
+        policy = _CAPABILITY_POLICIES.get(self.arch)
+        if policy is None:
+            return False, f"no capabilities are listed for {self.arch}"
+        if self.asic_revision is None:
+            return False, f"{self.arch} ASIC revision is unavailable"
+        if self.asic_revision < 0:
+            return False, f"invalid {self.arch} ASIC revision {self.asic_revision}"
+        if self.asic_revision not in policy.by_revision:
+            return False, f"unknown {self.arch} ASIC revision {self.asic_revision}"
+        return False, (
+            f"{capability.value} is unavailable on "
+            f"{self.arch} ASIC revision {self.asic_revision}"
+        )
+
+    def supports(self, capability: DeviceCapability) -> bool:
+        """Return whether ``capability`` is supported."""
+
+        return self.check(capability)[0]
+
+
+def _capabilities_for_properties(
+    arch: str | None, asic_revision: int | None
+) -> DeviceCapabilities:
+    """Build a capability result from device properties already read from HIP."""
+
+    policy = _CAPABILITY_POLICIES.get(arch)
+    if policy is None:
+        supported = frozenset()
+    else:
+        supported = policy.base | policy.by_revision.get(asic_revision, frozenset())
+    return DeviceCapabilities(
+        arch=arch,
+        asic_revision=asic_revision,
+        supported=supported,
+    )
+
+
+def get_device_capabilities(device: int = 0) -> DeviceCapabilities:
+    """Query one HIP device and return its named capabilities."""
+
+    return _capabilities_for_properties(
+        get_device_arch(device),
+        _get_device_asic_revision(device),
+    )
+
+
+__all__ = [
+    "DeviceCapabilities",
+    "DeviceCapability",
+    "get_device_capabilities",
+]

dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/hip_module.py

@@ -153,7 +153,7 @@ def _check(s: int, where: str) -> None:
 # ROCm releases (and the props symbol was versioned to ``...R0600`` in ROCm 6.x), so
 # we keep the raw bytes and let each query read the field it needs rather than mirror
 # the struct. See get_device_arch (gcnArchName) / get_device_name (name).
-_device_props_cache: Dict[int, Optional[bytes]] = {}
+_device_props_cache: Dict[int, bytes] = {}
 
 
 def _ensure_hip_init() -> None:
@@ -204,12 +204,11 @@ def _device_props(device: int = 0) -> Optional[bytes]:
         fn = _b(sym, ctypes.c_void_p, ctypes.c_int)
         try:
             rc = fn(buf, device)
-        except (AttributeError, OSError):
+        except (AttributeError, HipError, OSError):
             continue
         if rc == 0:
             _device_props_cache[device] = buf.raw
             return buf.raw
-    _device_props_cache[device] = None
     return None
 
 
@@ -273,6 +272,11 @@ def get_device_count() -> int:
 # is far more durable than reading a field offset out of the churny hipDeviceProp_t.
 _HIP_ATTR_MULTIPROCESSOR_COUNT = 63
 
+# hipDeviceAttributeAsicRevision. This lives in the AMD-specific portion of the stable
+# ``hipDeviceAttribute_t`` ABI. Querying it through ``hipDeviceGetAttribute`` avoids
+# depending on the ROCm-version-specific offset of ``hipDeviceProp_t::asicRevision``.
+_HIP_ATTR_ASIC_REVISION = 10012
+
 
 def get_device_num_cus(device: int = 0) -> Optional[int]:
     """CU (multiprocessor) count of a HIP device, or None if unqueryable.
@@ -295,6 +299,24 @@ def get_device_num_cus(device: int = 0) -> Optional[int]:
     return int(v.value) if rc == 0 and v.value > 0 else None
 
 
+def _get_device_asic_revision(device: int = 0) -> Optional[int]:
+    """ASIC revision of a HIP device, or ``None`` when it cannot be queried.
+
+    Revision ``0`` is a valid result and must remain distinguishable from a failed
+    query. This private value is used to look up named capabilities. Callers should
+    use that capability API instead of comparing revision numbers themselves.
+    """
+
+    v = ctypes.c_int(-1)
+    try:
+        rc = _hipDeviceGetAttribute(
+            ctypes.byref(v), _HIP_ATTR_ASIC_REVISION, int(device)
+        )
+    except (AttributeError, HipError, OSError):
+        return None
+    return int(v.value) if rc == 0 and v.value >= 0 else None
+
+
 @dataclass
 class Module:
     handle: _HipModuleHandle