feat(gpu): add assigning device for HIP and make a string for context

This commit is contained in:
Anthony Berg 2025-07-03 13:15:36 +02:00
parent 58f92c5ce9
commit a9c3a51556

View File

@ -14,7 +14,7 @@ class HIPContext(Context):
Class that manages the HIP context.
"""
def __init__(self, device=0, context_flags=None, use_cache=True, autotuning=False):
def __init__(self, device=None, context_flags=None, use_cache=True, autotuning=False):
"""
Creates a new HIP context.
"""
@ -25,6 +25,11 @@ class HIPContext(Context):
self.logger.info(f"HIP Python version {hip_main.HIP_VERSION_NAME}")
self.logger.info(f"ROCm version {hip_main.ROCM_VERSION_NAME}")
if device is None:
device = 0
hip_check(hip.hipSetDevice(device))
# Device information
props = hip.hipDeviceProp_t()
hip_check(hip.hipGetDeviceProperties(props, device))
@ -49,6 +54,10 @@ class HIPContext(Context):
for prog in self.prog.values():
hip_check(hiprtc.hiprtcDestroyProgram(prog.createRef()))
def __str__(self):
device_handle = hip_check(hip.hipGetDevice())
return f"HIPContext id {device_handle}"
def get_module(self, kernel_filename: str,
function: str,
include_dirs: list[str] = None,