fix/rocm properties export (#487)

* layer/device: fix Device properties type annotation and add ROCMProperties validation

* layer: export ROCMProperties from layer package

* kernels: export ROCMProperties from top-level package

* layer/device: add ROCMProperties type annotation and validation guard

---------

Co-authored-by: GitGlimpse895 <sayakmondal432@gmail.com>

Author: GitGlimpse895

kernels/src/kernels/__init__.py

@@ -16,6 +16,7 @@
     LockedFuncRepository,
     LockedLayerRepository,
     Mode,
+    ROCMProperties,
     kernelize,
     register_kernel_mapping,
     replace_kernel_forward_from_hub,
@@ -42,6 +43,7 @@
     "Benchmark",
     "CUDAProperties",
     "Device",
+    "ROCMProperties",
     "FuncRepository",
     "LayerRepository",
     "LoadedKernel",

kernels/src/kernels/layer/__init__.py

@@ -1,4 +1,4 @@
-from .device import CUDAProperties, Device
+from .device import CUDAProperties, Device, ROCMProperties
 from .func import (
     FuncRepository,
     LocalFuncRepository,
@@ -22,6 +22,7 @@
 __all__ = [
     "CUDAProperties",
     "Device",
+    "ROCMProperties",
     "FuncRepository",
     "LayerRepository",
     "LocalFuncRepository",

kernels/src/kernels/layer/device.py

@@ -141,12 +141,18 @@ class Device:
     """
 
     type: str
-    properties: CUDAProperties | None = None
+    properties: CUDAProperties | ROCMProperties | None = None
 
     def __post_init__(self):
         if self.properties is not None and isinstance(self.properties, CUDAProperties):
             if self.type != "cuda":
                 raise ValueError("CUDAProperties is only supported for 'cuda' devices.")
+        if self.properties is not None and isinstance(self.properties, ROCMProperties):
+            if self.type != "rocm":
+                raise ValueError("ROCMProperties is only supported for 'rocm' devices.")
+        if self.properties is not None and isinstance(self.properties, ROCMProperties):
+            if self.type != "rocm":
+                raise ValueError("ROCMProperties is only supported for 'rocm' devices.")
 
     def __eq__(self, other):
         if not isinstance(other, Device):