Parallelize build-time shader compilation in the descriptor generators

src/ComputeSharp.D2D1.SourceGenerators/D2DPixelShaderDescriptorGenerator.cs

@@ -26,8 +26,11 @@ public sealed partial class D2DPixelShaderDescriptorGenerator : IIncrementalGene
     /// <inheritdoc/>
     public void Initialize(IncrementalGeneratorInitializationContext context)
     {
-        // Discover all shader types and extract all the necessary info from each of them
-        IncrementalValuesProvider<D2D1ShaderInfo> shaderInfo =
+        // Discover all shader types and extract all the necessary info from each of them (with the exception
+        // of the compiled HLSL bytecode, which is processed later). The info needed to synthesize the diagnostics
+        // for the deferred compilation is returned alongside each model (as it holds syntax tree references, it
+        // is intentionally not part of the model itself, so it never affects the equality of the shader models).
+        IncrementalValuesProvider<(D2D1ShaderInfo Info, HlslBytecodeDiagnosticsInfo DiagnosticsInfo)> shaderInfoWithDiagnosticsInfo =
             context.SyntaxProvider
             .ForAttributeWithMetadataName(
                 "ComputeSharp.D2D1.D2DGeneratedPixelShaderDescriptorAttribute",
@@ -152,22 +155,19 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
 
                     token.ThrowIfCancellationRequested();
 
-                    // As the last steps in the pipeline, try to compile the shader if needed.
-                    // This is done last so that it can be skipped if any errors happened before.
+                    // Prepare the key to compile the shader afterwards. The compilation is deliberately not
+                    // done here: the incremental driver invokes transform callbacks sequentially, so compiling
+                    // here would serialize all shader compilations. Instead, compilation is deferred to a
+                    // dedicated node below, which can process all shaders in the compilation in parallel.
                     HlslBytecodeInfoKey hlslInfoKey = new(
                         hlslSource,
                         effectiveShaderProfile,
                         effectiveCompileOptions,
                         isCompilationEnabled);
 
-                    // Get the existing compiled shader, or compile the processed HLSL code
-                    HlslBytecodeInfo hlslInfo = HlslBytecodeSyntaxProcessor.GetInfo(ref hlslInfoKey, token);
-
-                    token.ThrowIfCancellationRequested();
-
-                    // Append any diagnostic for the shader compilation
-                    HlslBytecodeSyntaxProcessor.GetInfoDiagnostics(typeSymbol, hlslInfo, diagnostics);
-                    HlslBytecodeSyntaxProcessor.GetDoublePrecisionSupportDiagnostics(typeSymbol, hlslInfo, diagnostics);
+                    // Capture the info needed to synthesize the diagnostics for the deferred compilation,
+                    // as they cannot be created later (symbols must not be used past the transform node)
+                    HlslBytecodeDiagnosticsInfo hlslDiagnosticsInfo = HlslBytecodeSyntaxProcessor.GetDiagnosticsInfo(typeSymbol);
 
                     token.ThrowIfCancellationRequested();
 
@@ -176,7 +176,7 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
 
                     token.ThrowIfCancellationRequested();
 
-                    return new D2D1ShaderInfo(
+                    D2D1ShaderInfo info = new(
                         Hierarchy: hierarchyInfo,
                         EffectId: effectId,
                         EffectDisplayName: effectDisplayName,
@@ -192,11 +192,62 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
                         ChannelDepth: channelDepth,
                         PixelOptions: pixelOptions,
                         HlslInfoKey: hlslInfoKey,
-                        HlslInfo: hlslInfo,
+                        HlslInfo: HlslBytecodeInfo.Missing.Instance,
                         Diagnostcs: diagnostics.ToImmutable());
+
+                    return (Info: info, DiagnosticsInfo: hlslDiagnosticsInfo);
                 })
-            .WithTrackingName(WellKnownTrackingNames.Execute)
-            .Where(static item => item is not null)!;
+            .Where(static item => item.Info is not null);
+
+        // Compile all shaders in parallel in a single dedicated node, warming up the shared bytecode
+        // cache. The node produces no meaningful value: it only exists so that the join node below has
+        // an edge ordering it after all compilations are done (its input requires this node's output).
+        IncrementalValueProvider<bool> hlslBytecodeCache =
+            shaderInfoWithDiagnosticsInfo
+            .Select(static (item, _) => item.Info.HlslInfoKey)
+            .Collect()
+            .Select(static (keys, token) =>
+            {
+                HlslBytecodeSyntaxProcessor.CompileAllInParallel(keys, token);
+
+                return true;
+            });
+
+        // Join each shader with its compiled bytecode (guaranteed to be a cache hit, given the ordering edge on
+        // the node above), and synthesize the deferred diagnostics for the shader compilation. The resulting
+        // models are the same that used to be produced directly by the transform node (hence the tracking name).
+        IncrementalValuesProvider<D2D1ShaderInfo> shaderInfo =
+            shaderInfoWithDiagnosticsInfo
+            .Combine(hlslBytecodeCache)
+            .Select(static (pair, token) =>
+            {
+                (D2D1ShaderInfo info, HlslBytecodeDiagnosticsInfo diagnosticsInfo) = pair.Left;
+
+                HlslBytecodeInfoKey hlslInfoKey = info.HlslInfoKey;
+
+                // Get the compiled shader from the warmed up cache
+                HlslBytecodeInfo hlslInfo = HlslBytecodeSyntaxProcessor.GetInfo(ref hlslInfoKey, token);
+
+                token.ThrowIfCancellationRequested();
+
+                using ImmutableArrayBuilder<DiagnosticInfo> diagnostics = new();
+
+                diagnostics.AddRange(info.Diagnostcs.AsSpan());
+
+                // Append any diagnostic for the shader compilation
+                HlslBytecodeSyntaxProcessor.GetInfoDiagnostics(diagnosticsInfo, hlslInfo, diagnostics);
+                HlslBytecodeSyntaxProcessor.GetDoublePrecisionSupportDiagnostics(diagnosticsInfo, hlslInfo, diagnostics);
+
+                token.ThrowIfCancellationRequested();
+
+                return info with
+                {
+                    HlslInfoKey = hlslInfoKey,
+                    HlslInfo = hlslInfo,
+                    Diagnostcs = diagnostics.ToImmutable()
+                };
+            })
+            .WithTrackingName(WellKnownTrackingNames.Execute);
 
         // We need to create two more incremental steps to ensure we correctly emit diagnostics and re-generate sources.
         // First, select an incremental provider with just the diagnostics, which will trigger every time any of them changes.

src/ComputeSharp.SourceGeneration.Hlsl/ComputeSharp.SourceGeneration.Hlsl.projitems

@@ -26,6 +26,7 @@
     <Compile Include="$(MSBuildThisFileDirectory)Mappings\HlslKnownSizes.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Mappings\HlslKnownTypes.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Models\FieldInfo.cs" />
+    <Compile Include="$(MSBuildThisFileDirectory)Models\HlslBytecodeDiagnosticsInfo.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Models\HlslBytecodeInfo.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Models\Interfaces\IConstantBufferInfo.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Models\TypeAliases.cs" />

src/ComputeSharp.SourceGeneration.Hlsl/Models/HlslBytecodeDiagnosticsInfo.cs

@@ -0,0 +1,20 @@
+using Microsoft.CodeAnalysis;
+
+namespace ComputeSharp.SourceGeneration.Models;
+
+/// <summary>
+/// A model capturing the info needed to synthesize diagnostics for compiled HLSL bytecode. This makes it possible
+/// to create such diagnostics after the transform node has completed (symbols cannot be used past that point), which
+/// in turn allows deferring the bytecode compilation (so it can be parallelized). Note that this model holds source
+/// locations (which reference syntax trees), so it is deliberately kept separate from the shader models: it only
+/// travels alongside them up to the node that creates the diagnostics, and it is dropped from that point onwards.
+/// </summary>
+/// <param name="TypeName">The fully qualified name of the shader type.</param>
+/// <param name="TypeLocation">The location of the shader type, if available.</param>
+/// <param name="HasRequiresDoublePrecisionSupportAttribute">Whether the shader type is annotated to require double precision support.</param>
+/// <param name="RequiresDoublePrecisionSupportAttributeLocation">The location of the attribute requiring double precision support, if available.</param>
+internal sealed record HlslBytecodeDiagnosticsInfo(
+    string TypeName,
+    Location? TypeLocation,
+    bool HasRequiresDoublePrecisionSupportAttribute,
+    Location? RequiresDoublePrecisionSupportAttributeLocation);

src/ComputeSharp.SourceGeneration.Hlsl/SyntaxProcessors/HlslBytecodeSyntaxProcessor.cs

@@ -1,8 +1,11 @@
 using System;
 using System.Collections.Immutable;
 using System.ComponentModel;
+using System.Linq;
 using System.Runtime.CompilerServices;
+using System.Runtime.ExceptionServices;
 using System.Threading;
+using System.Threading.Tasks;
 using ComputeSharp.SourceGeneration.Extensions;
 using ComputeSharp.SourceGeneration.Helpers;
 using ComputeSharp.SourceGeneration.Models;
@@ -89,13 +92,89 @@ static unsafe HlslBytecodeInfo GetInfo(HlslBytecodeInfoKey key, CancellationToke
     }
 
     /// <summary>
-    /// Gets any diagnostics from a processed <see cref="HlslBytecodeInfo"/> instance.
+    /// Compiles the shaders for all input keys in parallel, warming up the shared cache.
+    /// After this call, <see cref="GetInfo"/> calls for any of the input keys will be cache hits.
+    /// </summary>
+    /// <param name="keys">The <see cref="HlslBytecodeInfoKey"/> instances for the shaders to compile.</param>
+    /// <param name="token">The <see cref="CancellationToken"/> used to cancel the operation, if needed.</param>
+    public static void CompileAllInParallel(ImmutableArray<HlslBytecodeInfoKey> keys, CancellationToken token)
+    {
+        static void Compile(HlslBytecodeInfoKey key, CancellationToken token)
+        {
+            _ = GetInfo(ref key, token);
+        }
+
+        // Skip the parallel dispatch entirely if there are less than two keys to process
+        if (keys.Length == 0)
+        {
+            return;
+        }
+
+        if (keys.Length == 1)
+        {
+            Compile(keys[0], token);
+
+            return;
+        }
+
+        try
+        {
+            // Compile all shaders in parallel (each compilation is independent, and both the shared cache and
+            // the native compilers support concurrent use). Duplicate keys are filtered out first: concurrent
+            // requests for the same key would be benign (one result would just be discarded), but there is no
+            // reason to schedule them at all. The order of compilations does not matter, as the results are
+            // only published to the cache here (callers will then retrieve them via cache hits afterwards).
+            _ = Parallel.ForEach(
+                keys.Distinct(),
+                new ParallelOptions { CancellationToken = token },
+                key => Compile(key, token));
+        }
+        catch (AggregateException e)
+        {
+            // If cancellation is requested, normalize to an OperationCanceledException for the incremental
+            // driver (a cancellation from the callbacks may be wrapped, depending on interleaving). Other
+            // exceptions cannot really occur, as the compilation callback catches all expected exceptions.
+            token.ThrowIfCancellationRequested();
+
+            // If there is a single unexpected exception, unwrap it and rethrow it (preserving its stack
+            // trace), so that the failure reported by the compiler names the actual exception type.
+            if (e.InnerExceptions.Count == 1)
+            {
+                ExceptionDispatchInfo.Capture(e.InnerExceptions[0]).Throw();
+            }
+
+            throw;
+        }
+    }
+
+    /// <summary>
+    /// Gets the <see cref="HlslBytecodeDiagnosticsInfo"/> instance for a given shader type.
+    /// This captures all info needed to synthesize compile diagnostics after the transform
+    /// node has completed (which is required, as symbols cannot be used past that point).
     /// </summary>
     /// <param name="structDeclarationSymbol">The input <see cref="INamedTypeSymbol"/> instance to process.</param>
+    /// <returns>The <see cref="HlslBytecodeDiagnosticsInfo"/> instance for the current shader.</returns>
+    public static HlslBytecodeDiagnosticsInfo GetDiagnosticsInfo(INamedTypeSymbol structDeclarationSymbol)
+    {
+        bool hasRequiresDoublePrecisionSupportAttribute = structDeclarationSymbol.TryGetAttributeWithFullyQualifiedMetadataName(
+            GetRequiresDoublePrecisionSupportAttributeName(),
+            out AttributeData? attributeData);
+
+        return new HlslBytecodeDiagnosticsInfo(
+            TypeName: structDeclarationSymbol.ToDisplayString(),
+            TypeLocation: structDeclarationSymbol.Locations.FirstOrDefault(),
+            HasRequiresDoublePrecisionSupportAttribute: hasRequiresDoublePrecisionSupportAttribute,
+            RequiresDoublePrecisionSupportAttributeLocation: attributeData?.GetLocation());
+    }
+
+    /// <summary>
+    /// Gets any diagnostics from a processed <see cref="HlslBytecodeInfo"/> instance.
+    /// </summary>
+    /// <param name="diagnosticsInfo">The <see cref="HlslBytecodeDiagnosticsInfo"/> instance for the current shader.</param>
     /// <param name="info">The source <see cref="HlslBytecodeInfo"/> instance.</param>
     /// <param name="diagnostics">The collection of produced <see cref="DiagnosticInfo"/> instances.</param>
     public static void GetInfoDiagnostics(
-        INamedTypeSymbol structDeclarationSymbol,
+        HlslBytecodeDiagnosticsInfo diagnosticsInfo,
         HlslBytecodeInfo info,
         ImmutableArrayBuilder<DiagnosticInfo> diagnostics)
     {
@@ -105,17 +184,17 @@ public static void GetInfoDiagnostics(
         {
             diagnostic = DiagnosticInfo.Create(
                 HlslBytecodeFailedWithWin32Exception,
-                structDeclarationSymbol,
-                structDeclarationSymbol,
+                diagnosticsInfo.TypeLocation,
+                diagnosticsInfo.TypeName,
                 win32Error.HResult,
                 win32Error.Message);
         }
         else if (info is HlslBytecodeInfo.CompilerError fxcError)
         {
             diagnostic = DiagnosticInfo.Create(
                 HlslBytecodeFailedWithCompilationException,
-                structDeclarationSymbol,
-                structDeclarationSymbol,
+                diagnosticsInfo.TypeLocation,
+                diagnosticsInfo.TypeName,
                 fxcError.Message);
         }
 
@@ -128,11 +207,11 @@ public static void GetInfoDiagnostics(
     /// <summary>
     /// Gets the diagnostics for when double precision support is configured incorrectly.
     /// </summary>
-    /// <param name="structDeclarationSymbol">The input <see cref="INamedTypeSymbol"/> instance to process.</param>
+    /// <param name="diagnosticsInfo">The <see cref="HlslBytecodeDiagnosticsInfo"/> instance for the current shader.</param>
     /// <param name="info">The source <see cref="HlslBytecodeInfo"/> instance.</param>
     /// <param name="diagnostics">The collection of produced <see cref="DiagnosticInfo"/> instances.</param>
     public static void GetDoublePrecisionSupportDiagnostics(
-        INamedTypeSymbol structDeclarationSymbol,
+        HlslBytecodeDiagnosticsInfo diagnosticsInfo,
         HlslBytecodeInfo info,
         ImmutableArrayBuilder<DiagnosticInfo> diagnostics)
     {
@@ -142,26 +221,22 @@ public static void GetDoublePrecisionSupportDiagnostics(
             return;
         }
 
-        bool hasRequiresDoublePrecisionSupportAttribute = structDeclarationSymbol.TryGetAttributeWithFullyQualifiedMetadataName(
-            GetRequiresDoublePrecisionSupportAttributeName(),
-            out AttributeData? attributeData);
-
         // Check the two cases where diagnostics are necessary:
         //   - The shader does not have [[D2D]RequiresDoublePrecisionSupport], but it needs it
         //   - The shader has [[D2D]RequiresDoublePrecisionSupport], but it does not need it
-        if (!hasRequiresDoublePrecisionSupportAttribute && success.RequiresDoublePrecisionSupport)
+        if (!diagnosticsInfo.HasRequiresDoublePrecisionSupportAttribute && success.RequiresDoublePrecisionSupport)
         {
             diagnostics.Add(DiagnosticInfo.Create(
                 MissingRequiresDoublePrecisionSupportAttribute,
-                structDeclarationSymbol,
-                structDeclarationSymbol));
+                diagnosticsInfo.TypeLocation,
+                diagnosticsInfo.TypeName));
         }
-        else if (hasRequiresDoublePrecisionSupportAttribute && !success.RequiresDoublePrecisionSupport)
+        else if (diagnosticsInfo.HasRequiresDoublePrecisionSupportAttribute && !success.RequiresDoublePrecisionSupport)
         {
             diagnostics.Add(DiagnosticInfo.Create(
                 UnnecessaryRequiresDoublePrecisionSupportAttribute,
-                attributeData!.GetLocation(),
-                structDeclarationSymbol));
+                diagnosticsInfo.RequiresDoublePrecisionSupportAttributeLocation,
+                diagnosticsInfo.TypeName));
         }
     }