Platform-agnostic ROCm Helper (Win/Linux) with hardened GPU detection + Enhancements

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") =>

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<bool> IsMemoryInfoAvailableLazy = new(() => TryGetMemoryInfo(out _));
     public static bool IsMemoryInfoAvailable => IsMemoryInfoAvailableLazy.Value;

StabilityMatrix.Core/Models/PackageModification/InstallWindowsRocmPackageCommandStep.cs → 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,21 +11,23 @@
 
 namespace StabilityMatrix.Core.Models.PackageModification;
 
-public enum WindowsRocmPackageCommandType
+public enum RocmPackageCommandType
 {
     SageAttention,
     DevelopmentSdk,
     BitsAndBytes,
     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<string, string>? 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<ProgressReport>? 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);
+    }
+
+    /// <summary>
+    /// Runs <c>rocm-sdk init</c> 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.
+    /// </summary>
+    private async Task TryInitializeRocmSdkAsync(
+        IPyVenvRunner venvRunner,
+        IProgress<ProgressReport>? 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(

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();
     }