diff --git a/StabilityMatrix.Core/Helper/HardwareInfo/GpuInfo.cs b/StabilityMatrix.Core/Helper/HardwareInfo/GpuInfo.cs index a19002dd8..beedf1876 100644 --- a/StabilityMatrix.Core/Helper/HardwareInfo/GpuInfo.cs +++ b/StabilityMatrix.Core/Helper/HardwareInfo/GpuInfo.cs @@ -62,9 +62,9 @@ public bool IsLegacyNvidiaGpu() return ComputeCapabilityValue < 7.5m; } - public bool IsWindowsRocmSupportedGpu() + public bool IsRocmSupportedGpu() { - return WindowsRocmSupport.IsSupportedGpu(this); + return RocmSupport.IsSupportedGpu(this); } public bool IsAmd => Name?.Contains("amd", StringComparison.OrdinalIgnoreCase) ?? false; @@ -83,7 +83,7 @@ public bool IsWindowsRocmSupportedGpu() { // RDNA4 _ when Has("R9700") || Has("R9600") || Has("9070") => "gfx1201", - _ when Has("9060") => "gfx1200", + _ when Has("9060") || Has("9050") => "gfx1200", // RDNA3.5 APUs _ when Has("860M") => "gfx1152", @@ -131,6 +131,10 @@ _ when Has("v520") => "gfx1011", // RDNA1 Navi10 XT _ when Has("5600") || Has("5700") => "gfx1010", + // CDNA (PCI-E slot Instincts) + _ when Has("mi100") => "gfx908", + _ when Has("mi210") => "gfx90a", + // Vega/GCN5 Dedicated GPUs _ when Has("rx vega") || Has("vega 64") || Has("vega 56") || Has("vega frontier") => "gfx900", _ when Has("radeon vii") || HasNoSpace("radeonvii") || Has("pro vii") || HasNoSpace("provii") => diff --git a/StabilityMatrix.Core/Helper/HardwareInfo/HardwareHelper.cs b/StabilityMatrix.Core/Helper/HardwareInfo/HardwareHelper.cs index 3efaa3c33..b8ce71d4b 100644 --- a/StabilityMatrix.Core/Helper/HardwareInfo/HardwareHelper.cs +++ b/StabilityMatrix.Core/Helper/HardwareInfo/HardwareHelper.cs @@ -320,11 +320,11 @@ public static bool HasAmdGpu() return IterGpuInfo().Any(gpu => gpu.IsAmd); } - public static bool HasWindowsRocmSupportedGpu() => IterGpuInfo().Any(WindowsRocmSupport.IsSupportedGpu); + public static bool HasRocmSupportedGpu() => IterGpuInfo().Any(RocmSupport.IsSupportedGpu); - public static GpuInfo? GetWindowsRocmSupportedGpu() + public static GpuInfo? GetRocmSupportedGpu() { - return IterGpuInfo().FirstOrDefault(WindowsRocmSupport.IsSupportedGpu); + return IterGpuInfo().FirstOrDefault(RocmSupport.IsSupportedGpu); } public static bool HasIntelGpu() => IterGpuInfo().Any(gpu => gpu.IsIntel); @@ -334,7 +334,7 @@ public static bool HasAmdGpu() // Set DirectML for default if AMD and Windows public static bool PreferDirectMLOrZluda() => - !HasNvidiaGpu() && HasAmdGpu() && Compat.IsWindows && !HasWindowsRocmSupportedGpu(); + !HasNvidiaGpu() && HasAmdGpu() && Compat.IsWindows && !HasRocmSupportedGpu(); private static readonly Lazy IsMemoryInfoAvailableLazy = new(() => TryGetMemoryInfo(out _)); public static bool IsMemoryInfoAvailable => IsMemoryInfoAvailableLazy.Value; diff --git a/StabilityMatrix.Core/Models/PackageModification/InstallWindowsRocmPackageCommandStep.cs b/StabilityMatrix.Core/Models/PackageModification/InstallRocmPackageCommandStep.cs similarity index 72% rename from StabilityMatrix.Core/Models/PackageModification/InstallWindowsRocmPackageCommandStep.cs rename to StabilityMatrix.Core/Models/PackageModification/InstallRocmPackageCommandStep.cs index 8621f6128..69a7705f3 100644 --- a/StabilityMatrix.Core/Models/PackageModification/InstallWindowsRocmPackageCommandStep.cs +++ b/StabilityMatrix.Core/Models/PackageModification/InstallRocmPackageCommandStep.cs @@ -1,3 +1,5 @@ +using NLog; +using StabilityMatrix.Core.Extensions; using StabilityMatrix.Core.Helper; using StabilityMatrix.Core.Models.FileInterfaces; using StabilityMatrix.Core.Models.Progress; @@ -9,7 +11,7 @@ namespace StabilityMatrix.Core.Models.PackageModification; -public enum WindowsRocmPackageCommandType +public enum RocmPackageCommandType { SageAttention, DevelopmentSdk, @@ -17,13 +19,15 @@ public enum WindowsRocmPackageCommandType FlashAttention, } -public class InstallWindowsRocmPackageCommandStep( +public class InstallRocmPackageCommandStep( IDownloadService downloadService, IPyInstallationManager pyInstallationManager, IPrerequisiteHelper prerequisiteHelper, IRocmPackageHelper rocmPackageHelper ) : IPackageStep { + private static readonly Logger Logger = LogManager.GetCurrentClassLogger(); + private const string BitsAndBytesWheelUrl = "https://github.com/0xDELUXA/bitsandbytes_win_rocm/releases/download/0.50.0.dev0-py3.12-rocm7.15-win_amd64_all/bitsandbytes-0.50.0.dev0-cp312-cp312-win_amd64.whl"; private const string AmdAiterWheelUrl = @@ -44,26 +48,25 @@ IRocmPackageHelper rocmPackageHelper public required InstalledPackage InstalledPackage { get; init; } public required DirectoryPath WorkingDirectory { get; init; } - public required WindowsRocmPackageCommandType CommandType { get; init; } + public required RocmPackageCommandType CommandType { get; init; } public IReadOnlyDictionary? EnvironmentVariables { get; init; } public string ProgressTitle => CommandType switch { - WindowsRocmPackageCommandType.SageAttention => "Installing Windows ROCm SageAttention", - WindowsRocmPackageCommandType.DevelopmentSdk => "Installing Windows ROCm Development SDK", - WindowsRocmPackageCommandType.BitsAndBytes => "Installing Windows ROCm bitsandbytes", - WindowsRocmPackageCommandType.FlashAttention => "Installing Windows ROCm Flash Attention", + RocmPackageCommandType.SageAttention => "Installing Windows ROCm SageAttention", + RocmPackageCommandType.DevelopmentSdk => "Installing ROCm Development SDK", + RocmPackageCommandType.BitsAndBytes => "Installing Windows ROCm bitsandbytes", + RocmPackageCommandType.FlashAttention => "Installing Windows ROCm Flash Attention", _ => "Running Windows ROCm package command", }; public async Task ExecuteAsync(IProgress? progress = null) { - if (!OperatingSystem.IsWindows()) + // The development SDK is published for Linux as well; the compiled extras are Windows-only. + if (!OperatingSystem.IsWindows() && CommandType != RocmPackageCommandType.DevelopmentSdk) { - throw new PlatformNotSupportedException( - "Windows ROCm package commands are only supported on Windows." - ); + throw new PlatformNotSupportedException($"{CommandType} is only available on Windows."); } if (InstalledPackage.FullPath is null) @@ -97,16 +100,16 @@ await pyInstallationManager.GetInstallationAsync(pyVersion).ConfigureAwait(false switch (CommandType) { - case WindowsRocmPackageCommandType.SageAttention: + case RocmPackageCommandType.SageAttention: await ExecuteSageAttentionAsync(venvRunner, progress).ConfigureAwait(false); break; - case WindowsRocmPackageCommandType.DevelopmentSdk: + case RocmPackageCommandType.DevelopmentSdk: await ExecuteDevelopmentSdkAsync(venvRunner, progress).ConfigureAwait(false); break; - case WindowsRocmPackageCommandType.BitsAndBytes: + case RocmPackageCommandType.BitsAndBytes: await ExecuteBitsAndBytesAsync(venvRunner, pyVersion, progress).ConfigureAwait(false); break; - case WindowsRocmPackageCommandType.FlashAttention: + case RocmPackageCommandType.FlashAttention: await ExecuteFlashAttentionAsync(venvRunner, progress).ConfigureAwait(false); break; default: @@ -144,7 +147,70 @@ private async Task ExecuteDevelopmentSdkAsync( ) { EnsureRocmCompatibility(); - await rocmPackageHelper.EnsureWindowsSdkDevelAsync(venvRunner, progress).ConfigureAwait(false); + await rocmPackageHelper + .EnsureRocmSdkDevelAsync(venvRunner, progress, progress.AsProcessOutputHandler()) + .ConfigureAwait(false); + await TryInitializeRocmSdkAsync(venvRunner, progress).ConfigureAwait(false); + } + + /// + /// Runs rocm-sdk init so the installed devel tree is expanded eagerly. Failures are logged and + /// surfaced as a warning rather than thrown, because the SDK also initialises lazily on first use. + /// + private async Task TryInitializeRocmSdkAsync( + IPyVenvRunner venvRunner, + IProgress? progress + ) + { + var sdkPath = venvRunner.RootPath.JoinFile( + Compat.IsWindows ? Path.Combine("Scripts", "rocm-sdk.exe") : Path.Combine("bin", "rocm-sdk") + ); + + if (!sdkPath.Exists) + { + Logger.Warn("rocm-sdk was not found at '{Path}'; skipping initialization.", sdkPath.FullPath); + return; + } + + var environment = venvRunner.EnvironmentVariables.ToDictionary( + pair => pair.Key, + pair => pair.Value, + StringComparer.OrdinalIgnoreCase + ); + + // Keep a stale system ROCm or HIP SDK from taking over the venv's toolchain. + foreach (var key in RocmSupport.ConflictingRocmEnvironmentVariables) + { + environment[key] = ProcessRunner.UnsetEnvironmentVariable; + } + + progress?.Report(new ProgressReport(-1f, "Initializing ROCm SDK...", isIndeterminate: true)); + + using var process = ProcessRunner.StartAnsiProcess( + fileName: sdkPath.FullPath, + arguments: "init", + workingDirectory: venvRunner.RootPath.FullPath, + outputDataReceived: progress.AsProcessOutputHandler(), + environmentVariables: environment + ); + + await process.WaitForExitAsync().ConfigureAwait(false); + await process.WaitUntilOutputEOF().ConfigureAwait(false); + + if (process.ExitCode == 0) + { + Logger.Info("rocm-sdk init completed."); + return; + } + + Logger.Warn("rocm-sdk init exited with {Code}.", process.ExitCode); + progress?.Report( + new ProgressReport( + -1f, + "ROCm SDK initialization reported an error; continuing.", + isIndeterminate: true + ) + ); } private async Task ExecuteSageAttentionAsync( @@ -154,7 +220,9 @@ private async Task ExecuteSageAttentionAsync( { EnsureRocmCompatibility(); await EnsureVcBuildToolsAsync(progress).ConfigureAwait(false); - await rocmPackageHelper.EnsureWindowsSdkDevelAsync(venvRunner, progress).ConfigureAwait(false); + await rocmPackageHelper + .EnsureRocmSdkDevelAsync(venvRunner, progress, progress.AsProcessOutputHandler()) + .ConfigureAwait(false); progress?.Report( new ProgressReport( diff --git a/StabilityMatrix.Core/Models/Packages/BaseGitPackage.cs b/StabilityMatrix.Core/Models/Packages/BaseGitPackage.cs index 67c877bb0..948815b58 100644 --- a/StabilityMatrix.Core/Models/Packages/BaseGitPackage.cs +++ b/StabilityMatrix.Core/Models/Packages/BaseGitPackage.cs @@ -84,18 +84,13 @@ IPipWheelService pipWheelService PipWheelService = pipWheelService; } - protected bool HasWindowsRocmSupport(IRocmPackageHelper rocmPackageHelper) + protected bool HasRocmSupport(IRocmPackageHelper rocmPackageHelper) { - return GetWindowsRocmCompatibility(rocmPackageHelper).IsCompatible; + return GetRocmCompatibility(rocmPackageHelper).IsCompatible; } - protected RocmCompatibilityResult GetWindowsRocmCompatibility(IRocmPackageHelper rocmPackageHelper) + protected RocmCompatibilityResult GetRocmCompatibility(IRocmPackageHelper rocmPackageHelper) { - if (!Compat.IsWindows) - { - return new RocmCompatibilityResult { IsCompatible = false }; - } - return rocmPackageHelper.GetCompatibility(); } diff --git a/StabilityMatrix.Core/Models/Packages/ComfyUI.cs b/StabilityMatrix.Core/Models/Packages/ComfyUI.cs index 65142f1d0..cc550f051 100644 --- a/StabilityMatrix.Core/Models/Packages/ComfyUI.cs +++ b/StabilityMatrix.Core/Models/Packages/ComfyUI.cs @@ -1,4 +1,4 @@ -using System.Collections.Immutable; +using System.Collections.Immutable; using System.Text.Json; using System.Text.RegularExpressions; using Injectio.Attributes; @@ -316,9 +316,7 @@ private static DirectoryPath GetSharedWorkflowsLinkDir(DirectoryPath installDire Name = "Enable DirectML", Type = LaunchOptionType.Bool, InitialValue = - !HasWindowsRocmSupport() - && HardwareHelper.PreferDirectMLOrZluda() - && this is not ComfyZluda, + !HasRocmSupport() && HardwareHelper.PreferDirectMLOrZluda() && this is not ComfyZluda, Options = ["--directml"], }, new() @@ -415,7 +413,7 @@ public override List GetExtraCommands() ); } - if (Compat.IsWindows && HasWindowsRocmSupport()) + if (Compat.IsWindows && HasRocmSupport()) { commands.Add( new ExtraPackageCommand @@ -430,18 +428,7 @@ public override List GetExtraCommands() { CommandName = "Install Flash Attention (ROCm)", Command = InstallWindowsRocmFlashAttention, - IsVisible = _ => - WindowsRocmSupport.IsLegacyArchitecture( - GetWindowsRocmCompatibility().ResolvedGfxArch - ), - } - ); - - commands.Add( - new ExtraPackageCommand - { - CommandName = "Install ROCm Development SDK", - Command = InstallWindowsRocmDevelopmentSdk, + IsVisible = _ => RocmSupport.IsLegacyArchitecture(GetRocmCompatibility().ResolvedGfxArch), } ); @@ -461,6 +448,17 @@ public override List GetExtraCommands() ); } + if (HasRocmSupport()) + { + commands.Add( + new ExtraPackageCommand + { + CommandName = "Install ROCm Development SDK", + Command = InstallRocmDevelopmentSdk, + } + ); + } + if (!Compat.IsMacOS && SettingsManager.Settings.PreferredGpu?.ComputeCapabilityValue is >= 7.5m) { commands.Add( @@ -495,9 +493,9 @@ public override async Task InstallPackage( ?? HardwareHelper.HasLegacyNvidiaGpu() ); - if (Compat.IsWindows && torchIndex == TorchIndex.Rocm && HasWindowsRocmSupport()) + if (torchIndex == TorchIndex.Rocm && HasRocmSupport()) { - var config = rocmPackageHelper.BuildWindowsNativeInstallConfig(ComfyWindowsRocmProfile.Default); + var config = rocmPackageHelper.BuildRocmInstallConfig(ComfyRocmProfile.Default); await StandardPipInstallProcessAsync( venvRunner, @@ -511,10 +509,10 @@ await StandardPipInstallProcessAsync( .ConfigureAwait(false); await rocmPackageHelper - .InstallWindowsNativeTorchAsync( + .InstallRocmTorchAsync( venvRunner, installedPackage, - ComfyWindowsRocmProfile.Default, + ComfyRocmProfile.Default, progress, onConsoleOutput, cancellationToken @@ -574,7 +572,7 @@ internal PipInstallConfig GetStandardInstallConfig(PyVersion pythonVersion, bool ExtraPipArgs = pythonVersion < new PyVersion(3, 13, 0) ? ["numpy<2"] : [], TorchaudioVersion = " ", // Request torchaudio without a specific version CudaIndex = isLegacyNvidia ? "cu126" : "cu130", - RocmIndex = "rocm7.2", + RocmIndex = "rocm7.14", XpuIndex = "xpu", UpgradePackages = true, PostInstallPipArgs = ["typing-extensions>=4.15.0"], @@ -649,7 +647,7 @@ private void EmitWindowsRocmLaunchNotice( ) { var torchIndex = installedPackage.PreferredTorchIndex ?? GetRecommendedTorchVersion(); - foreach (var line in rocmPackageHelper.GetWindowsLaunchNoticeLines(torchIndex)) + foreach (var line in rocmPackageHelper.GetLaunchNoticeLines(torchIndex)) { onConsoleOutput?.Invoke(ProcessOutput.FromStdOutLine($"{line}{Environment.NewLine}")); } @@ -687,7 +685,7 @@ public override TorchIndex GetRecommendedTorchVersion() { var preferRocm = (Compat.IsLinux && (SettingsManager.Settings.PreferredGpu?.IsAmd ?? HardwareHelper.PreferRocm())) - || HasWindowsRocmSupport(); + || HasRocmSupport(); if (AvailableTorchIndices.Contains(TorchIndex.Rocm) && preferRocm) { @@ -697,26 +695,26 @@ public override TorchIndex GetRecommendedTorchVersion() return base.GetRecommendedTorchVersion(); } - /// Uses the shared ROCm helper for Windows ROCm eligibility checks so ComfyUI does not maintain its own support matrix. - private bool HasWindowsRocmSupport() + /// Uses the shared ROCm helper for ROCm eligibility checks so ComfyUI does not maintain its own support matrix. + private bool HasRocmSupport() { - return HasWindowsRocmSupport(rocmPackageHelper); + return HasRocmSupport(rocmPackageHelper); } - private RocmCompatibilityResult GetWindowsRocmCompatibility() + private RocmCompatibilityResult GetRocmCompatibility() { - return GetWindowsRocmCompatibility(rocmPackageHelper); + return GetRocmCompatibility(rocmPackageHelper); } - /// Defaults legacy Windows ROCm GPUs to quad cross-attention because PyTorch cross-attention is considerably slower + /// Defaults legacy ROCm GPUs to quad cross-attention because PyTorch cross-attention is considerably slower /// and not as supported on older AMD architectures. private bool DefaultToQuadCrossAttention() { - var compatibility = GetWindowsRocmCompatibility(); + var compatibility = GetRocmCompatibility(); if (!compatibility.IsCompatible) return false; - return WindowsRocmSupport.PreferLegacyAttentionFallback(compatibility.ResolvedGfxArch); + return RocmSupport.PreferLegacyAttentionFallback(compatibility.ResolvedGfxArch); } public override IPackageExtensionManager ExtensionManager => @@ -973,7 +971,7 @@ internal async Task UpdateSageAttentionAsync( IProgress? progress = null ) { - if (Compat.IsWindows && torchIndex == TorchIndex.Rocm && HasWindowsRocmSupport()) + if (torchIndex == TorchIndex.Rocm && HasRocmSupport()) return; try @@ -1054,9 +1052,9 @@ await PipWheelService private async Task InstallWindowsRocmSageAttention(InstalledPackage? installedPackage) { - var succeeded = await RunWindowsRocmPackageCommandAsync( + var succeeded = await RunRocmPackageCommandAsync( installedPackage, - WindowsRocmPackageCommandType.SageAttention, + RocmPackageCommandType.SageAttention, "Windows ROCm SageAttention installed successfully", includeEnvironmentVariables: true ) @@ -1068,12 +1066,12 @@ private async Task InstallWindowsRocmSageAttention(InstalledPackage? installedPa await EnableSageAttentionAsync(installedPackage).ConfigureAwait(false); } - private async Task InstallWindowsRocmDevelopmentSdk(InstalledPackage? installedPackage) + private async Task InstallRocmDevelopmentSdk(InstalledPackage? installedPackage) { - await RunWindowsRocmPackageCommandAsync( + await RunRocmPackageCommandAsync( installedPackage, - WindowsRocmPackageCommandType.DevelopmentSdk, - "Windows ROCm Development SDK installed successfully", + RocmPackageCommandType.DevelopmentSdk, + "ROCm Development SDK installed successfully", includeEnvironmentVariables: false ) .ConfigureAwait(false); @@ -1081,9 +1079,9 @@ await RunWindowsRocmPackageCommandAsync( private async Task InstallWindowsRocmBitsAndBytes(InstalledPackage? installedPackage) { - await RunWindowsRocmPackageCommandAsync( + await RunRocmPackageCommandAsync( installedPackage, - WindowsRocmPackageCommandType.BitsAndBytes, + RocmPackageCommandType.BitsAndBytes, "Windows ROCm bitsandbytes installed successfully", includeEnvironmentVariables: true ) @@ -1092,18 +1090,18 @@ await RunWindowsRocmPackageCommandAsync( private async Task InstallWindowsRocmFlashAttention(InstalledPackage? installedPackage) { - await RunWindowsRocmPackageCommandAsync( + await RunRocmPackageCommandAsync( installedPackage, - WindowsRocmPackageCommandType.FlashAttention, + RocmPackageCommandType.FlashAttention, "Windows ROCm Flash Attention installed successfully", includeEnvironmentVariables: true ) .ConfigureAwait(false); } - private async Task RunWindowsRocmPackageCommandAsync( + private async Task RunRocmPackageCommandAsync( InstalledPackage? installedPackage, - WindowsRocmPackageCommandType commandType, + RocmPackageCommandType commandType, string completionMessage, bool includeEnvironmentVariables ) @@ -1130,7 +1128,7 @@ bool includeEnvironmentVariables await runner .ExecuteSteps( [ - new InstallWindowsRocmPackageCommandStep( + new InstallRocmPackageCommandStep( downloadService, pyInstallationManager, prerequisiteHelper, @@ -1231,8 +1229,6 @@ private ImmutableDictionary GetEnvVars( InstalledPackage installedPackage ) { - // if we're not on windows or we don't have a windows rocm gpu, return original env - var hasRocmGpu = HasWindowsRocmSupport(); var selectedTorchIndex = installedPackage.PreferredTorchIndex ?? GetRecommendedTorchVersion(); // Add FFmpeg to PATH if it's installed (optional - for video processing) @@ -1247,10 +1243,10 @@ InstalledPackage installedPackage } } - if (!Compat.IsWindows || !hasRocmGpu || selectedTorchIndex != TorchIndex.Rocm) + if (!rocmPackageHelper.ShouldApplyLaunchEnvironment(selectedTorchIndex)) return env; - var rocmEnvironment = rocmPackageHelper.BuildLaunchEnvironment(ComfyWindowsRocmProfile.Default); + var rocmEnvironment = rocmPackageHelper.BuildLaunchEnvironment(ComfyRocmProfile.Default); return env.SetItems(rocmEnvironment); } diff --git a/StabilityMatrix.Core/Models/Packages/ForgeNeo.cs b/StabilityMatrix.Core/Models/Packages/ForgeNeo.cs index 0dc0c29eb..7c8904ff9 100644 --- a/StabilityMatrix.Core/Models/Packages/ForgeNeo.cs +++ b/StabilityMatrix.Core/Models/Packages/ForgeNeo.cs @@ -49,7 +49,7 @@ IRocmPackageHelper rocmPackageHelper /// Windows installs use AMD repo wheels; Linux uses the upstream PyTorch ROCm index. /// private RocmPackageProfile ActiveRocmProfile => - Compat.IsWindows ? ForgeNeoWindowsRocmProfile.Default : LinuxRocmProfile; + Compat.IsWindows ? ForgeNeoRocmProfile.Default : LinuxRocmProfile; public override string Name => "forge-neo"; public override string DisplayName { get; set; } = "Stable Diffusion WebUI Forge - Neo"; @@ -103,7 +103,7 @@ public override List LaunchOptions } ); - ForgeNeoWindowsRocmProfile.Default.ApplyWindowsRocmLaunchDefaults(options, rocmPackageHelper); + ForgeNeoRocmProfile.Default.ApplyRocmLaunchDefaults(options, rocmPackageHelper); return options; } @@ -144,10 +144,10 @@ CancellationToken cancellationToken if (Compat.IsWindows) { - // Install the AMD multi-arch torch build before launch.py runs, so its is_installed("torch") + // Install the helper-managed ROCm torch build before launch.py runs, so its is_installed("torch") // gate skips the default NVIDIA CUDA torch install. await rocmPackageHelper - .InstallWindowsNativeTorchAsync( + .InstallRocmTorchAsync( venvRunner, installedPackage, profile, @@ -214,8 +214,8 @@ InstalledPackage installedPackage protected override IReadOnlyList GetLaunchNoticeLines(InstalledPackage installedPackage) { var selectedTorchIndex = installedPackage.PreferredTorchIndex ?? GetRecommendedTorchVersion(); - // The helper returns an empty list on non-Windows platforms. - return rocmPackageHelper.GetWindowsLaunchNoticeLines(selectedTorchIndex); + // Launch notices only apply where the helper manages a native ROCm install. + return rocmPackageHelper.GetLaunchNoticeLines(selectedTorchIndex); } public override string Blurb => diff --git a/StabilityMatrix.Core/Models/Packages/InvokeAI.cs b/StabilityMatrix.Core/Models/Packages/InvokeAI.cs index 870f598ff..ade44a62e 100644 --- a/StabilityMatrix.Core/Models/Packages/InvokeAI.cs +++ b/StabilityMatrix.Core/Models/Packages/InvokeAI.cs @@ -232,7 +232,7 @@ public override async Task InstallPackage( progress?.Report(new ProgressReport(-1f, "Installing Package", isIndeterminate: true)); var torchVersion = options.PythonOptions.TorchIndex ?? GetRecommendedTorchVersion(); - var isSupportedWindowsRocmInstall = + var isSupportedRocmInstall = Compat.IsWindows && torchVersion == TorchIndex.Rocm && rocmPackageHelper.GetCompatibility().IsCompatible; @@ -268,7 +268,7 @@ public override async Task InstallPackage( torchVersion.ToString().ToLowerInvariant() ]?.GetValue(); - if (isSupportedWindowsRocmInstall) + if (isSupportedRocmInstall) { index = null; } @@ -286,9 +286,9 @@ public override async Task InstallPackage( await venvRunner.PipInstall(invokeInstallArgs, onConsoleOutput).ConfigureAwait(false); - if (isSupportedWindowsRocmInstall) + if (isSupportedRocmInstall) { - var installProfile = InvokeAiWindowsRocmProfile.CreateInstallProfile( + var installProfile = InvokeAiRocmProfile.CreateInstallProfile( options.PythonOptions.PythonVersion ?? RecommendedPythonVersion ); @@ -297,7 +297,7 @@ public override async Task InstallPackage( ); await rocmPackageHelper - .InstallWindowsNativeTorchAsync( + .InstallRocmTorchAsync( venvRunner, installedPackage, installProfile, @@ -662,17 +662,18 @@ InstalledPackage installedPackage env = GetEnvVars(env, installPath); var selectedTorchIndex = installedPackage.PreferredTorchIndex ?? GetRecommendedTorchVersion(); - if (!rocmPackageHelper.ShouldApplyWindowsLaunchEnvironment(selectedTorchIndex)) + // Gated to Windows until this package installs through the helper on Linux as well. + if (!Compat.IsWindows || !rocmPackageHelper.ShouldApplyLaunchEnvironment(selectedTorchIndex)) { return env; } - return env.SetItems(rocmPackageHelper.BuildLaunchEnvironment(InvokeAiWindowsRocmProfile.Default)); + return env.SetItems(rocmPackageHelper.BuildLaunchEnvironment(InvokeAiRocmProfile.Default)); } private IReadOnlyList GetLaunchNoticeLines(InstalledPackage installedPackage) { var selectedTorchIndex = installedPackage.PreferredTorchIndex ?? GetRecommendedTorchVersion(); - return rocmPackageHelper.GetWindowsLaunchNoticeLines(selectedTorchIndex); + return rocmPackageHelper.GetLaunchNoticeLines(selectedTorchIndex); } } diff --git a/StabilityMatrix.Core/Models/Packages/OneTrainer.cs b/StabilityMatrix.Core/Models/Packages/OneTrainer.cs index f6b69278c..6c862fcf9 100644 --- a/StabilityMatrix.Core/Models/Packages/OneTrainer.cs +++ b/StabilityMatrix.Core/Models/Packages/OneTrainer.cs @@ -93,9 +93,7 @@ public override async Task InstallPackage( if (isWindowsRocm) { - var config = rocmPackageHelper.BuildWindowsNativeInstallConfig( - OneTrainerWindowsRocmProfile.Default - ); + var config = rocmPackageHelper.BuildRocmInstallConfig(OneTrainerRocmProfile.Default); await StandardPipInstallProcessAsync( venvRunner, @@ -116,10 +114,10 @@ config with ); await rocmPackageHelper - .InstallWindowsNativeTorchAsync( + .InstallRocmTorchAsync( venvRunner, installedPackage, - OneTrainerWindowsRocmProfile.CreateInstallProfile( + OneTrainerRocmProfile.CreateInstallProfile( options.PythonOptions.PythonVersion ?? RecommendedPythonVersion ), progress, @@ -174,10 +172,10 @@ await SetupVenv(installLocation, pythonVersion: PyVersion.Parse(installedPackage var selectedTorchIndex = installedPackage.PreferredTorchIndex ?? GetRecommendedTorchVersion(); - if (rocmPackageHelper.ShouldApplyWindowsLaunchEnvironment(selectedTorchIndex)) + if (rocmPackageHelper.ShouldApplyLaunchEnvironment(selectedTorchIndex)) { VenvRunner.UpdateEnvironmentVariables(env => - env.SetItems(rocmPackageHelper.BuildLaunchEnvironment(OneTrainerWindowsRocmProfile.Default)) + env.SetItems(rocmPackageHelper.BuildLaunchEnvironment(OneTrainerRocmProfile.Default)) ); } @@ -197,12 +195,12 @@ private IReadOnlyList GetLaunchNoticeLines(InstalledPackage installedPac { var selectedTorchIndex = installedPackage.PreferredTorchIndex ?? GetRecommendedTorchVersion(); - if (!rocmPackageHelper.ShouldApplyWindowsLaunchEnvironment(selectedTorchIndex)) + if (!rocmPackageHelper.ShouldApplyLaunchEnvironment(selectedTorchIndex)) { return []; } - return [.. rocmPackageHelper.GetWindowsLaunchNoticeLines(selectedTorchIndex)]; + return [.. rocmPackageHelper.GetLaunchNoticeLines(selectedTorchIndex)]; } public override List LaunchOptions => [LaunchOptionDefinition.Extras]; diff --git a/StabilityMatrix.Core/Models/Packages/Reforge.cs b/StabilityMatrix.Core/Models/Packages/Reforge.cs index f31cb1dd8..30b7df825 100644 --- a/StabilityMatrix.Core/Models/Packages/Reforge.cs +++ b/StabilityMatrix.Core/Models/Packages/Reforge.cs @@ -58,10 +58,7 @@ public override List LaunchOptions // Adjust inherited launch defaults for the Windows ROCm path before inserting reForge-specific options. // Makes the reForge-specific attention options visible in the UI and leaves them unset by default // for non-Windows-ROCm installs so reForge can keep using its normal internal attention selection. - ReforgeWindowsRocmProfile.Default.ApplyWindowsRocmLaunchDefaults( - baseLaunchOptions, - rocmPackageHelper - ); + ReforgeRocmProfile.Default.ApplyRocmLaunchDefaults(baseLaunchOptions, rocmPackageHelper); baseLaunchOptions.Insert( extrasIndex >= 0 ? extrasIndex : baseLaunchOptions.Count, @@ -69,7 +66,7 @@ public override List LaunchOptions { Name = "Cross Attention Method", Type = LaunchOptionType.Bool, - InitialValue = ReforgeWindowsRocmProfile.Default.GetPreferredCrossAttentionArgument( + InitialValue = ReforgeRocmProfile.Default.GetPreferredCrossAttentionArgument( rocmPackageHelper ), Options = ["--attention-split", "--attention-quad", "--attention-pytorch"], @@ -106,7 +103,7 @@ public override async Task InstallPackage( { var torchIndex = options.PythonOptions.TorchIndex ?? GetRecommendedTorchVersion(); - if (!rocmPackageHelper.ShouldApplyWindowsLaunchEnvironment(torchIndex)) + if (!rocmPackageHelper.ShouldApplyLaunchEnvironment(torchIndex)) { await base.InstallPackage( installLocation, @@ -127,8 +124,8 @@ await base.InstallPackage( ) .ConfigureAwait(false); - var profile = ReforgeWindowsRocmProfile.CreateProfile(GetRequirementsPaths(installLocation)); - var config = rocmPackageHelper.BuildWindowsNativeInstallConfig(profile); + var profile = ReforgeRocmProfile.CreateProfile(GetRequirementsPaths(installLocation)); + var config = rocmPackageHelper.BuildRocmInstallConfig(profile); await StandardPipInstallProcessAsync( venvRunner, @@ -142,7 +139,7 @@ await StandardPipInstallProcessAsync( .ConfigureAwait(false); await rocmPackageHelper - .InstallWindowsNativeTorchAsync( + .InstallRocmTorchAsync( venvRunner, installedPackage, profile, @@ -173,17 +170,17 @@ InstalledPackage installedPackage env = base.GetEnvVars(env, installedPackage); env = env.SetItem("STABLE_DIFFUSION_REPO", StableDiffusionRepoOverride); - if (!rocmPackageHelper.ShouldApplyWindowsLaunchEnvironment(selectedTorchIndex)) + if (!rocmPackageHelper.ShouldApplyLaunchEnvironment(selectedTorchIndex)) { return env; } - return env.SetItems(rocmPackageHelper.BuildLaunchEnvironment(ReforgeWindowsRocmProfile.Default)); + return env.SetItems(rocmPackageHelper.BuildLaunchEnvironment(ReforgeRocmProfile.Default)); } protected override IReadOnlyList GetLaunchNoticeLines(InstalledPackage installedPackage) { var selectedTorchIndex = installedPackage.PreferredTorchIndex ?? GetRecommendedTorchVersion(); - return rocmPackageHelper.GetWindowsLaunchNoticeLines(selectedTorchIndex); + return rocmPackageHelper.GetLaunchNoticeLines(selectedTorchIndex); } } diff --git a/StabilityMatrix.Core/Models/Packages/StableSwarm.cs b/StabilityMatrix.Core/Models/Packages/StableSwarm.cs index b9d43abbb..49dcb5c06 100644 --- a/StabilityMatrix.Core/Models/Packages/StableSwarm.cs +++ b/StabilityMatrix.Core/Models/Packages/StableSwarm.cs @@ -593,7 +593,7 @@ x.PackageName is nameof(ComfyUI) or "ComfyUI-Zluda" return new Dictionary(); } - return rocmPackageHelper.BuildLaunchEnvironment(ComfyWindowsRocmProfile.Default); + return rocmPackageHelper.BuildLaunchEnvironment(ComfyRocmProfile.Default); } /// diff --git a/StabilityMatrix.Core/Models/Packages/Wan2GP.cs b/StabilityMatrix.Core/Models/Packages/Wan2GP.cs index 5ef817459..99fc2f7cf 100644 --- a/StabilityMatrix.Core/Models/Packages/Wan2GP.cs +++ b/StabilityMatrix.Core/Models/Packages/Wan2GP.cs @@ -67,8 +67,7 @@ IRocmPackageHelper rocmPackageHelper public override IEnumerable AvailableTorchIndices => [TorchIndex.Cuda, TorchIndex.Rocm]; public override bool IsCompatible => - HardwareHelper.HasNvidiaGpu() - || (Compat.IsWindows ? HasWindowsRocmSupport() : HardwareHelper.HasAmdGpu()); + HardwareHelper.HasNvidiaGpu() || (Compat.IsWindows ? HasRocmSupport() : HardwareHelper.HasAmdGpu()); public override string MainBranch => "main"; public override bool ShouldIgnoreReleases => true; @@ -95,14 +94,14 @@ IRocmPackageHelper rocmPackageHelper private bool IsLegacyNvidia => SettingsManager.Settings.PreferredGpu?.IsLegacyNvidiaGpu() ?? HardwareHelper.HasLegacyNvidiaGpu(); - private bool HasWindowsRocmSupport() + private bool HasRocmSupport() { - return HasWindowsRocmSupport(rocmPackageHelper); + return HasRocmSupport(rocmPackageHelper); } - private RocmCompatibilityResult GetWindowsRocmCompatibility() + private RocmCompatibilityResult GetRocmCompatibility() { - return GetWindowsRocmCompatibility(rocmPackageHelper); + return GetRocmCompatibility(rocmPackageHelper); } /// @@ -234,7 +233,7 @@ public override TorchIndex GetRecommendedTorchVersion() { // Check for AMD ROCm support (Windows or Linux) var preferRocm = - (Compat.IsWindows && HasWindowsRocmSupport()) + (Compat.IsWindows && HasRocmSupport()) || ( Compat.IsLinux && (SettingsManager.Settings.PreferredGpu?.IsAmd ?? HardwareHelper.PreferRocm()) @@ -493,7 +492,7 @@ CancellationToken cancellationToken { if (Compat.IsWindows) { - var config = rocmPackageHelper.BuildWindowsNativeInstallConfig(Wan2GpWindowsRocmProfile.Default); + var config = rocmPackageHelper.BuildRocmInstallConfig(Wan2GpRocmProfile.Default); await StandardPipInstallProcessAsync( venvRunner, @@ -507,10 +506,10 @@ await StandardPipInstallProcessAsync( .ConfigureAwait(false); await rocmPackageHelper - .InstallWindowsNativeTorchAsync( + .InstallRocmTorchAsync( venvRunner, installedPackage, - Wan2GpWindowsRocmProfile.Default, + Wan2GpRocmProfile.Default, progress, onConsoleOutput, cancellationToken @@ -559,9 +558,9 @@ await SetupVenv(installLocation, pythonVersion: PyVersion.Parse(installedPackage var selectedTorchIndex = installedPackage.PreferredTorchIndex ?? GetRecommendedTorchVersion(); - if (Compat.IsWindows && selectedTorchIndex == TorchIndex.Rocm && HasWindowsRocmSupport()) + if (Compat.IsWindows && selectedTorchIndex == TorchIndex.Rocm && HasRocmSupport()) { - var rocmEnvironment = rocmPackageHelper.BuildLaunchEnvironment(Wan2GpWindowsRocmProfile.Default); + var rocmEnvironment = rocmPackageHelper.BuildLaunchEnvironment(Wan2GpRocmProfile.Default); VenvRunner.UpdateEnvironmentVariables(env => env.SetItems(rocmEnvironment)); } diff --git a/StabilityMatrix.Core/Models/Rocm/ComfyWindowsRocmProfile.cs b/StabilityMatrix.Core/Models/Rocm/ComfyRocmProfile.cs similarity index 65% rename from StabilityMatrix.Core/Models/Rocm/ComfyWindowsRocmProfile.cs rename to StabilityMatrix.Core/Models/Rocm/ComfyRocmProfile.cs index 70cddbbb9..b6e955548 100644 --- a/StabilityMatrix.Core/Models/Rocm/ComfyWindowsRocmProfile.cs +++ b/StabilityMatrix.Core/Models/Rocm/ComfyRocmProfile.cs @@ -1,13 +1,13 @@ -using StabilityMatrix.Core.Models.Packages; +using StabilityMatrix.Core.Models.Packages; namespace StabilityMatrix.Core.Models.Rocm; /// -/// Shared Windows ROCm profile for Comfy backends launched either directly by Stability Matrix or indirectly via SwarmUI. +/// Shared ROCm profile for Comfy backends launched either directly by Stability Matrix or indirectly via SwarmUI. /// -public class ComfyWindowsRocmProfile : RocmPackageProfile +public class ComfyRocmProfile : RocmPackageProfile { - public ComfyWindowsRocmProfile() + public ComfyRocmProfile() { InstallConfig = new PipInstallConfig { @@ -22,10 +22,10 @@ public ComfyWindowsRocmProfile() private IReadOnlyDictionary BuildEnvironment(RocmRuntimeContext runtimeContext) { - return WindowsRocmSupport.IsModernArchitecture(runtimeContext.RuntimeGfxArch) + return RocmSupport.IsModernArchitecture(runtimeContext.RuntimeGfxArch) ? new Dictionary { ["COMFYUI_ENABLE_MIOPEN"] = "1" } : new Dictionary(); } - public static RocmPackageProfile Default { get; } = new ComfyWindowsRocmProfile(); + public static RocmPackageProfile Default { get; } = new ComfyRocmProfile(); } diff --git a/StabilityMatrix.Core/Models/Rocm/ForgeNeoWindowsRocmProfile.cs b/StabilityMatrix.Core/Models/Rocm/ForgeNeoRocmProfile.cs similarity index 82% rename from StabilityMatrix.Core/Models/Rocm/ForgeNeoWindowsRocmProfile.cs rename to StabilityMatrix.Core/Models/Rocm/ForgeNeoRocmProfile.cs index 69d5e56b3..2a14b2a5b 100644 --- a/StabilityMatrix.Core/Models/Rocm/ForgeNeoWindowsRocmProfile.cs +++ b/StabilityMatrix.Core/Models/Rocm/ForgeNeoRocmProfile.cs @@ -1,15 +1,15 @@ -using StabilityMatrix.Core.Helper; +using StabilityMatrix.Core.Helper; using StabilityMatrix.Core.Models.Packages; using StabilityMatrix.Core.Services.Rocm; namespace StabilityMatrix.Core.Models.Rocm; /// -/// Windows ROCm profile for Forge Neo. +/// ROCm profile for Forge Neo. /// Forge Neo has no native AMD support, so the helper owns the ROCm torch install and the launch /// environment, while Forge Neo's own launch.py still installs the rest of its dependencies. /// -public class ForgeNeoWindowsRocmProfile : RocmPackageProfile +public class ForgeNeoRocmProfile : RocmPackageProfile { private const string DisableSmartMemoryOptionName = "Disable Smart Memory"; private const string UsePytorchCrossAttentionOptionName = "Use PyTorch Cross Attention"; @@ -18,7 +18,7 @@ public class ForgeNeoWindowsRocmProfile : RocmPackageProfile private static readonly string[] EnableLaunchOptionNames = ["Pin Shared Memory", "CUDA Stream"]; private static readonly string[] DisableLaunchOptionNames = ["CUDA Malloc"]; - public ForgeNeoWindowsRocmProfile() + public ForgeNeoRocmProfile() { InstallConfig = new PipInstallConfig { @@ -31,14 +31,13 @@ public ForgeNeoWindowsRocmProfile() ExtraEnvironmentFactory = BuildEnvironment; } - public static ForgeNeoWindowsRocmProfile Default { get; } = new ForgeNeoWindowsRocmProfile(); + public static ForgeNeoRocmProfile Default { get; } = new ForgeNeoRocmProfile(); private static IReadOnlyDictionary BuildEnvironment(RocmRuntimeContext runtimeContext) { - var indexUrl = WindowsRocmSupport.GetMultiArchPythonPackageIndexUrl(runtimeContext.RuntimeGfxArch); - var deviceExtra = WindowsRocmSupport.TryGetMultiArchDeviceExtra(runtimeContext.RuntimeGfxArch); + var target = RocmSupport.ResolveTarget(runtimeContext.RuntimeGfxArch); - if (string.IsNullOrWhiteSpace(deviceExtra)) + if (string.IsNullOrWhiteSpace(target.DeviceExtra)) { return new Dictionary(); } @@ -46,17 +45,17 @@ private static IReadOnlyDictionary BuildEnvironment(RocmRuntimeC // Persisted so that if Forge Neo re-runs its own torch install step it targets AMD's index // instead of pulling the NVIDIA CUDA build from PyTorch. var torchCommand = - $"pip install --index-url {indexUrl} " - + $"\"torch[{deviceExtra}]\" \"torchvision[{deviceExtra}]\" torchaudio \"rocm[devel]\""; + $"pip install --index-url {target.IndexUrl} " + + $"\"torch[{target.DeviceExtra}]\" \"torchvision[{target.DeviceExtra}]\" torchaudio \"rocm[devel]\""; return new Dictionary { - ["TORCH_INDEX_URL"] = indexUrl, + ["TORCH_INDEX_URL"] = target.IndexUrl, ["TORCH_COMMAND"] = torchCommand, }; } - public override void ApplyWindowsRocmLaunchDefaults( + public override void ApplyRocmLaunchDefaults( List launchOptions, IRocmPackageHelper rocmPackageHelper ) diff --git a/StabilityMatrix.Core/Models/Rocm/InvokeAiWindowsRocmProfile.cs b/StabilityMatrix.Core/Models/Rocm/InvokeAiRocmProfile.cs similarity index 87% rename from StabilityMatrix.Core/Models/Rocm/InvokeAiWindowsRocmProfile.cs rename to StabilityMatrix.Core/Models/Rocm/InvokeAiRocmProfile.cs index 23dae333b..503b43fca 100644 --- a/StabilityMatrix.Core/Models/Rocm/InvokeAiWindowsRocmProfile.cs +++ b/StabilityMatrix.Core/Models/Rocm/InvokeAiRocmProfile.cs @@ -1,14 +1,14 @@ -using StabilityMatrix.Core.Models.Packages; +using StabilityMatrix.Core.Models.Packages; using StabilityMatrix.Core.Python; namespace StabilityMatrix.Core.Models.Rocm; /// -/// Shared Windows ROCm profile for InvokeAI. +/// Shared ROCm profile for InvokeAI. /// -public class InvokeAiWindowsRocmProfile : RocmPackageProfile +public class InvokeAiRocmProfile : RocmPackageProfile { - public static RocmPackageProfile Default { get; } = new InvokeAiWindowsRocmProfile(); + public static RocmPackageProfile Default { get; } = new InvokeAiRocmProfile(); //Restores flop counter functionality requiring triton module private const string TritonWindowsPackage = "triton-windows"; diff --git a/StabilityMatrix.Core/Models/Rocm/OneTrainerWindowsRocmProfile.cs b/StabilityMatrix.Core/Models/Rocm/OneTrainerRocmProfile.cs similarity index 87% rename from StabilityMatrix.Core/Models/Rocm/OneTrainerWindowsRocmProfile.cs rename to StabilityMatrix.Core/Models/Rocm/OneTrainerRocmProfile.cs index 10f3c6a86..fc88a264e 100644 --- a/StabilityMatrix.Core/Models/Rocm/OneTrainerWindowsRocmProfile.cs +++ b/StabilityMatrix.Core/Models/Rocm/OneTrainerRocmProfile.cs @@ -1,15 +1,15 @@ -using StabilityMatrix.Core.Models.Packages; +using StabilityMatrix.Core.Models.Packages; using StabilityMatrix.Core.Python; namespace StabilityMatrix.Core.Models.Rocm; /// -/// Shared Windows ROCm profile for OneTrainer. -/// Python 3.12 only - required by compatible bitsandbytes wheel. +/// Shared ROCm profile for OneTrainer. +/// Python 3.12 only - required by compatible bitsandbytes wheel on Windows. /// -public class OneTrainerWindowsRocmProfile : RocmPackageProfile +public class OneTrainerRocmProfile : RocmPackageProfile { - public static RocmPackageProfile Default { get; } = new OneTrainerWindowsRocmProfile(); + public static RocmPackageProfile Default { get; } = new OneTrainerRocmProfile(); // Restores flop counter functionality requiring triton module private const string TritonWindowsPackage = "triton-windows"; diff --git a/StabilityMatrix.Core/Models/Rocm/ReforgeWindowsRocmProfile.cs b/StabilityMatrix.Core/Models/Rocm/ReforgeRocmProfile.cs similarity index 82% rename from StabilityMatrix.Core/Models/Rocm/ReforgeWindowsRocmProfile.cs rename to StabilityMatrix.Core/Models/Rocm/ReforgeRocmProfile.cs index 1ff866ca1..f2674cf48 100644 --- a/StabilityMatrix.Core/Models/Rocm/ReforgeWindowsRocmProfile.cs +++ b/StabilityMatrix.Core/Models/Rocm/ReforgeRocmProfile.cs @@ -1,20 +1,20 @@ -using StabilityMatrix.Core.Helper; +using StabilityMatrix.Core.Helper; using StabilityMatrix.Core.Models.Packages; using StabilityMatrix.Core.Services.Rocm; namespace StabilityMatrix.Core.Models.Rocm; /// -/// Shared Windows ROCm profile for reForge. +/// Shared ROCm profile for reForge. /// -public class ReforgeWindowsRocmProfile : RocmPackageProfile +public class ReforgeRocmProfile : RocmPackageProfile { // reForge currently pins accelerate 0.21.0, but 0.22.0 avoids the early distributed.torch import that breaks on Windows ROCm torch builds // caused by the marigold depth controlnet preprocessor - public const string WindowsRocmAccelerateVersion = "0.22.0"; + public const string AccelerateVersion = "0.22.0"; private static readonly string[] DisabledCudaLaunchOptionNames = ["CUDA Malloc", "CUDA Stream"]; - public static ReforgeWindowsRocmProfile Default { get; } = new ReforgeWindowsRocmProfile(); + public static ReforgeRocmProfile Default { get; } = new ReforgeRocmProfile(); public static RocmPackageProfile CreateProfile(IEnumerable requirementsFilePaths) { @@ -32,7 +32,7 @@ public static RocmPackageProfile CreateProfile(IEnumerable requirementsF [ "numpy==1.26.4", "setuptools==69.5.1", - $"accelerate=={WindowsRocmAccelerateVersion}", + $"accelerate=={AccelerateVersion}", ], PostTorchInstallPipArgs = [ @@ -46,7 +46,7 @@ public static RocmPackageProfile CreateProfile(IEnumerable requirementsF }; } - public override void ApplyWindowsRocmLaunchDefaults( + public override void ApplyRocmLaunchDefaults( List launchOptions, IRocmPackageHelper rocmPackageHelper ) @@ -76,7 +76,7 @@ IRocmPackageHelper rocmPackageHelper return null; } - return WindowsRocmSupport.PreferLegacyAttentionFallback(compatibility.ResolvedGfxArch) + return RocmSupport.PreferLegacyAttentionFallback(compatibility.ResolvedGfxArch) ? "--attention-quad" : "--attention-pytorch"; } diff --git a/StabilityMatrix.Core/Models/Rocm/RocmEnvironmentOptions.cs b/StabilityMatrix.Core/Models/Rocm/RocmEnvironmentOptions.cs index 88216ec12..b97eba32a 100644 --- a/StabilityMatrix.Core/Models/Rocm/RocmEnvironmentOptions.cs +++ b/StabilityMatrix.Core/Models/Rocm/RocmEnvironmentOptions.cs @@ -1,4 +1,4 @@ -namespace StabilityMatrix.Core.Models.Rocm; +namespace StabilityMatrix.Core.Models.Rocm; /// /// Controls how ROCm helper defaults, package-specific variables, and user overrides are layered at launch. @@ -15,6 +15,13 @@ public class RocmEnvironmentOptions /// public bool IncludeUserOverrides { get; init; } = true; + /// + /// When true, inherited system ROCm/HIP SDK variables are neutralized so a helper-managed, + /// self-contained ROCm SDK is used. Disable for installs that intentionally rely on a system + /// ROCm or a PyTorch-bundled runtime instead. + /// + public bool SanitizeRocmSdkEnvironment { get; init; } = true; + /// /// When set, overrides the default PyTorch allocator tuning string added by the ROCm helper. /// @@ -41,7 +48,7 @@ public class RocmEnvironmentOptions public string? FlashAttentionTritonAmdEnable { get; init; } = "TRUE"; /// - /// When true, helper-managed defaults will enable ROCm AOTriton on modern Windows ROCm architectures. + /// When true, helper-managed defaults will enable ROCm AOTriton on modern ROCm architectures. /// public bool ApplyAotritonExperimental { get; init; } = true; diff --git a/StabilityMatrix.Core/Models/Rocm/RocmInstallTarget.cs b/StabilityMatrix.Core/Models/Rocm/RocmInstallTarget.cs new file mode 100644 index 000000000..a9c263fad --- /dev/null +++ b/StabilityMatrix.Core/Models/Rocm/RocmInstallTarget.cs @@ -0,0 +1,39 @@ +namespace StabilityMatrix.Core.Models.Rocm; + +/// +/// Package index a helper-managed ROCm install resolves torch from. AMD's feeds carry per-architecture +/// device packages and the ROCm SDK; the upstream PyTorch feeds carry neither. +/// +public enum RocmIndexKind +{ + /// AMD's stable multi-arch feed, the default on both platforms. + AmdStable, + + /// AMD's nightly multi-arch feed, used for targets the stable feed does not publish. + AmdNightly, + + /// Upstream PyTorch ROCm 7.2 channel, available to a package as an override. + PyTorchRocm72, + + /// Upstream PyTorch ROCm 7.14 channel, available to a package as an override. + PyTorchRocm714, +} + +/// +/// Everything the helper needs to install ROCm torch: which index to use, which device +/// package applies, and which index features are available. Resolved from the architecture and the +/// package profile. +/// +public sealed record RocmInstallTarget +{ + public required RocmIndexKind IndexKind { get; init; } + + public required string IndexUrl { get; init; } + + /// Per-architecture device package extra, e.g. device-gfx1201. Null when the index has none. + public string? DeviceExtra { get; init; } + + public required bool SupportsDeviceExtras { get; init; } + + public required bool SupportsSdkDevel { get; init; } +} diff --git a/StabilityMatrix.Core/Models/Rocm/RocmMachineState.cs b/StabilityMatrix.Core/Models/Rocm/RocmMachineState.cs index bf6c07dc5..b1c801b6d 100644 --- a/StabilityMatrix.Core/Models/Rocm/RocmMachineState.cs +++ b/StabilityMatrix.Core/Models/Rocm/RocmMachineState.cs @@ -1,9 +1,9 @@ -using StabilityMatrix.Core.Helper.HardwareInfo; +using StabilityMatrix.Core.Helper.HardwareInfo; namespace StabilityMatrix.Core.Models.Rocm; /// -/// Captures the canonical Windows ROCm machine state resolved by the helper. +/// Captures the canonical ROCm machine state resolved by the helper. /// public class RocmMachineState { @@ -15,5 +15,9 @@ public class RocmMachineState public string? RuntimeGfxArch { get; init; } - public string? MultiArchDeviceExtra { get; init; } + /// Number of AMD GPUs the machine exposes, GPU agents only. + public int AmdDeviceCount { get; init; } + + /// ROCm device index of the selected GPU, when known. + public int? DeviceIndex { get; init; } } diff --git a/StabilityMatrix.Core/Models/Rocm/RocmPackageProfile.cs b/StabilityMatrix.Core/Models/Rocm/RocmPackageProfile.cs index 41e13f8f4..37b1af6a3 100644 --- a/StabilityMatrix.Core/Models/Rocm/RocmPackageProfile.cs +++ b/StabilityMatrix.Core/Models/Rocm/RocmPackageProfile.cs @@ -1,4 +1,4 @@ -using StabilityMatrix.Core.Models.Packages; +using StabilityMatrix.Core.Models.Packages; using StabilityMatrix.Core.Models.Progress; using StabilityMatrix.Core.Services.Rocm; @@ -11,10 +11,16 @@ namespace StabilityMatrix.Core.Models.Rocm; public class RocmPackageProfile { /// - /// Standard package install configuration used before the helper-managed Windows ROCm torch step. + /// Standard package install configuration used before the helper-managed ROCm torch step. /// public PipInstallConfig InstallConfig { get; init; } = new(); + /// + /// Index to install torch from, overriding the platform default. Leave null unless the package needs a + /// specific ROCm channel. + /// + public RocmIndexKind? Index { get; init; } + /// /// Optional callback for package-specific environment variables derived from a resolved ROCm context. /// @@ -29,17 +35,17 @@ public Func< public RocmEnvironmentOptions EnvironmentOptions { get; init; } = new(); /// - /// Allows a package profile to adjust default launch options for Windows ROCm. + /// Allows a package profile to adjust default launch options for ROCm launches. /// Default implementation is a no-op; profiles that need package-specific adjustments /// should override this method. /// - public virtual void ApplyWindowsRocmLaunchDefaults( + public virtual void ApplyRocmLaunchDefaults( List launchOptions, IRocmPackageHelper rocmPackageHelper ) { } /// - /// Returns a package-specific preferred cross-attention argument for Windows ROCm launches. + /// Returns a package-specific preferred cross-attention argument for ROCm launches. /// Default implementation returns null indicating no preference. /// public virtual string? GetPreferredCrossAttentionArgument(IRocmPackageHelper rocmPackageHelper) diff --git a/StabilityMatrix.Core/Models/Rocm/RocmRuntimeContext.cs b/StabilityMatrix.Core/Models/Rocm/RocmRuntimeContext.cs index 1fdda7914..0a6d4e61e 100644 --- a/StabilityMatrix.Core/Models/Rocm/RocmRuntimeContext.cs +++ b/StabilityMatrix.Core/Models/Rocm/RocmRuntimeContext.cs @@ -1,4 +1,4 @@ -using StabilityMatrix.Core.Helper.HardwareInfo; +using StabilityMatrix.Core.Helper.HardwareInfo; namespace StabilityMatrix.Core.Models.Rocm; @@ -15,4 +15,10 @@ public class RocmRuntimeContext public GpuInfo? SelectedGpu { get; init; } public string? RuntimeGfxArch { get; init; } + + /// Number of AMD GPUs the machine exposes, GPU agents only. + public int AmdDeviceCount { get; init; } + + /// ROCm device index of the selected GPU, when known. + public int? DeviceIndex { get; init; } } diff --git a/StabilityMatrix.Core/Models/Rocm/RocmSupport.cs b/StabilityMatrix.Core/Models/Rocm/RocmSupport.cs new file mode 100644 index 000000000..be382da11 --- /dev/null +++ b/StabilityMatrix.Core/Models/Rocm/RocmSupport.cs @@ -0,0 +1,172 @@ +using StabilityMatrix.Core.Helper.HardwareInfo; + +namespace StabilityMatrix.Core.Models.Rocm; + +/// +/// Centralizes ROCm support and architecture policy so hardware detection, package selection, +/// installation, and shared launch decisions use the same support map. +/// +public static class RocmSupport +{ + public const string StableMultiArchPythonPackageIndexUrl = "https://stable.repo.amd.com/rocm/whl-next/"; + public const string NightlyMultiArchPythonPackageIndexUrl = "https://nightly.repo.amd.com/rocm/whl-next/"; + + // Used to exclude modern gfxarches from AOTriton activation EnVar as AOTriton does not currently support them. + // This is a temporary measure until AOTriton adds support for these architectures. + private static readonly HashSet AotritonExperimentalExcludedArchitectures = + [ + "gfx1152", + "gfx1153", + ]; + + /// + /// System-level ROCm/HIP SDK variables that conflict with the self-contained pip ROCm SDK the + /// helper manages. They are removed from helper-managed child environments so a stale system + /// ROCm or HIP SDK install cannot shadow the venv toolchain or its compiled extensions. + /// + public static readonly string[] ConflictingRocmEnvironmentVariables = + [ + "ROCM_PATH", + "ROCM_HOME", + "HIP_PATH", + "ROCM_DIR", + "HIP_DIR", + ]; + + public static bool IsSupportedGpu(GpuInfo? gpu) + { + if (gpu is null || !gpu.IsAmd || string.IsNullOrWhiteSpace(gpu.Name)) + return false; + + return IsSupportedArchitecture(gpu.GetAmdGfxArch()); + } + + public static bool IsSupportedArchitecture(string? gfxArch) + { + return TryGetCanonicalArchitecture(gfxArch) is not null; + } + + public static bool IsModernArchitecture(string? gfxArch) + { + return gfxArch?.StartsWith("gfx110", StringComparison.OrdinalIgnoreCase) == true + || gfxArch?.StartsWith("gfx115", StringComparison.OrdinalIgnoreCase) == true + || gfxArch?.StartsWith("gfx120", StringComparison.OrdinalIgnoreCase) == true; + } + + public static bool SupportsAotritonExperimental(string? gfxArch) + { + var canonicalArch = TryGetCanonicalArchitecture(gfxArch); + return canonicalArch is not null + && IsModernArchitecture(canonicalArch) + && !AotritonExperimentalExcludedArchitectures.Contains(canonicalArch); + } + + public static bool IsLegacyArchitecture(string? gfxArch) + { + return IsSupportedArchitecture(gfxArch) && !IsModernArchitecture(gfxArch); + } + + public static bool PreferLegacyAttentionFallback(string? gfxArch) + { + return IsLegacyArchitecture(gfxArch); + } + + /// + /// True for targets that ship fused SDPA kernels, which on ROCm means the target has an AOTriton + /// image: RDNA3, RDNA3.5, RDNA4 and CDNA2 or newer. gfx908 and the GCN/RDNA1/RDNA2 targets have + /// no fused kernels and fall back to the math backend. + /// + public static bool SupportsFusedSdpa(string? gfxArch) + { + var canonicalArch = TryGetCanonicalArchitecture(gfxArch); + return canonicalArch is not null + && (IsModernArchitecture(canonicalArch) || canonicalArch is "gfx90a" or "gfx942" or "gfx950"); + } + + /// + /// True when the target has no fused SDPA kernels and the math backend has to be forced. + /// + public static bool PreferMathSdpaFallback(string? gfxArch) => + IsSupportedArchitecture(gfxArch) && !SupportsFusedSdpa(gfxArch); + + public static bool UsesNightlyMultiArchPythonPackageIndex(string? gfxArch) + { + var canonicalArch = TryGetCanonicalArchitecture(gfxArch); + // TheRock currently publishes Vega multi-arch device packages on the nightly feed but not the stable one. + return canonicalArch is "gfx900" or "gfx906"; + } + + public static string? TryGetCanonicalArchitecture(string? gfxArch) + { + if (string.IsNullOrWhiteSpace(gfxArch)) + return null; + + var normalizedArch = gfxArch.ToLowerInvariant(); + + return normalizedArch switch + { + "gfx900" + or "gfx906" + or "gfx908" + or "gfx90a" + or "gfx942" + or "gfx950" + or "gfx1150" + or "gfx1151" + or "gfx1152" + or "gfx1153" => normalizedArch, + var s + when s.StartsWith("gfx101", StringComparison.Ordinal) + || s.StartsWith("gfx103", StringComparison.Ordinal) + || s.StartsWith("gfx110", StringComparison.Ordinal) + || s.StartsWith("gfx120", StringComparison.Ordinal) => normalizedArch, + _ => null, + }; + } + + public static string? TryGetMultiArchDeviceExtra(string? gfxArch) + { + var canonicalArch = TryGetCanonicalArchitecture(gfxArch); + return canonicalArch is null ? null : $"device-{canonicalArch}"; + } + + /// URL of the index a target resolves torch from. + public static string GetIndexUrl(RocmIndexKind kind) => + kind switch + { + RocmIndexKind.AmdStable => StableMultiArchPythonPackageIndexUrl, + RocmIndexKind.AmdNightly => NightlyMultiArchPythonPackageIndexUrl, + RocmIndexKind.PyTorchRocm72 => "https://download.pytorch.org/whl/rocm7.2", + RocmIndexKind.PyTorchRocm714 => "https://download.pytorch.org/whl/rocm7.14", + _ => throw new ArgumentOutOfRangeException(nameof(kind), kind, "Unknown ROCm index kind."), + }; + + /// + /// Resolves everything an install needs: the index to use, whether per-architecture + /// device packages apply, and whether torch has to be paired with torchvision by the helper. + /// + public static RocmInstallTarget ResolveTarget(string? gfxArch, RocmIndexKind? indexOverride = null) + { + // Vega is only published on AMD's nightly feed. Every other target uses stable unless the package + // profile pins a specific index. + var indexKind = + indexOverride + ?? ( + UsesNightlyMultiArchPythonPackageIndex(gfxArch) + ? RocmIndexKind.AmdNightly + : RocmIndexKind.AmdStable + ); + var usesAmdIndex = indexKind is RocmIndexKind.AmdStable or RocmIndexKind.AmdNightly; + + return new RocmInstallTarget + { + IndexKind = indexKind, + IndexUrl = GetIndexUrl(indexKind), + // AMD's feeds publish per-architecture device packages for both win_amd64 and linux_x86_64. + // The upstream PyTorch includes GPU libraries for all supported device architectures. + DeviceExtra = usesAmdIndex ? TryGetMultiArchDeviceExtra(gfxArch) : null, + SupportsDeviceExtras = usesAmdIndex, + SupportsSdkDevel = usesAmdIndex, + }; + } +} diff --git a/StabilityMatrix.Core/Models/Rocm/Wan2GpWindowsRocmProfile.cs b/StabilityMatrix.Core/Models/Rocm/Wan2GpRocmProfile.cs similarity index 65% rename from StabilityMatrix.Core/Models/Rocm/Wan2GpWindowsRocmProfile.cs rename to StabilityMatrix.Core/Models/Rocm/Wan2GpRocmProfile.cs index 91fa8e1db..1657616e6 100644 --- a/StabilityMatrix.Core/Models/Rocm/Wan2GpWindowsRocmProfile.cs +++ b/StabilityMatrix.Core/Models/Rocm/Wan2GpRocmProfile.cs @@ -1,13 +1,13 @@ -using StabilityMatrix.Core.Models.Packages; +using StabilityMatrix.Core.Models.Packages; namespace StabilityMatrix.Core.Models.Rocm; /// -/// Shared Windows ROCm profile for Wan2GP. +/// Shared ROCm profile for Wan2GP. /// -public class Wan2GpWindowsRocmProfile : RocmPackageProfile +public class Wan2GpRocmProfile : RocmPackageProfile { - public Wan2GpWindowsRocmProfile() + public Wan2GpRocmProfile() { InstallConfig = new PipInstallConfig { @@ -17,5 +17,5 @@ public Wan2GpWindowsRocmProfile() }; } - public static RocmPackageProfile Default { get; } = new Wan2GpWindowsRocmProfile(); + public static RocmPackageProfile Default { get; } = new Wan2GpRocmProfile(); } diff --git a/StabilityMatrix.Core/Models/Rocm/WindowsRocmSupport.cs b/StabilityMatrix.Core/Models/Rocm/WindowsRocmSupport.cs deleted file mode 100644 index 01d8b0a1c..000000000 --- a/StabilityMatrix.Core/Models/Rocm/WindowsRocmSupport.cs +++ /dev/null @@ -1,98 +0,0 @@ -using StabilityMatrix.Core.Helper.HardwareInfo; - -namespace StabilityMatrix.Core.Models.Rocm; - -/// -/// Centralizes Windows ROCm support and architecture policy so hardware detection, package selection, -/// installation, and shared launch decisions use the same support map. -/// -public static class WindowsRocmSupport -{ - public const string StableMultiArchPythonPackageIndexUrl = "https://stable.repo.amd.com/rocm/whl-next/"; - public const string NightlyMultiArchPythonPackageIndexUrl = "https://nightly.repo.amd.com/rocm/whl-next/"; - - // Used to exclude modern gfxarches from AOTriton activation EnVar as AOTriton does not currently support them. - // This is a temporary measure until AOTriton adds support for these architectures. - private static readonly HashSet AotritonExperimentalExcludedArchitectures = - [ - "gfx1152", - "gfx1153", - ]; - - public static bool IsSupportedGpu(GpuInfo? gpu) - { - if (gpu is null || !gpu.IsAmd || string.IsNullOrWhiteSpace(gpu.Name)) - return false; - - return IsSupportedArchitecture(gpu.GetAmdGfxArch()); - } - - public static bool IsSupportedArchitecture(string? gfxArch) - { - return TryGetCanonicalArchitecture(gfxArch) is not null; - } - - public static bool IsModernArchitecture(string? gfxArch) - { - return gfxArch?.StartsWith("gfx110", StringComparison.OrdinalIgnoreCase) == true - || gfxArch?.StartsWith("gfx115", StringComparison.OrdinalIgnoreCase) == true - || gfxArch?.StartsWith("gfx120", StringComparison.OrdinalIgnoreCase) == true; - } - - public static bool SupportsAotritonExperimental(string? gfxArch) - { - var canonicalArch = TryGetCanonicalArchitecture(gfxArch); - return canonicalArch is not null - && IsModernArchitecture(canonicalArch) - && !AotritonExperimentalExcludedArchitectures.Contains(canonicalArch); - } - - public static bool IsLegacyArchitecture(string? gfxArch) - { - return IsSupportedArchitecture(gfxArch) && !IsModernArchitecture(gfxArch); - } - - public static bool PreferLegacyAttentionFallback(string? gfxArch) - { - return IsLegacyArchitecture(gfxArch); - } - - public static bool UsesNightlyMultiArchPythonPackageIndex(string? gfxArch) - { - var canonicalArch = TryGetCanonicalArchitecture(gfxArch); - // TheRock currently publishes Vega multi-arch device packages on the nightly feed but not the stable one. - return canonicalArch is "gfx900" or "gfx906"; - } - - public static string? TryGetCanonicalArchitecture(string? gfxArch) - { - if (string.IsNullOrWhiteSpace(gfxArch)) - return null; - - var normalizedArch = gfxArch.ToLowerInvariant(); - - return normalizedArch switch - { - "gfx900" or "gfx906" or "gfx1150" or "gfx1151" or "gfx1152" or "gfx1153" => normalizedArch, - var s - when s.StartsWith("gfx101", StringComparison.Ordinal) - || s.StartsWith("gfx103", StringComparison.Ordinal) - || s.StartsWith("gfx110", StringComparison.Ordinal) - || s.StartsWith("gfx120", StringComparison.Ordinal) => normalizedArch, - _ => null, - }; - } - - public static string? TryGetMultiArchDeviceExtra(string? gfxArch) - { - var canonicalArch = TryGetCanonicalArchitecture(gfxArch); - return canonicalArch is null ? null : $"device-{canonicalArch}"; - } - - public static string GetMultiArchPythonPackageIndexUrl(string? gfxArch) - { - return UsesNightlyMultiArchPythonPackageIndex(gfxArch) - ? NightlyMultiArchPythonPackageIndexUrl - : StableMultiArchPythonPackageIndexUrl; - } -} diff --git a/StabilityMatrix.Core/Processes/ProcessRunner.cs b/StabilityMatrix.Core/Processes/ProcessRunner.cs index 665e91fc6..f606bafac 100644 --- a/StabilityMatrix.Core/Processes/ProcessRunner.cs +++ b/StabilityMatrix.Core/Processes/ProcessRunner.cs @@ -10,6 +10,13 @@ public static class ProcessRunner { private static readonly Logger Logger = LogManager.GetCurrentClassLogger(); + /// + /// Sentinel value for an environment variable entry. When used as a value, the named variable is + /// removed from the child process environment instead of being set, which is how inherited system + /// settings are neutralized without passing an empty string that would still shadow the default. + /// + public const string UnsetEnvironmentVariable = "\0SM_UNSET\0"; + // Managed Python/uv subprocesses should not inherit ambient Python/venv activation state. private static readonly string[] PythonEnvironmentVariablesToSanitize = [ @@ -115,7 +122,15 @@ private static void ApplyEnvironmentVariables( var keys = new List(); foreach (var (key, value) in environmentVariables) { - info.EnvironmentVariables[key] = value; + if (value == UnsetEnvironmentVariable) + { + info.Environment.Remove(key); + } + else + { + info.EnvironmentVariables[key] = value; + } + keys.Add(key); } @@ -651,9 +666,10 @@ public static Task RunBashCommand( /// public static string Quote(string argument) { - var inner = argument.Length >= 2 && argument.StartsWith('"') && argument.EndsWith('"') - ? argument[1..^1] - : argument; + var inner = + argument.Length >= 2 && argument.StartsWith('"') && argument.EndsWith('"') + ? argument[1..^1] + : argument; if (!inner.Contains(' ') && !inner.Contains('"')) return argument; diff --git a/StabilityMatrix.Core/Services/Rocm/IRocmPackageHelper.cs b/StabilityMatrix.Core/Services/Rocm/IRocmPackageHelper.cs index 84b60b2ee..086640cd3 100644 --- a/StabilityMatrix.Core/Services/Rocm/IRocmPackageHelper.cs +++ b/StabilityMatrix.Core/Services/Rocm/IRocmPackageHelper.cs @@ -1,4 +1,4 @@ -using StabilityMatrix.Core.Models; +using StabilityMatrix.Core.Models; using StabilityMatrix.Core.Models.Packages; using StabilityMatrix.Core.Models.Progress; using StabilityMatrix.Core.Models.Rocm; @@ -23,22 +23,22 @@ public interface IRocmPackageHelper IReadOnlyDictionary BuildLaunchEnvironment(RocmPackageProfile profile); /// - /// Returns shared Windows ROCm launch notice lines if the current machine and selected torch index - /// qualify for the Windows native ROCm launch environment; otherwise returns an empty list. + /// Returns shared ROCm launch notice lines when the selected torch index qualifies for the + /// helper-managed ROCm launch environment on this machine; otherwise returns an empty list. /// - IReadOnlyList GetWindowsLaunchNoticeLines(TorchIndex selectedTorchIndex); + IReadOnlyList GetLaunchNoticeLines(TorchIndex selectedTorchIndex); /// - /// Returns true when the current machine is Windows, the selected torch index is ROCm, - /// and the machine is compatible with Windows native ROCm. + /// Returns true when the selected torch index is ROCm and the helper can manage a native ROCm + /// install on the current machine. /// - bool ShouldApplyWindowsLaunchEnvironment(TorchIndex selectedTorchIndex); + bool ShouldApplyLaunchEnvironment(TorchIndex selectedTorchIndex); /// - /// Ensures a usable Windows ROCm SDK devel package is installed from the ROCm multi-arch index, + /// Ensures a usable ROCm SDK devel package is installed from the ROCm index, /// preferring the same build date token as the installed torch build and falling back to the latest available build. /// - Task EnsureWindowsSdkDevelAsync( + Task EnsureRocmSdkDevelAsync( IPyVenvRunner venvRunner, IProgress? progress = null, Action? onConsoleOutput = null, @@ -46,14 +46,16 @@ Task EnsureWindowsSdkDevelAsync( ); /// - /// Builds the standard pip install config for a Windows ROCm package, ensuring torch installation stays helper-managed. + /// Builds the standard pip install config for a ROCm package, ensuring torch installation stays + /// helper-managed. /// - PipInstallConfig BuildWindowsNativeInstallConfig(RocmPackageProfile profile); + PipInstallConfig BuildRocmInstallConfig(RocmPackageProfile profile); /// - /// Installs the Windows-native ROCm torch wheel set for a package using helper-resolved multi-arch device extras. + /// Installs the ROCm torch wheel set for a package from the resolved index, selecting device code + /// through the index's device extra when it has one. /// - Task InstallWindowsNativeTorchAsync( + Task InstallRocmTorchAsync( IPyVenvRunner venvRunner, InstalledPackage installedPackage, RocmPackageProfile profile, diff --git a/StabilityMatrix.Core/Services/Rocm/RocmArchDetector.cs b/StabilityMatrix.Core/Services/Rocm/RocmArchDetector.cs new file mode 100644 index 000000000..54cb5ff74 --- /dev/null +++ b/StabilityMatrix.Core/Services/Rocm/RocmArchDetector.cs @@ -0,0 +1,164 @@ +using NLog; + +namespace StabilityMatrix.Core.Services.Rocm; + +/// +/// One AMD GPU exposed through the KFD topology. +/// +internal sealed record RocmDetectedDevice +{ + public required string GfxArch { get; init; } + + /// + /// ROCm device index: the position among GPU nodes in KFD node order, which is the order the runtime + /// enumerates devices in. + /// + public required int DeviceIndex { get; init; } + + /// SIMD count reported by the driver. Used to tell a discrete card from an APU. + public required int SimdCount { get; init; } +} + +/// +/// Reads the GPUs the amdgpu kernel module exports through the KFD topology. On Linux this is the +/// driver's own answer, so no ROCm userspace and no marketing-name table is involved. Windows has no KFD +/// node and keeps resolving architectures from the adapter name table. +/// +internal static class RocmArchDetector +{ + private const string KfdNodesPath = "/sys/class/kfd/kfd/topology/nodes"; + + private static readonly Logger Logger = LogManager.GetCurrentClassLogger(); + private static readonly object Gate = new(); + private static IReadOnlyList? cached; + + /// GPUs the driver reports, in node order. Empty when KFD is unavailable. + internal static IReadOnlyList Devices + { + get + { + if (cached is { } existing) + { + return existing; + } + + lock (Gate) + { + return cached ??= Detect(); + } + } + } + + /// Forces the next access to re-read the topology. + internal static void ResetCache() + { + lock (Gate) + { + cached = null; + } + } + + private static IReadOnlyList Detect() + { + try + { + if (!OperatingSystem.IsLinux() || !Directory.Exists(KfdNodesPath)) + { + return []; + } + + List devices = []; + + // Node ids ascend with PCI topology, so sort numerically rather than by directory name. The + // device index counts GPU nodes only, because that is what the runtime exposes. + var nodePaths = Directory + .EnumerateDirectories(KfdNodesPath) + .OrderBy(path => int.TryParse(Path.GetFileName(path), out var node) ? node : int.MaxValue); + + foreach (var nodePath in nodePaths) + { + var propertiesPath = Path.Combine(nodePath, "properties"); + if (!File.Exists(propertiesPath)) + { + continue; + } + + var device = TryParseNode(File.ReadAllText(propertiesPath), devices.Count); + if (device is not null) + { + devices.Add(device); + } + } + + return devices; + } + catch (Exception exception) + { + Logger.Warn(exception, "Failed to read the KFD topology; falling back to adapter names."); + return []; + } + } + + /// + /// Decodes one node's properties text, or returns null for CPU nodes and unparsable nodes. + /// + internal static RocmDetectedDevice? TryParseNode(string properties, int deviceIndex) + { + Dictionary values = new(StringComparer.Ordinal); + + foreach (var line in properties.Split('\n')) + { + var parts = line.Split( + ' ', + StringSplitOptions.RemoveEmptyEntries | StringSplitOptions.TrimEntries + ); + if (parts.Length == 2 && int.TryParse(parts[1], out var value)) + { + values[parts[0]] = value; + } + } + + // CPU nodes carry simd_count 0 and no target version. + if (!values.TryGetValue("simd_count", out var simdCount) || simdCount == 0) + { + return null; + } + + if (!values.TryGetValue("gfx_target_version", out var targetVersion) || targetVersion == 0) + { + return null; + } + + return new RocmDetectedDevice + { + GfxArch = FormatGfxTarget(targetVersion), + DeviceIndex = deviceIndex, + SimdCount = simdCount, + }; + } + + /// + /// Picks the device an install should target when the user has no preference: the one with the most + /// compute. An APU reports a handful of SIMD units, a discrete card reports tens to hundreds, and KFD + /// node order is not a reliable indicator of which is which. + /// + internal static RocmDetectedDevice? SelectPreferred(IReadOnlyList devices) => + devices + .OrderByDescending(device => device.SimdCount) + .ThenBy(device => device.DeviceIndex) + .FirstOrDefault(); + + /// + /// Decodes the KFD gfx_target_version integer into its ISA name: 110000 to gfx1100, + /// 120001 to gfx1201, 100300 to gfx1030, 110500 to gfx1150, 90010 to gfx90a. The name is always + /// major, then minor, then the stepping as a single hex digit, including when the stepping is zero. + /// + internal static string FormatGfxTarget(int gfxTargetVersion) + { + var major = gfxTargetVersion / 10000; + var minor = gfxTargetVersion / 100 % 100; + var stepping = gfxTargetVersion % 100; + + return $"gfx{major}{minor}{stepping:x}"; + } +} diff --git a/StabilityMatrix.Core/Services/Rocm/RocmPackageHelper.cs b/StabilityMatrix.Core/Services/Rocm/RocmPackageHelper.cs index 94e9a9b26..bde4f1d2b 100644 --- a/StabilityMatrix.Core/Services/Rocm/RocmPackageHelper.cs +++ b/StabilityMatrix.Core/Services/Rocm/RocmPackageHelper.cs @@ -1,4 +1,4 @@ -using System.Text.Json; +using System.Text.Json; using Injectio.Attributes; using NLog; using StabilityMatrix.Core.Exceptions; @@ -25,9 +25,21 @@ public class RocmPackageHelper : IRocmPackageHelper private static readonly StringComparer EnvComparer = StringComparer.OrdinalIgnoreCase; private const string RocmSdkDevelPackageName = "rocm-sdk-devel"; - private static readonly string[] WindowsLaunchNoticeLines = + + // The ROCm torch install only ever needs to move the torch trio. Bundling or upgrading anything + // wider lets the resolver re-evaluate transitive dependencies against AMD's feed. + private static readonly string[] RocmTorchPackageNames = ["torch", "torchvision", "torchaudio"]; + private static readonly string[] RocmTorchUpgradePackageArgs = BuildRepeatablePackageArgs( + "--upgrade-package", + RocmTorchPackageNames + ); + private static readonly string[] RocmTorchReinstallPackageArgs = BuildRepeatablePackageArgs( + "--reinstall-package", + RocmTorchPackageNames + ); + private static readonly string[] LaunchNoticeLines = [ - "Stability Matrix Windows ROCm Notice: Windows AMD ROCm support is experimental. Please report any issues to Stability Matrix first so it can be determined whether the issue is package-specific.", + "Stability Matrix ROCm Notice: AMD ROCm support is experimental. Please report any issues to Stability Matrix first so it can be determined whether the issue is package-specific.", "Because this setup may not be officially supported by package developers, only contact upstream support for issues clearly caused by the package itself.", ]; @@ -40,11 +52,11 @@ public class RocmPackageHelper : IRocmPackageHelper public RocmPackageHelper(ISettingsManager settingsManager) { this.settingsManager = settingsManager; - machineState = new Lazy(ResolveWindowsMachineState); + machineState = new Lazy(ResolveMachineState); } /// - /// Evaluates the current Windows machine state and returns the resolved ROCm compatibility result. + /// Evaluates the current machine state and returns the resolved ROCm compatibility result. /// public RocmCompatibilityResult GetCompatibility() { @@ -52,7 +64,7 @@ public RocmCompatibilityResult GetCompatibility() } /// - /// Resolves launch-time ROCm runtime details from the current Windows machine state. + /// Resolves launch-time ROCm runtime details from the current machine state. /// This is used to build helper-managed environment variables for package launch. /// private RocmRuntimeContext ResolveRuntimeContext() @@ -74,6 +86,8 @@ private RocmRuntimeContext ResolveRuntimeContext() IsSupported = true, SelectedGpu = state.SelectedGpu, RuntimeGfxArch = state.RuntimeGfxArch, + AmdDeviceCount = state.AmdDeviceCount, + DeviceIndex = state.DeviceIndex, }; } @@ -101,30 +115,31 @@ public IReadOnlyDictionary BuildLaunchEnvironment(RocmPackagePro return mergedEnvironment; } - private IReadOnlyList GetWindowsLaunchNoticeLines() + private IReadOnlyList GetLaunchNoticeLines() { - return WindowsLaunchNoticeLines; + return LaunchNoticeLines; } /// - public bool ShouldApplyWindowsLaunchEnvironment(TorchIndex selectedTorchIndex) + public bool ShouldApplyLaunchEnvironment(TorchIndex selectedTorchIndex) { - if (!Compat.IsWindows || selectedTorchIndex != TorchIndex.Rocm) + if (selectedTorchIndex != TorchIndex.Rocm) return false; return GetCompatibility().IsCompatible; } /// - public IReadOnlyList GetWindowsLaunchNoticeLines(TorchIndex selectedTorchIndex) + public IReadOnlyList GetLaunchNoticeLines(TorchIndex selectedTorchIndex) { - return ShouldApplyWindowsLaunchEnvironment(selectedTorchIndex) ? GetWindowsLaunchNoticeLines() : []; + return ShouldApplyLaunchEnvironment(selectedTorchIndex) ? GetLaunchNoticeLines() : []; } /// - /// Ensures rocm-sdk-devel is installed from the ROCm multi-arch index. - /// It prefers a build whose date token matches the installed ROCm torch build and falls back to the latest available build when no exact match is available. + /// Ensures rocm-sdk-devel is installed from the resolved ROCm index, on the same ROCm line as + /// the installed torch build. Nightly builds match on their build date token, stable builds on the + /// ROCm version in torch's local label, and only when neither matches is the newest build used. /// - public async Task EnsureWindowsSdkDevelAsync( + public async Task EnsureRocmSdkDevelAsync( IPyVenvRunner venvRunner, IProgress? progress = null, Action? onConsoleOutput = null, @@ -132,70 +147,98 @@ public async Task EnsureWindowsSdkDevelAsync( ) { var state = machineState.Value; - var multiArchPythonPackageIndexUrl = WindowsRocmSupport.GetMultiArchPythonPackageIndexUrl( - state.RuntimeGfxArch - ); + var target = RocmSupport.ResolveTarget(state.RuntimeGfxArch); + + if (!target.SupportsSdkDevel) + { + Logger.Info( + "The {IndexKind} index does not publish {Package}; skipping the SDK devel step.", + target.IndexKind, + RocmSdkDevelPackageName + ); + return; + } + + var indexUrl = target.IndexUrl; var torchInfo = await venvRunner.PipShow("torch").ConfigureAwait(false); if (torchInfo is null) { throw new InvalidOperationException( - "torch is not installed in this environment. Install the Windows ROCm torch build first." + "torch is not installed in this environment. Install the ROCm torch build first." ); } - if (!IsUsableWindowsNativeTorchBuild(torchInfo.Version, null)) + if (!IsUsableRocmTorchBuild(torchInfo.Version, null)) { throw new InvalidOperationException( - $"Installed torch is not a usable Windows ROCm build (detected version: {torchInfo.Version})." + $"Installed torch is not a usable ROCm build (detected version: {torchInfo.Version})." ); } var nightlyBuildDateToken = TryGetNightlyBuildDateToken(torchInfo.Version); + var rocmVersion = TryGetRocmVersionLabel(torchInfo.Version); var installedRocmSdkDevel = await venvRunner.PipShow(RocmSdkDevelPackageName).ConfigureAwait(false); + + // Already on the right line: a nightly build with the matching date token, or a stable build + // matching the ROCm version the torch wheel was compiled against. if ( - !string.IsNullOrWhiteSpace(nightlyBuildDateToken) - && HasNightlyBuildDateToken(installedRocmSdkDevel?.Version, nightlyBuildDateToken) + ( + !string.IsNullOrWhiteSpace(nightlyBuildDateToken) + && HasNightlyBuildDateToken(installedRocmSdkDevel?.Version, nightlyBuildDateToken) + ) || IsRocmSdkDevelMatch(installedRocmSdkDevel?.Version, rocmVersion) ) { return; } - var indexResult = await venvRunner - .PipIndex(RocmSdkDevelPackageName, multiArchPythonPackageIndexUrl) - .ConfigureAwait(false); + var indexResult = await venvRunner.PipIndex(RocmSdkDevelPackageName, indexUrl).ConfigureAwait(false); - var latestVersion = indexResult?.AvailableVersions.FirstOrDefault(); - var matchingVersion = string.IsNullOrWhiteSpace(nightlyBuildDateToken) - ? null - : indexResult?.AvailableVersions.FirstOrDefault(version => - HasNightlyBuildDateToken(version, nightlyBuildDateToken) - ); - var versionToInstall = matchingVersion ?? latestVersion; + var selection = SelectSdkDevelVersion( + indexResult?.AvailableVersions, + rocmVersion, + nightlyBuildDateToken + ); - if (string.IsNullOrWhiteSpace(versionToInstall)) + if (selection is not { } chosen) { throw new InvalidOperationException( - $"No {RocmSdkDevelPackageName} builds were found on the ROCm multi-arch index." + $"No {RocmSdkDevelPackageName} builds were found on the ROCm index." ); } + var versionToInstall = chosen.Version; + var matchingVersion = chosen.IsExactMatch ? versionToInstall : null; + if (!string.IsNullOrWhiteSpace(matchingVersion)) { progress?.Report( new ProgressReport( -1f, - $"Installing {RocmSdkDevelPackageName} {matchingVersion} for Windows ROCm...", + $"Installing {RocmSdkDevelPackageName} {matchingVersion} for ROCm...", isIndeterminate: true ) ); } else { + if (!string.IsNullOrWhiteSpace(rocmVersion)) + { + Logger.Warn( + "No {Package} build matches the ROCm {RocmVersion} line of the installed torch " + + "({TorchVersion}); installing {Version} instead. Builds against the development " + + "SDK may then not load in this environment.", + RocmSdkDevelPackageName, + rocmVersion, + torchInfo.Version, + versionToInstall + ); + } + progress?.Report( new ProgressReport( -1f, - $"Falling back to latest available {RocmSdkDevelPackageName} build {versionToInstall} for Windows ROCm...", + $"Falling back to latest available {RocmSdkDevelPackageName} build {versionToInstall} for ROCm...", isIndeterminate: true ) ); @@ -205,7 +248,7 @@ await venvRunner .PipInstall( new PipInstallArgs() .AddArg("--upgrade") - .AddKeyedArgs("--index-url", ["--index-url", multiArchPythonPackageIndexUrl]) + .AddKeyedArgs("--index-url", ["--index-url", indexUrl]) .AddArg($"{RocmSdkDevelPackageName}=={versionToInstall}"), onConsoleOutput ) @@ -213,17 +256,18 @@ await venvRunner } /// - /// Builds the standard pip install config for helper-managed Windows ROCm package installs. + /// Builds the standard pip install config for helper-managed ROCm package installs. /// - public PipInstallConfig BuildWindowsNativeInstallConfig(RocmPackageProfile profile) + public PipInstallConfig BuildRocmInstallConfig(RocmPackageProfile profile) { return profile.InstallConfig with { SkipTorchInstall = true }; } /// - /// Installs the ROCm torch wheel set from the multi-arch index and verifies that the resulting torch installation reports usable ROCm metadata. + /// Installs the ROCm torch wheel set from the resolved index and verifies that the resulting torch + /// installation reports usable ROCm metadata. /// - public async Task InstallWindowsNativeTorchAsync( + public async Task InstallRocmTorchAsync( IPyVenvRunner venvRunner, InstalledPackage installedPackage, RocmPackageProfile profile, @@ -236,26 +280,26 @@ public async Task InstallWindowsNativeTorchAsync( if (!state.IsCompatible) { throw new InvalidOperationException( - state.FailureReason ?? "Windows ROCm installation is not supported for the current machine." + state.FailureReason ?? "ROCm installation is not supported for the current machine." ); } - var multiArchDeviceExtra = state.MultiArchDeviceExtra; + var target = RocmSupport.ResolveTarget(state.RuntimeGfxArch, profile.Index); - if (string.IsNullOrWhiteSpace(multiArchDeviceExtra)) + // AMD's feeds select device code through a per-architecture extra; the upstream PyTorch feeds + // either ship one multi-architecture wheel (e.g: 7.2) or explicitly force-installs all device extras (e.g 7.14), + // so an index without device extras is still installable. + if (target.SupportsDeviceExtras && string.IsNullOrWhiteSpace(target.DeviceExtra)) { throw new InvalidOperationException( - $"No Windows ROCm multi-arch device extra is available for '{state.RuntimeGfxArch ?? "unknown"}'." + $"No ROCm device package is available for '{state.RuntimeGfxArch ?? "unknown"}' " + + $"on the {target.IndexKind} index." ); } var installConfig = profile.InstallConfig; - var multiArchPythonPackageIndexUrl = WindowsRocmSupport.GetMultiArchPythonPackageIndexUrl( - state.RuntimeGfxArch - ); - var usesNightlyIndex = WindowsRocmSupport.UsesNightlyMultiArchPythonPackageIndex( - state.RuntimeGfxArch - ); + var indexUrl = target.IndexUrl; + var usesNightlyIndex = target.IndexKind == RocmIndexKind.AmdNightly; // AMD's torchvision wheels carry no torch constraint, and the matched build may be published only // as a pre-release that pip skips. Resolve the newest paired torch/torchvision and pin both. @@ -263,8 +307,8 @@ public async Task InstallWindowsNativeTorchAsync( var torchResolution = await ResolveTorchInstallAsync( venvRunner, - multiArchPythonPackageIndexUrl, - multiArchDeviceExtra, + indexUrl, + target.DeviceExtra, usesNightlyIndex, onConsoleOutput, cancellationToken @@ -275,12 +319,12 @@ public async Task InstallWindowsNativeTorchAsync( var torchArgs = ApplyRocmTorchInstallOptions( BuildRocmTorchArgs( - multiArchPythonPackageIndexUrl, + indexUrl, torchResolution.TorchSpecifier, torchResolution.TorchvisionSpecifier ), installConfig, - installedPackage + installedPackage.PipOverrides ); if (torchResolution.AllowPrerelease) @@ -304,49 +348,76 @@ public async Task InstallWindowsNativeTorchAsync( await venvRunner.PipInstall(postTorchInstallPipArgs, onConsoleOutput).ConfigureAwait(false); } - await VerifyWindowsNativeTorchInstallAsync(venvRunner, onConsoleOutput, cancellationToken) + await VerifyRocmTorchInstallAsync(venvRunner, onConsoleOutput, cancellationToken) .ConfigureAwait(false); + + // A development SDK that is already present has to follow the torch build it sits beside, or the + // headers and libraries come from a different ROCm line than the runtime loading them. + if (await venvRunner.PipShow(RocmSdkDevelPackageName).ConfigureAwait(false) is not null) + { + await EnsureRocmSdkDevelAsync(venvRunner, progress, onConsoleOutput, cancellationToken) + .ConfigureAwait(false); + } } /// /// Applies the profile's upgrade/reinstall flags and any user pip overrides to a ROCm torch install command. + /// Scope is deliberately limited to the torch trio. A blanket --upgrade lets the resolver move + /// transitive dependencies against AMD's feed, numpy foremost, where the only build may be + /// numpy 2.5.2 (cp312+). A blanket --force-reinstall reinstalls the whole graph for the same reason. + /// Packages that must follow the torch build, amd-torch-device-gfx* and triton, are exact-pinned by the + /// wheel metadata, so they are still replaced to satisfy the resolved pin. /// - private static PipInstallArgs ApplyRocmTorchInstallOptions( + internal static PipInstallArgs ApplyRocmTorchInstallOptions( PipInstallArgs args, PipInstallConfig installConfig, - InstalledPackage installedPackage + IEnumerable? pipOverrides ) { if (installConfig.UpgradePackages) { - args = args.AddArg("--upgrade"); + args = args.AddKeyedArgs("--upgrade-package", RocmTorchUpgradePackageArgs); } if (installConfig.ForceReinstallTorch) { - args = args.AddArg("--force-reinstall"); + args = args.AddKeyedArgs("--reinstall-package", RocmTorchReinstallPackageArgs); } - if (installedPackage.PipOverrides != null) + if (pipOverrides != null) { - args = args.WithUserOverrides(installedPackage.PipOverrides); + args = args.WithUserOverrides(pipOverrides.ToList()); } return args; } + /// + /// Expands a repeatable flag so each package gets its own flag value pair. + /// + private static string[] BuildRepeatablePackageArgs(string flag, IEnumerable packages) => + packages.SelectMany(package => new[] { flag, package }).ToArray(); + /// /// Builds the pip args for the resolved torch/torchvision pair. The specs are keyed so a user /// override replaces the pin instead of being appended next to it (which pip cannot satisfy). /// + /// + /// Requirement form for a torch package on the resolved index: a per-architecture extra on AMD's + /// feeds, a plain requirement on the upstream PyTorch feeds, which ship one multi-architecture wheel + /// or all supported device libraries. + /// + internal static string BuildTorchSpecifier(string packageName, string? deviceExtra) => + string.IsNullOrWhiteSpace(deviceExtra) ? packageName : $"{packageName}[{deviceExtra}]"; + internal static PipInstallArgs BuildRocmTorchArgs( - string multiArchPythonPackageIndexUrl, + string indexUrl, string torchSpecifier, string torchvisionSpecifier ) { return new PipInstallArgs() - .AddKeyedArgs("--index-url", ["--index-url", multiArchPythonPackageIndexUrl]) + .AddKeyedArgs("--index-url", ["--index-url", indexUrl]) .AddArg(new Argument(key: "torch", value: torchSpecifier)) .AddArg(new Argument(key: "torchvision", value: torchvisionSpecifier)) .AddArg("torchaudio"); @@ -368,27 +439,27 @@ bool AllowPrerelease /// private async Task ResolveTorchInstallAsync( IPyVenvRunner venvRunner, - string multiArchPythonPackageIndexUrl, - string multiArchDeviceExtra, + string indexUrl, + string? deviceExtra, bool usesNightlyIndex, Action? onConsoleOutput, CancellationToken cancellationToken ) { - var torchSpecifierBase = $"torch[{multiArchDeviceExtra}]"; - var torchvisionSpecifierBase = $"torchvision[{multiArchDeviceExtra}]"; + var torchSpecifierBase = BuildTorchSpecifier("torch", deviceExtra); + var torchvisionSpecifierBase = BuildTorchSpecifier("torchvision", deviceExtra); try { // torch: finals only on the stable channel, pre/dev builds on the nightly channel. var torchIndex = await venvRunner - .PipIndex("torch", multiArchPythonPackageIndexUrl, includePrerelease: usesNightlyIndex) + .PipIndex("torch", indexUrl, includePrerelease: usesNightlyIndex) .ConfigureAwait(false); cancellationToken.ThrowIfCancellationRequested(); // torchvision: allow pre-releases so the matched alpha (e.g. 0.29.0a0) is visible. var torchvisionIndex = await venvRunner - .PipIndex("torchvision", multiArchPythonPackageIndexUrl, includePrerelease: true) + .PipIndex("torchvision", indexUrl, includePrerelease: true) .ConfigureAwait(false); cancellationToken.ThrowIfCancellationRequested(); @@ -506,7 +577,7 @@ private RocmCompatibilityResult BuildCompatibilityResult() }; } - private RocmMachineState ResolveWindowsMachineState() + private RocmMachineState ResolveMachineState() { if (settingsManager.Settings.PreferredGpu is { } preferredGpu && !preferredGpu.IsAmd) { @@ -529,21 +600,34 @@ private RocmMachineState ResolveWindowsMachineState() }; } - var supportedAmdGpus = amdGpus.Where(WindowsRocmSupport.IsSupportedGpu).ToList(); - if (supportedAmdGpus.Count == 0) + var supportedAmdGpus = amdGpus.Where(RocmSupport.IsSupportedGpu).ToList(); + var linuxDevice = TrySelectLinuxDevice(); + var detectedArch = linuxDevice is null + ? null + : RocmSupport.TryGetCanonicalArchitecture(linuxDevice.GfxArch); + // Prefer the driver's count when it has one, since it counts GPU agents rather than adapters. + var driverDeviceCount = RocmArchDetector.Devices.Count; + var amdDeviceCount = driverDeviceCount > 0 ? driverDeviceCount : amdGpus.Count; + + // The driver can resolve an architecture the adapter-name table has never heard of, so only fail + // when neither source produced one. + if (supportedAmdGpus.Count == 0 && detectedArch is null) { return new RocmMachineState { IsCompatible = false, - FailureReason = "No AMD GPU with a supported Windows ROCm architecture was detected.", + FailureReason = "No AMD GPU with a supported ROCm architecture was detected.", }; } var selectedGpu = TryResolvePreferredAmdGpu(supportedAmdGpus, settingsManager.Settings.PreferredGpu) - ?? supportedAmdGpus.First(); + ?? SelectPreferredGpu(supportedAmdGpus) + ?? SelectPreferredGpu(amdGpus) + ?? amdGpus[0]; var runtimeGfxArch = - WindowsRocmSupport.TryGetCanonicalArchitecture(selectedGpu.GetAmdGfxArch()) + detectedArch + ?? RocmSupport.TryGetCanonicalArchitecture(selectedGpu.GetAmdGfxArch()) ?? GetSupportedFallbackGfxArch(supportedAmdGpus); var isCompatible = !string.IsNullOrWhiteSpace(runtimeGfxArch); @@ -555,10 +639,36 @@ private RocmMachineState ResolveWindowsMachineState() : "No supported AMD GFX architecture could be resolved for ROCm.", SelectedGpu = selectedGpu, RuntimeGfxArch = runtimeGfxArch, - MultiArchDeviceExtra = WindowsRocmSupport.TryGetMultiArchDeviceExtra(runtimeGfxArch), + AmdDeviceCount = amdDeviceCount, + DeviceIndex = linuxDevice?.DeviceIndex, }; } + /// + /// ISA target of the GPU the kernel reports through KFD, or null when KFD is unavailable or reports + /// nothing supported. The driver's answer takes precedence over the adapter-name table on Linux. + /// + private static RocmDetectedDevice? TrySelectLinuxDevice() + { + if (!OperatingSystem.IsLinux()) + { + return null; + } + + var supportedDevices = RocmArchDetector + .Devices.Where(device => RocmSupport.TryGetCanonicalArchitecture(device.GfxArch) is not null) + .ToList(); + + return RocmArchDetector.SelectPreferred(supportedDevices); + } + + /// + /// Picks between candidate adapters when the user has no preference: the one with the most memory, + /// which separates a discrete card from an APU's shared carveout. + /// + internal static GpuInfo? SelectPreferredGpu(IEnumerable candidates) => + candidates.OrderByDescending(gpu => gpu.MemoryBytes).FirstOrDefault(); + /// /// Returns AMD GPUs from Stability Matrix's internal hardware model. /// This is the canonical GPU source for the ROCm helper and intentionally avoids package-local probing. @@ -600,15 +710,15 @@ private static IReadOnlyList GetAmdGpuCandidates() private static string? GetSupportedFallbackGfxArch(IEnumerable availableGpus) { return availableGpus - .Where(WindowsRocmSupport.IsSupportedGpu) - .Select(gpu => WindowsRocmSupport.TryGetCanonicalArchitecture(gpu.GetAmdGfxArch())) - .FirstOrDefault(WindowsRocmSupport.IsSupportedArchitecture); + .Where(RocmSupport.IsSupportedGpu) + .Select(gpu => RocmSupport.TryGetCanonicalArchitecture(gpu.GetAmdGfxArch())) + .FirstOrDefault(RocmSupport.IsSupportedArchitecture); } /// /// Verifies that the installed torch build still reports usable ROCm metadata after helper-managed installs complete. /// - private static async Task VerifyWindowsNativeTorchInstallAsync( + private static async Task VerifyRocmTorchInstallAsync( IPyVenvRunner venvRunner, Action? onConsoleOutput, CancellationToken cancellationToken @@ -619,7 +729,7 @@ CancellationToken cancellationToken var torchInfo = await venvRunner.PipShow("torch").ConfigureAwait(false); if (torchInfo is null) { - throw new InvalidOperationException("torch was not installed after Windows ROCm setup."); + throw new InvalidOperationException("torch was not installed after ROCm setup."); } var verificationResult = await venvRunner @@ -672,14 +782,14 @@ CancellationToken cancellationToken var hipVersion = root.TryGetProperty("hip", out var hipElement) ? hipElement.GetString() : null; var cudaAvailable = root.TryGetProperty("cuda", out var cudaElement) && cudaElement.GetBoolean(); - if (!IsUsableWindowsNativeTorchBuild(version, hipVersion)) + if (!IsUsableRocmTorchBuild(version, hipVersion)) { throw new InvalidOperationException( $"Installed torch is not a usable ROCm build. Verification output: {verificationOutput}" ); } - // AMD's torchvision wheels declare no torch version constraint, so an ABI-mismatched pair only + // If AMD's torchvision wheels declare no torch version constraint, so an ABI-mismatched pair only // fails at import time. Importing it here surfaces that as an install failure instead of a // confusing runtime error later. if (string.IsNullOrWhiteSpace(torchvisionVersion)) @@ -706,7 +816,7 @@ CancellationToken cancellationToken } } - internal static bool IsUsableWindowsNativeTorchBuild(string? version, string? hipVersion) + internal static bool IsUsableRocmTorchBuild(string? version, string? hipVersion) { if (!string.IsNullOrWhiteSpace(hipVersion)) return true; @@ -739,6 +849,87 @@ private static bool HasNightlyBuildDateToken(string? version, string nightlyBuil && version.Contains($"dev{nightlyBuildDateToken}", StringComparison.OrdinalIgnoreCase); } + /// + /// ROCm version a torch build was compiled against, taken from its local label: 2.14.0+rocm10.0.0 + /// becomes 10.0.0. Returns null when the build carries no rocm local label. + /// + internal static string? TryGetRocmVersionLabel(string? torchVersion) + { + if (string.IsNullOrWhiteSpace(torchVersion)) + { + return null; + } + + var plusIndex = torchVersion.IndexOf('+'); + if (plusIndex < 0 || plusIndex == torchVersion.Length - 1) + { + return null; + } + + var localLabel = torchVersion[(plusIndex + 1)..].Trim(); + const string prefix = "rocm"; + if (!localLabel.StartsWith(prefix, StringComparison.OrdinalIgnoreCase)) + { + return null; + } + + var version = localLabel[prefix.Length..].Trim(); + return version.Length == 0 ? null : version; + } + + /// True when a rocm-sdk-devel build sits on the given ROCm version line. + internal static bool IsRocmSdkDevelMatch(string? sdkDevelVersion, string? rocmVersion) + { + return !string.IsNullOrWhiteSpace(sdkDevelVersion) + && !string.IsNullOrWhiteSpace(rocmVersion) + && string.Equals(sdkDevelVersion.Trim(), rocmVersion.Trim(), StringComparison.OrdinalIgnoreCase); + } + + /// A rocm-sdk-devel build to install, and whether it matched the torch build exactly. + internal readonly record struct SdkDevelSelection(string Version, bool IsExactMatch); + + /// + /// Chooses the rocm-sdk-devel build for a torch install: the nightly build date token first, then the + /// ROCm line of the torch wheel, and the newest published build only when neither matches. + /// + internal static SdkDevelSelection? SelectSdkDevelVersion( + IReadOnlyList? availableVersions, + string? rocmVersion, + string? nightlyBuildDateToken + ) + { + if (availableVersions is null || availableVersions.Count == 0) + { + return null; + } + + if (!string.IsNullOrWhiteSpace(nightlyBuildDateToken)) + { + var nightly = availableVersions.FirstOrDefault(version => + HasNightlyBuildDateToken(version, nightlyBuildDateToken) + ); + + if (nightly is not null) + { + return new SdkDevelSelection(nightly, true); + } + } + + if (!string.IsNullOrWhiteSpace(rocmVersion)) + { + var exact = availableVersions.FirstOrDefault(version => + IsRocmSdkDevelMatch(version, rocmVersion) + ); + + if (exact is not null) + { + return new SdkDevelSelection(exact, true); + } + } + + return new SdkDevelSelection(availableVersions[0], false); + } + internal static string? TryExtractJsonObject(string output) { if (string.IsNullOrWhiteSpace(output)) @@ -775,10 +966,35 @@ RocmPackageProfile profile var gfxArch = runtimeContext.RuntimeGfxArch; ApplyDefaultLaunchEnvironment(environment, gfxArch, options); + ApplyDeviceVisibility(environment, runtimeContext); return environment; } + /// + /// Pins the runtime to the GPU the helper selected when the machine exposes more than one, so an APU + /// that enumerates first cannot become the device an app opens. ROCR_VISIBLE_DEVICES is the + /// lowest device filter and the runtime renumbers what remains from zero. Single-GPU machines are + /// left alone, as are machines whose selection has no device index. + /// + internal static void ApplyDeviceVisibility( + IDictionary environment, + RocmRuntimeContext runtimeContext + ) + { + if (runtimeContext.AmdDeviceCount <= 1 || runtimeContext.DeviceIndex is not { } deviceIndex) + { + return; + } + + environment["ROCR_VISIBLE_DEVICES"] = deviceIndex.ToString(); + Logger.Info( + "Pinning ROCm device visibility to device {Index} of {Count}.", + deviceIndex, + runtimeContext.AmdDeviceCount + ); + } + private void ApplyDefaultLaunchEnvironment( IDictionary environment, string? gfxArch, @@ -791,12 +1007,12 @@ RocmEnvironmentOptions options SetIfNotNull(environment, "MIOPEN_FIND_ENFORCE", options.MiopenFindEnforce); SetIfNotNull(environment, "PYTORCH_ALLOC_CONF", options.PyTorchAllocConf); - if (options.ApplyAotritonExperimental && WindowsRocmSupport.SupportsAotritonExperimental(gfxArch)) + if (options.ApplyAotritonExperimental && RocmSupport.SupportsAotritonExperimental(gfxArch)) { environment["TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL"] = "1"; } - if (options.ApplyLegacySdpFallback && WindowsRocmSupport.IsLegacyArchitecture(gfxArch)) + if (options.ApplyLegacySdpFallback && RocmSupport.PreferMathSdpaFallback(gfxArch)) { environment["TORCH_BACKENDS_CUDA_FLASH_SDP_ENABLED"] = "0"; environment["TORCH_BACKENDS_CUDA_MEM_EFF_SDP_ENABLED"] = "0"; @@ -834,6 +1050,13 @@ RocmEnvironmentOptions options } } + // Neutralize inherited system ROCm/HIP SDK variables so the venv's self-contained SDK is used. + // User overrides below can still intentionally re-set these keys. + if (options.SanitizeRocmSdkEnvironment) + { + ApplyRocmSdkEnvironmentSanitization(merged); + } + if ( options.IncludeUserOverrides && settingsManager.Settings.EnvironmentVariables is { Count: > 0 } userOverrides @@ -847,4 +1070,16 @@ RocmEnvironmentOptions options return merged; } + + /// + /// Marks the inherited system ROCm/HIP SDK variables for removal from the child process, so the + /// venv's self-contained SDK wins over any system install that happens to be on PATH. + /// + internal static void ApplyRocmSdkEnvironmentSanitization(IDictionary environment) + { + foreach (var key in RocmSupport.ConflictingRocmEnvironmentVariables) + { + environment[key] = ProcessRunner.UnsetEnvironmentVariable; + } + } } diff --git a/StabilityMatrix.Tests/Core/ProcessRunnerTests.cs b/StabilityMatrix.Tests/Core/ProcessRunnerTests.cs index 9309303f9..137fd1086 100644 --- a/StabilityMatrix.Tests/Core/ProcessRunnerTests.cs +++ b/StabilityMatrix.Tests/Core/ProcessRunnerTests.cs @@ -6,6 +6,24 @@ namespace StabilityMatrix.Tests.Core; [TestClass] public class ProcessRunnerTests { + [TestMethod] + public void PrepareEnvironment_RemovesInheritedVariable_WhenUnsetSentinelIsUsed() + { + var info = new ProcessStartInfo(); + info.Environment["ROCM_PATH"] = "/opt/rocm"; + info.Environment["HIP_PATH"] = "/opt/rocm"; + + ProcessRunner.PrepareEnvironment( + info, + "python3", + new Dictionary { ["ROCM_PATH"] = ProcessRunner.UnsetEnvironmentVariable } + ); + + // The sentinel removes the inherited value instead of shadowing it with an empty string. + Assert.IsFalse(info.Environment.ContainsKey("ROCM_PATH")); + Assert.AreEqual("/opt/rocm", info.Environment["HIP_PATH"]); + } + [TestMethod] public void PrepareEnvironment_SanitizesInheritedPythonVariables_ForPythonProcesses() { diff --git a/StabilityMatrix.Tests/Core/RocmPackageHelperTests.cs b/StabilityMatrix.Tests/Core/RocmPackageHelperTests.cs index 16c52eebf..d5fb8c641 100644 --- a/StabilityMatrix.Tests/Core/RocmPackageHelperTests.cs +++ b/StabilityMatrix.Tests/Core/RocmPackageHelperTests.cs @@ -1,7 +1,9 @@ -using System.Text.Json; +using System.Text.Json; using StabilityMatrix.Core.Helper; using StabilityMatrix.Core.Helper.HardwareInfo; +using StabilityMatrix.Core.Models.Packages; using StabilityMatrix.Core.Models.Rocm; +using StabilityMatrix.Core.Processes; using StabilityMatrix.Core.Python; using StabilityMatrix.Core.Services.Rocm; @@ -11,33 +13,33 @@ namespace StabilityMatrix.Tests.Core; public class RocmPackageHelperTests { [TestMethod] - public void WindowsRocmSupport_TryGetMultiArchDeviceExtra_ReturnsExpectedExtra_ForSupportedArch() + public void RocmSupport_TryGetMultiArchDeviceExtra_ReturnsExpectedExtra_ForSupportedArch() { - var deviceExtra = WindowsRocmSupport.TryGetMultiArchDeviceExtra("gfx1201"); + var deviceExtra = RocmSupport.TryGetMultiArchDeviceExtra("gfx1201"); Assert.AreEqual("device-gfx1201", deviceExtra); } [TestMethod] - public void WindowsRocmSupport_TryGetMultiArchDeviceExtra_ReturnsExpectedExtra_ForCanonicalVega20Arch() + public void RocmSupport_TryGetMultiArchDeviceExtra_ReturnsExpectedExtra_ForCanonicalVega20Arch() { - var deviceExtra = WindowsRocmSupport.TryGetMultiArchDeviceExtra("gfx906"); + var deviceExtra = RocmSupport.TryGetMultiArchDeviceExtra("gfx906"); Assert.AreEqual("device-gfx906", deviceExtra); } [TestMethod] - public void WindowsRocmSupport_TryGetCanonicalArchitecture_ReturnsCanonicalArch_WhenAlreadyCanonical() + public void RocmSupport_TryGetCanonicalArchitecture_ReturnsCanonicalArch_WhenAlreadyCanonical() { - var canonicalArch = WindowsRocmSupport.TryGetCanonicalArchitecture("gfx906"); + var canonicalArch = RocmSupport.TryGetCanonicalArchitecture("gfx906"); Assert.AreEqual("gfx906", canonicalArch); } [TestMethod] - public void IsUsableWindowsNativeTorchBuild_ReturnsTrue_WhenHipMetadataExists() + public void IsUsableRocmTorchBuild_ReturnsTrue_WhenHipMetadataExists() { - var isUsable = RocmPackageHelper.IsUsableWindowsNativeTorchBuild( + var isUsable = RocmPackageHelper.IsUsableRocmTorchBuild( version: "test-version", hipVersion: "test-hip-version" ); @@ -46,9 +48,9 @@ public void IsUsableWindowsNativeTorchBuild_ReturnsTrue_WhenHipMetadataExists() } [TestMethod] - public void IsUsableWindowsNativeTorchBuild_ReturnsTrue_WhenVersionContainsRocm() + public void IsUsableRocmTorchBuild_ReturnsTrue_WhenVersionContainsRocm() { - var isUsable = RocmPackageHelper.IsUsableWindowsNativeTorchBuild( + var isUsable = RocmPackageHelper.IsUsableRocmTorchBuild( version: "test-version+rocm", hipVersion: null ); @@ -57,12 +59,9 @@ public void IsUsableWindowsNativeTorchBuild_ReturnsTrue_WhenVersionContainsRocm( } [TestMethod] - public void IsUsableWindowsNativeTorchBuild_ReturnsFalse_WhenNoRocmMetadataExists() + public void IsUsableRocmTorchBuild_ReturnsFalse_WhenNoRocmMetadataExists() { - var isUsable = RocmPackageHelper.IsUsableWindowsNativeTorchBuild( - version: "test-version", - hipVersion: null - ); + var isUsable = RocmPackageHelper.IsUsableRocmTorchBuild(version: "test-version", hipVersion: null); Assert.IsFalse(isUsable); } @@ -100,19 +99,19 @@ public void TryExtractJsonObject_ReturnsNull_WhenOutputContainsNoJson() } [TestMethod] - public void WindowsRocmSupport_TryGetMultiArchDeviceExtra_ReturnsExpectedExtra_ForKrakenPoint() + public void RocmSupport_TryGetMultiArchDeviceExtra_ReturnsExpectedExtra_ForKrakenPoint() { - var deviceExtra = WindowsRocmSupport.TryGetMultiArchDeviceExtra("gfx1152"); + var deviceExtra = RocmSupport.TryGetMultiArchDeviceExtra("gfx1152"); Assert.AreEqual("device-gfx1152", deviceExtra); } [TestMethod] - public void WindowsRocmSupport_IsSupportedGpu_ReturnsTrue_ForSupportedAmdGpu() + public void RocmSupport_IsSupportedGpu_ReturnsTrue_ForSupportedAmdGpu() { var gpu = new GpuInfo { Name = "AMD Radeon RX 9070 XT", MemoryBytes = 16UL * Size.GiB }; - Assert.IsTrue(WindowsRocmSupport.IsSupportedGpu(gpu)); + Assert.IsTrue(RocmSupport.IsSupportedGpu(gpu)); } [TestMethod] @@ -131,6 +130,396 @@ public void BuildRocmTorchArgs_PinsTorchAndTorchvisionWithDeviceExtras() StringAssert.Contains(text, "torchaudio"); } + [TestMethod] + [DataRow("gfx908")] + [DataRow("gfx90a")] + [DataRow("gfx942")] + [DataRow("gfx950")] + public void RocmSupport_IsSupportedArchitecture_AcceptsCdnaTargets(string arch) + { + Assert.IsTrue(RocmSupport.IsSupportedArchitecture(arch)); + } + + [TestMethod] + [DataRow("AMD Instinct MI100", "gfx908")] + [DataRow("AMD Instinct MI210", "gfx90a")] + public void GpuInfo_GetAmdGfxArch_ResolvesInstinctCdnaTargets(string name, string expected) + { + var gpu = new GpuInfo { Name = name }; + + Assert.AreEqual(expected, gpu.GetAmdGfxArch()); + } + + [TestMethod] + [DataRow("gfx908", true)] + [DataRow("gfx90a", false)] + [DataRow("gfx942", false)] + [DataRow("gfx950", false)] + [DataRow("gfx1030", true)] + [DataRow("gfx1100", false)] + [DataRow("gfx1201", false)] + public void RocmSupport_PreferMathSdpaFallback_MatchesAotritonImageList(string arch, bool expectFallback) + { + // AOTriton ships images for gfx90a, gfx942, gfx950, gfx11xx and gfx120x. Everything else has no + // fused SDPA kernels and must be forced onto the math backend. + Assert.AreEqual(expectFallback, RocmSupport.PreferMathSdpaFallback(arch)); + } + + [TestMethod] + [DataRow(110000, "gfx1100")] + [DataRow(120001, "gfx1201")] + [DataRow(100300, "gfx1030")] + [DataRow(90402, "gfx942")] + [DataRow(90010, "gfx90a")] + [DataRow(110500, "gfx1150")] + public void RocmArchDetector_FormatGfxTarget_DecodesKfdTargetVersion(int version, string expected) + { + Assert.AreEqual(expected, RocmArchDetector.FormatGfxTarget(version)); + } + + [TestMethod] + public void RocmArchDetector_TryParseNode_SkipsCpuNodes() + { + const string cpuNode = "cpu_cores_count 16\nsimd_count 0\ngfx_target_version 0\n"; + + Assert.IsNull(RocmArchDetector.TryParseNode(cpuNode, 0)); + } + + [TestMethod] + public void RocmArchDetector_TryParseNode_ParsesGpuNodes() + { + const string gpuNode = "simd_count 192\ngfx_target_version 110000\ndevice_id 29772\n"; + + var device = RocmArchDetector.TryParseNode(gpuNode, 1); + + Assert.IsNotNull(device); + Assert.AreEqual("gfx1100", device.GfxArch); + Assert.AreEqual(1, device.DeviceIndex); + Assert.AreEqual(192, device.SimdCount); + } + + [TestMethod] + public void RocmArchDetector_SelectPreferred_PrefersTheDiscreteDevice() + { + // An APU is a couple of SIMD units; the discrete card is two orders of magnitude more. Node order + // carries no information about which is which, so both orders have to resolve the same way. + var apuFirst = new[] + { + new RocmDetectedDevice + { + GfxArch = "gfx1036", + DeviceIndex = 0, + SimdCount = 4, + }, + new RocmDetectedDevice + { + GfxArch = "gfx1100", + DeviceIndex = 1, + SimdCount = 192, + }, + }; + var dgpuFirst = new[] + { + new RocmDetectedDevice + { + GfxArch = "gfx1100", + DeviceIndex = 0, + SimdCount = 192, + }, + new RocmDetectedDevice + { + GfxArch = "gfx1036", + DeviceIndex = 1, + SimdCount = 4, + }, + }; + + Assert.AreEqual("gfx1100", RocmArchDetector.SelectPreferred(apuFirst)!.GfxArch); + Assert.AreEqual("gfx1100", RocmArchDetector.SelectPreferred(dgpuFirst)!.GfxArch); + } + + [TestMethod] + public void RocmArchDetector_SelectPreferred_BreaksTiesOnDeviceIndex() + { + var first = new RocmDetectedDevice + { + GfxArch = "gfx1100", + DeviceIndex = 0, + SimdCount = 192, + }; + var second = new RocmDetectedDevice + { + GfxArch = "gfx1100", + DeviceIndex = 1, + SimdCount = 192, + }; + + Assert.AreEqual(0, RocmArchDetector.SelectPreferred([first, second])!.DeviceIndex); + } + + [TestMethod] + public void SelectPreferredGpu_PrefersTheCardWithMoreMemory() + { + // On Windows the APU reports its carveout, so memory is the only discrete signal available. + var apu = new GpuInfo { Name = "AMD Radeon(TM) Graphics", MemoryBytes = 512UL * 1024 * 1024 }; + var dgpu = new GpuInfo { Name = "AMD Radeon RX 7900 XTX", MemoryBytes = 24UL * 1024 * 1024 * 1024 }; + + Assert.AreEqual(dgpu, RocmPackageHelper.SelectPreferredGpu([apu, dgpu])); + Assert.AreEqual(dgpu, RocmPackageHelper.SelectPreferredGpu([dgpu, apu])); + } + + [TestMethod] + [DataRow("2.14.0+rocm10.0.0", "10.0.0")] + [DataRow("2.14.0+rocm10.1.0", "10.1.0")] + [DataRow("2.13.0+rocm7.2.0", "7.2.0")] + [DataRow("2.14.0", null)] + [DataRow("2.14.0+cpu", null)] + [DataRow("2.14.0+cu130", null)] + [DataRow("", null)] + public void TryGetRocmVersionLabel_ReadsTheLocalLabel(string torchVersion, string? expected) + { + Assert.AreEqual(expected, RocmPackageHelper.TryGetRocmVersionLabel(torchVersion)); + } + + [TestMethod] + public void SelectSdkDevelVersion_PrefersTheRocmLineOfTheTorchBuild() + { + // The exact case from a gfx1100 install: torch landed on 10.0.0 because torchvision 0.29.0 only + // exists for that line, and the newest devel on the index was 10.1.0. + string[] available = ["10.1.0", "10.0.0"]; + + var selection = RocmPackageHelper.SelectSdkDevelVersion(available, "10.0.0", null); + + Assert.IsNotNull(selection); + Assert.AreEqual("10.0.0", selection.Value.Version); + Assert.IsTrue(selection.Value.IsExactMatch); + } + + [TestMethod] + public void SelectSdkDevelVersion_PrefersTheNightlyDateToken() + { + string[] available = ["10.1.0.dev20261006", "10.1.0.dev20261001", "10.0.0"]; + + var selection = RocmPackageHelper.SelectSdkDevelVersion(available, "10.1.0", "20261001"); + + Assert.IsNotNull(selection); + Assert.AreEqual("10.1.0.dev20261001", selection.Value.Version); + Assert.IsTrue(selection.Value.IsExactMatch); + } + + [TestMethod] + public void SelectSdkDevelVersion_FallsBackToNewestWhenTheLineIsMissing() + { + string[] available = ["10.1.0"]; + + var selection = RocmPackageHelper.SelectSdkDevelVersion(available, "10.0.0", null); + + Assert.IsNotNull(selection); + Assert.AreEqual("10.1.0", selection.Value.Version); + Assert.IsFalse(selection.Value.IsExactMatch); + } + + [TestMethod] + public void SelectSdkDevelVersion_ReturnsNullWhenNothingIsPublished() + { + Assert.IsNull(RocmPackageHelper.SelectSdkDevelVersion([], "10.0.0", null)); + Assert.IsNull(RocmPackageHelper.SelectSdkDevelVersion(null, "10.0.0", null)); + } + + [TestMethod] + [DataRow("10.0.0", "10.0.0", true)] + [DataRow("10.1.0", "10.0.0", false)] + [DataRow("10.0.0", null, false)] + [DataRow(null, "10.0.0", false)] + public void IsRocmSdkDevelMatch_ComparesTheRocmLine(string? sdkDevel, string? rocm, bool expected) + { + Assert.AreEqual(expected, RocmPackageHelper.IsRocmSdkDevelMatch(sdkDevel, rocm)); + } + + [TestMethod] + public void ResolveTarget_WindowsStable_KeepsAmdFeedAndDeviceExtra() + { + var target = RocmSupport.ResolveTarget("gfx1201"); + + Assert.AreEqual(RocmIndexKind.AmdStable, target.IndexKind); + Assert.AreEqual(RocmSupport.StableMultiArchPythonPackageIndexUrl, target.IndexUrl); + Assert.AreEqual("device-gfx1201", target.DeviceExtra); + Assert.IsTrue(target.SupportsDeviceExtras); + Assert.IsTrue(target.SupportsSdkDevel); + } + + [TestMethod] + public void ResolveTarget_VegaUsesNightlyFeed() + { + var target = RocmSupport.ResolveTarget("gfx906"); + + Assert.AreEqual(RocmIndexKind.AmdNightly, target.IndexKind); + Assert.AreEqual(RocmSupport.NightlyMultiArchPythonPackageIndexUrl, target.IndexUrl); + } + + [TestMethod] + public void ResolveTarget_LinuxAmdFeed_KeepsDeviceExtras() + { + // The feed publishes linux_x86_64 device packages, so the extra applies on both platforms. + var target = RocmSupport.ResolveTarget("gfx1100"); + + Assert.AreEqual(RocmSupport.StableMultiArchPythonPackageIndexUrl, target.IndexUrl); + Assert.AreEqual("device-gfx1100", target.DeviceExtra); + } + + [TestMethod] + public void ResolveTarget_PyTorchOverride_DropsDeviceExtrasAndSdkDevel() + { + var target = RocmSupport.ResolveTarget("gfx1100", RocmIndexKind.PyTorchRocm72); + + Assert.AreEqual("https://download.pytorch.org/whl/rocm7.2", target.IndexUrl); + Assert.IsNull(target.DeviceExtra); + Assert.IsFalse(target.SupportsDeviceExtras); + Assert.IsFalse(target.SupportsSdkDevel); + } + + [TestMethod] + public void ResolveTarget_UnsupportedArch_HasNoDeviceExtra() + { + var target = RocmSupport.ResolveTarget("gfx9999"); + + Assert.IsNull(target.DeviceExtra); + Assert.AreEqual(RocmIndexKind.AmdStable, target.IndexKind); + } + + [TestMethod] + [DataRow("torch", "device-gfx1201", "torch[device-gfx1201]")] + [DataRow("torchvision", "device-gfx1201", "torchvision[device-gfx1201]")] + [DataRow("torch", null, "torch")] + [DataRow("torchvision", "", "torchvision")] + public void BuildTorchSpecifier_UsesDeviceExtrasOnlyWhenTheIndexHasThem( + string package, + string? deviceExtra, + string expected + ) + { + Assert.AreEqual(expected, RocmPackageHelper.BuildTorchSpecifier(package, deviceExtra)); + } + + [TestMethod] + public void ApplyDeviceVisibility_PinsSelectedDeviceWhenSeveralAmdGpusExist() + { + var environment = new Dictionary(StringComparer.OrdinalIgnoreCase); + + RocmPackageHelper.ApplyDeviceVisibility( + environment, + new RocmRuntimeContext + { + IsSupported = true, + AmdDeviceCount = 2, + DeviceIndex = 1, + } + ); + + Assert.AreEqual("1", environment["ROCR_VISIBLE_DEVICES"]); + } + + [TestMethod] + public void ApplyDeviceVisibility_LeavesSingleGpuMachinesAlone() + { + var environment = new Dictionary(StringComparer.OrdinalIgnoreCase); + + RocmPackageHelper.ApplyDeviceVisibility( + environment, + new RocmRuntimeContext + { + IsSupported = true, + AmdDeviceCount = 1, + DeviceIndex = 0, + } + ); + + Assert.IsFalse(environment.ContainsKey("ROCR_VISIBLE_DEVICES")); + } + + [TestMethod] + public void ApplyDeviceVisibility_SkipsWhenThePlatformReportsNoIndex() + { + var environment = new Dictionary(StringComparer.OrdinalIgnoreCase); + + RocmPackageHelper.ApplyDeviceVisibility( + environment, + new RocmRuntimeContext + { + IsSupported = true, + AmdDeviceCount = 2, + DeviceIndex = null, + } + ); + + Assert.IsFalse(environment.ContainsKey("ROCR_VISIBLE_DEVICES")); + } + + [TestMethod] + public void ApplyRocmSdkEnvironmentSanitization_MarksInheritedSdkVariablesForRemoval() + { + var environment = new Dictionary(StringComparer.OrdinalIgnoreCase) + { + ["ROCM_PATH"] = "/opt/rocm", + ["PATH"] = "/usr/bin", + }; + + RocmPackageHelper.ApplyRocmSdkEnvironmentSanitization(environment); + + foreach (var key in RocmSupport.ConflictingRocmEnvironmentVariables) + { + Assert.AreEqual(ProcessRunner.UnsetEnvironmentVariable, environment[key], key); + } + + Assert.AreEqual("/usr/bin", environment["PATH"]); + } + + [TestMethod] + public void ApplyRocmTorchInstallOptions_ScopesUpgradeAndReinstallToTorchTrio() + { + var args = RocmPackageHelper.ApplyRocmTorchInstallOptions( + new PipInstallArgs().AddArg("torch[device-gfx1201]==2.14.0"), + new PipInstallConfig { UpgradePackages = true, ForceReinstallTorch = true }, + pipOverrides: null + ); + + var values = args.Arguments.Select(argument => argument.Value).ToList(); + var expectedTargets = new[] { "torch", "torchvision", "torchaudio" }; + + // A blanket --upgrade would re-resolve numpy (and friends) against the multi-arch index; a blanket + // --force-reinstall would reinstall the whole graph. Neither may come back. + CollectionAssert.DoesNotContain(values, "--upgrade"); + CollectionAssert.DoesNotContain(values, "--force-reinstall"); + + Assert.AreEqual(3, values.Count(value => value == "--upgrade-package")); + Assert.AreEqual(3, values.Count(value => value == "--reinstall-package")); + CollectionAssert.AreEqual(expectedTargets, TargetsOf(values, "--upgrade-package")); + CollectionAssert.AreEqual(expectedTargets, TargetsOf(values, "--reinstall-package")); + } + + [TestMethod] + public void ApplyRocmTorchInstallOptions_WithoutUpgradePackages_ReinstallsTrioWithoutUpgrading() + { + var args = RocmPackageHelper.ApplyRocmTorchInstallOptions( + new PipInstallArgs().AddArg("torch[device-gfx1201]==2.14.0"), + new PipInstallConfig { UpgradePackages = false, ForceReinstallTorch = true }, + pipOverrides: null + ); + + var values = args.Arguments.Select(argument => argument.Value).ToList(); + + CollectionAssert.DoesNotContain(values, "--upgrade"); + CollectionAssert.DoesNotContain(values, "--upgrade-package"); + Assert.AreEqual(3, values.Count(value => value == "--reinstall-package")); + } + + private static string[] TargetsOf(List values, string flag) => + values + .Select((value, index) => (value, index)) + .Where(pair => pair.value == flag) + .Select(pair => values[pair.index + 1]) + .ToArray(); + [TestMethod] public void BuildRocmTorchArgs_UserTorchOverrideReplacesPinnedTorch() { diff --git a/StabilityMatrix.Tests/Models/Packages/ComfyUIInstallConfigTests.cs b/StabilityMatrix.Tests/Models/Packages/ComfyUIInstallConfigTests.cs index 1b5e84fa6..112b80212 100644 --- a/StabilityMatrix.Tests/Models/Packages/ComfyUIInstallConfigTests.cs +++ b/StabilityMatrix.Tests/Models/Packages/ComfyUIInstallConfigTests.cs @@ -1,4 +1,4 @@ -using NSubstitute; +using NSubstitute; using StabilityMatrix.Core.Helper; using StabilityMatrix.Core.Helper.Cache; using StabilityMatrix.Core.Models.Packages; @@ -58,7 +58,7 @@ public void Python313_PreservesTorchAndRequirementsConfiguration(bool isLegacyNv CollectionAssert.AreEqual(new[] { "requirements.txt" }, config.RequirementsFilePaths.ToArray()); Assert.AreEqual(cudaIndex, config.CudaIndex); - Assert.AreEqual("rocm7.2", config.RocmIndex); + Assert.AreEqual("rocm7.14", config.RocmIndex); Assert.AreEqual("xpu", config.XpuIndex); Assert.AreEqual(" ", config.TorchaudioVersion); Assert.IsTrue(config.UpgradePackages);