diff --git a/installer.py b/installer.py index b5ddca5b7..b5e9ba1ad 100644 --- a/installer.py +++ b/installer.py @@ -678,32 +678,42 @@ def install_rocm_zluda(): return torch_command from modules import rocm - log.info('ROCm: AMD toolkit detected') + amd_gpus = [] + try: + if sys.platform == "win32" and not rocm.is_installed: + amd_gpus = rocm.driver_get_agents() + else: + amd_gpus = rocm.get_agents() + log.info('ROCm: AMD toolkit detected') + except Exception as e: + log.warning(f'ROCm agent enumerator failed: {e}') + #os.environ.setdefault('TENSORFLOW_PACKAGE', 'tensorflow') device = None - try: - amd_gpus = rocm.get_agents() - if len(amd_gpus) == 0: - log.warning('ROCm: no agent was found') + if len(amd_gpus) == 0: + log.warning('ROCm: no agent was found') + else: + log.info(f'ROCm: agents={[gpu.name for gpu in amd_gpus]}') + if args.device_id is None: + index = 0 + for idx, gpu in enumerate(amd_gpus): + index = idx + if not gpu.is_apu: + # although apu was found, there can be a dedicated card. do not break loop. + # if no dedicated card was found, apu will be used. + break + os.environ.setdefault('HIP_VISIBLE_DEVICES', str(index)) + device = amd_gpus[index] else: - log.info(f'ROCm: agents={[gpu.name for gpu in amd_gpus]}') - if args.device_id is None: - index = 0 - for idx, gpu in enumerate(amd_gpus): - index = idx - if not gpu.is_apu: - # although apu was found, there can be a dedicated card. do not break loop. - # if no dedicated card was found, apu will be used. - break - os.environ.setdefault('HIP_VISIBLE_DEVICES', str(index)) - device = amd_gpus[index] - else: - device_id = int(args.device_id) - if device_id < len(amd_gpus): - device = amd_gpus[device_id] - except Exception as e: - log.warning(f'ROCm agent enumerator failed: {e}') + device_id = int(args.device_id) + if device_id < len(amd_gpus): + device = amd_gpus[device_id] + + if sys.platform == "win32" and args.use_rocm and not rocm.is_installed: + check_python(supported_minors=[11, 12, 13], reason='ROCm backend requires a Python version between 3.11 and 3.13') + install(f"rocm rocm-sdk-core --index-url https://rocm.nightlies.amd.com/v2-staging/{device.therock}") + rocm.refresh() msg = f'ROCm: version={rocm.version}' if device is not None: @@ -714,7 +724,7 @@ def install_rocm_zluda(): if args.use_rocm: # TODO install: switch to pytorch source when it becomes available if device is not None and isinstance(rocm.environment, rocm.PythonPackageEnvironment): # TheRock check_python(supported_minors=[11, 12, 13], reason='ROCm backend requires a Python version between 3.11 and 3.13') - torch_command = os.environ.get('TORCH_COMMAND', f'torch torchvision --index-url https://rocm.nightlies.amd.com/v2-staging/{rocm.get_distribution(device)}') + torch_command = os.environ.get('TORCH_COMMAND', f'torch torchvision --index-url https://rocm.nightlies.amd.com/v2-staging/{device.therock}') else: check_python(supported_minors=[12], reason='ROCm Windows preview requires Python version 3.12') torch_command = os.environ.get('TORCH_COMMAND', '--no-cache-dir https://repo.radeon.com/rocm/windows/rocm-rel-6.4.4/torch-2.8.0a0%2Bgitfc14c65-cp312-cp312-win_amd64.whl https://repo.radeon.com/rocm/windows/rocm-rel-6.4.4/torchvision-0.24.0a0%2Bc85f008-cp312-cp312-win_amd64.whl') diff --git a/modules/rocm.py b/modules/rocm.py index 5188409d4..8dfcd9ec8 100644 --- a/modules/rocm.py +++ b/modules/rocm.py @@ -105,6 +105,20 @@ class Agent: self.is_apu = (self.gfx_version & 0xFFF0 == 0x1150) or self.gfx_version in (0x801, 0x902, 0x90c, 0x1013, 0x1033, 0x1035, 0x1036, 0x1103,) self.blaslt_supported = os.path.exists(os.path.join(blaslt_tensile_libpath, f"Kernels.so-000-{name}.hsaco" if sys.platform == "win32" else f"extop_{name}.co")) + @property + def therock(self) -> str: + if (self.gfx_version & 0xFFF0) == 0x1100: + return "gfx110X-dgpu" + if self.gfx_version == 0x1151: + return "gfx1151" + if (self.gfx_version & 0xFFF0) == 0x1200: + return "gfx120X-all" + if (self.gfx_version & 0xFFF0) == 0x940: + return "gfx94X-dcgpu" + if self.gfx_version == 0x950: + return "gfx950-dcgpu" + raise Exception(f"Unsupported GPU architecture: {self.name}") + def get_gfx_version(self) -> Union[str, None]: if self.gfx_version >= 0x1100 and self.gfx_version < 0x1200: return "11.0.0" @@ -204,29 +218,29 @@ def get_flash_attention_command(agent: Agent) -> str: return "--no-build-isolation " + os.environ.get("FLASH_ATTENTION_PACKAGE", default) -def get_distribution(agent: Agent) -> str: - if (agent.gfx_version & 0xFFF0) == 0x1100: - return "gfx110X-dgpu" - if agent.gfx_version == 0x1151: - return "gfx1151" - if (agent.gfx_version & 0xFFF0) == 0x1200: - return "gfx120X-all" - if (agent.gfx_version & 0xFFF0) == 0x940: - return "gfx94X-dcgpu" - if agent.gfx_version == 0x950: - return "gfx950-dcgpu" - raise Exception(f"Unsupported GPU architecture: {agent.name}") - - if sys.platform == "win32": def get_agents() -> List[Agent]: if isinstance(environment, ROCmEnvironment): out = spawn("amdgpu-arch", cwd=os.path.join(environment.path, 'bin')) else: + # Assume that amdgpu-arch is in PATH (venv/Scripts/amdgpu-arch.exe) out = spawn("amdgpu-arch") out = out.strip() return [Agent(x.split(' ')[-1].strip()) for x in out.split("\n")] + def driver_get_agents() -> List[Agent]: + # unsafe and experimental feature + from modules import windows_hip_ffi + hip = windows_hip_ffi.HIP() + count = hip.get_device_count() + agents = [None] * count + for i in range(count): + prop = hip.get_device_properties(i) + name = prop.gcnArchName.decode('utf-8').strip('\x00') + agents[i] = Agent(name) + del hip + return agents + is_wsl: bool = False else: def get_agents() -> List[Agent]: @@ -250,16 +264,16 @@ else: is_wsl: bool = os.environ.get('WSL_DISTRO_NAME', 'unknown' if spawn('wslpath -w /') else None) is not None environment = None -err = None -try: - environment = find() -except Exception as e: - err = e blaslt_tensile_libpath = "" is_installed = False version = None -if environment is not None: - if isinstance(environment, ROCmEnvironment): - blaslt_tensile_libpath = os.environ.get("HIPBLASLT_TENSILE_LIBPATH", os.path.join(environment.path, "bin" if sys.platform == "win32" else "lib", "hipblaslt", "library")) - is_installed = True - version = get_version() + +def refresh(): + global environment, blaslt_tensile_libpath, is_installed, version + environment = find() + if environment is not None: + if isinstance(environment, ROCmEnvironment): + blaslt_tensile_libpath = os.environ.get("HIPBLASLT_TENSILE_LIBPATH", os.path.join(environment.path, "bin" if sys.platform == "win32" else "lib", "hipblaslt", "library")) + is_installed = True + version = get_version() +refresh() diff --git a/modules/windows_hip_ffi.py b/modules/windows_hip_ffi.py new file mode 100644 index 000000000..cc9e4a1e4 --- /dev/null +++ b/modules/windows_hip_ffi.py @@ -0,0 +1,44 @@ +import sys + +if sys.platform == "win32": + import ctypes + import ctypes.wintypes + + class hipDeviceProp(ctypes.Structure): + _fields_ = [ + ('__front__', ctypes.c_byte * 396), + ('gcnArchName', ctypes.c_char * 256), + ('__rear__', ctypes.c_byte * 820) + ] + + class HIP: + def __init__(self): + ctypes.windll.kernel32.LoadLibraryA.restype = ctypes.wintypes.HMODULE + ctypes.windll.kernel32.LoadLibraryA.argtypes = [ctypes.c_char_p] + # amdhip64.dll is a part of AMDGPU drivers + self.handle = ctypes.windll.kernel32.LoadLibraryA(b"amdhip64.dll") + ctypes.windll.kernel32.GetLastError.restype = ctypes.wintypes.DWORD + ctypes.windll.kernel32.GetLastError.argtypes = [] + assert ctypes.windll.kernel32.GetLastError() == 0 + ctypes.windll.kernel32.GetProcAddress.restype = ctypes.c_void_p + ctypes.windll.kernel32.GetProcAddress.argtypes = [ctypes.wintypes.HMODULE, ctypes.c_char_p] + self.hipGetDeviceCount = ctypes.CFUNCTYPE(ctypes.c_int, ctypes.POINTER(ctypes.c_int))( + ctypes.windll.kernel32.GetProcAddress(self.handle, b"hipGetDeviceCount")) + self.hipGetDeviceProperties = ctypes.CFUNCTYPE(ctypes.c_int, ctypes.POINTER(hipDeviceProp), ctypes.c_int)( + ctypes.windll.kernel32.GetProcAddress(self.handle, b"hipGetDeviceProperties")) + + def __del__(self): + #ctypes.windll.kernel32.FreeLibrary.argtypes = [ctypes.wintypes.HMODULE] + #ctypes.windll.kernel32.FreeLibrary(self.handle) + # Hopefully it does not make conflicts with amdhip64_7.dll + pass + + def get_device_count(self): + count = ctypes.c_int() + assert self.hipGetDeviceCount(ctypes.byref(count)) == 0 + return count.value + + def get_device_properties(self, device_id): + prop = hipDeviceProp() + assert self.hipGetDeviceProperties(ctypes.byref(prop), device_id) == 0 + return prop