Parallelize build-time shader compilation in the descriptor generators

src/ComputeSharp.D2D1.SourceGenerators/D2DPixelShaderDescriptorGenerator.cs

@@ -27,7 +27,8 @@ public sealed partial class D2DPixelShaderDescriptorGenerator : IIncrementalGene
     public void Initialize(IncrementalGeneratorInitializationContext context)
     {
         // Discover all shader types and extract all the necessary info from each of them
-        IncrementalValuesProvider<D2D1ShaderInfo> shaderInfo =
+        // (with the exception of the compiled HLSL bytecode, which is processed later)
+        IncrementalValuesProvider<D2D1ShaderInfo> shaderInfoWithNoHlslBytecode =
             context.SyntaxProvider
             .ForAttributeWithMetadataName(
                 "ComputeSharp.D2D1.D2DGeneratedPixelShaderDescriptorAttribute",
@@ -152,22 +153,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();
 
@@ -192,12 +190,65 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
                         ChannelDepth: channelDepth,
                         PixelOptions: pixelOptions,
                         HlslInfoKey: hlslInfoKey,
-                        HlslInfo: hlslInfo,
+                        HlslInfo: HlslBytecodeInfo.Missing.Instance,
+                        HlslDiagnosticsInfo: hlslDiagnosticsInfo,
                         Diagnostcs: diagnostics.ToImmutable());
                 })
             .WithTrackingName(WellKnownTrackingNames.Execute)
             .Where(static item => item 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 =
+            shaderInfoWithNoHlslBytecode
+            .Select(static (item, _) => item.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
+        IncrementalValuesProvider<D2D1ShaderInfo> shaderInfo =
+            shaderInfoWithNoHlslBytecode
+            .Combine(hlslBytecodeCache)
+            .Select(static (pair, token) =>
+            {
+                D2D1ShaderInfo item = pair.Left;
+
+                HlslBytecodeInfoKey hlslInfoKey = item.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(item.Diagnostcs.AsSpan());
+
+                // Append any diagnostic for the shader compilation
+                HlslBytecodeSyntaxProcessor.GetInfoDiagnostics(item.HlslDiagnosticsInfo!, hlslInfo, diagnostics);
+                HlslBytecodeSyntaxProcessor.GetDoublePrecisionSupportDiagnostics(item.HlslDiagnosticsInfo!, hlslInfo, diagnostics);
+
+                token.ThrowIfCancellationRequested();
+
+                // The diagnostics info is dropped here, as it has served its purpose. This also improves
+                // incrementality, as it holds a reference to the syntax tree of the shader type, which
+                // would otherwise cause spurious changes in the resulting models on unrelated edits.
+                return item with
+                {
+                    HlslInfoKey = hlslInfoKey,
+                    HlslInfo = hlslInfo,
+                    HlslDiagnosticsInfo = null,
+                    Diagnostcs = diagnostics.ToImmutable()
+                };
+            });
+
         // 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.
         IncrementalValuesProvider<EquatableArray<DiagnosticInfo>> diagnosticInfo =

src/ComputeSharp.D2D1.SourceGenerators/Models/D2D1ShaderInfo.cs

@@ -22,6 +22,7 @@ namespace ComputeSharp.D2D1.SourceGenerators.Models;
 /// <param name="PixelOptions">The pixel options used by the shader.</param>
 /// <param name="HlslInfoKey">The key with processed info on the shader.</param>
 /// <param name="HlslInfo">The value with processed info on the shader.</param>
+/// <param name="HlslDiagnosticsInfo">The captured info to synthesize diagnostics for the compiled shader (only present until the bytecode is processed).</param>
 /// <param name="Diagnostcs">The discovered diagnostics, if any.</param>
 internal sealed record D2D1ShaderInfo(
     HierarchyInfo Hierarchy,
@@ -40,4 +41,5 @@ internal sealed record D2D1ShaderInfo(
     D2D1PixelOptions PixelOptions,
     HlslBytecodeInfoKey HlslInfoKey,
     HlslBytecodeInfo HlslInfo,
+    HlslBytecodeDiagnosticsInfo? HlslDiagnosticsInfo,
     EquatableArray<DiagnosticInfo> Diagnostcs) : IConstantBufferInfo;

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,16 @@
+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,
+/// which in turn allows deferring the bytecode compilation (so it can be parallelized).
+/// </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 present.</param>
+internal sealed record HlslBytecodeDiagnosticsInfo(
+    string TypeName,
+    LocationInfo? TypeLocation,
+    bool HasRequiresDoublePrecisionSupportAttribute,
+    LocationInfo? RequiresDoublePrecisionSupportAttributeLocation);

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

@@ -1,8 +1,10 @@
 using System;
 using System.Collections.Immutable;
 using System.ComponentModel;
+using System.Linq;
 using System.Runtime.CompilerServices;
 using System.Threading;
+using System.Threading.Tasks;
 using ComputeSharp.SourceGeneration.Extensions;
 using ComputeSharp.SourceGeneration.Helpers;
 using ComputeSharp.SourceGeneration.Models;
@@ -89,13 +91,82 @@ 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)
+        {
+            // 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();
+
+            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.ToString(),
+            TypeLocation: LocationInfo.From(structDeclarationSymbol),
+            HasRequiresDoublePrecisionSupportAttribute: hasRequiresDoublePrecisionSupportAttribute,
+            RequiresDoublePrecisionSupportAttributeLocation: LocationInfo.From(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 +176,17 @@ public static void GetInfoDiagnostics(
         {
             diagnostic = DiagnosticInfo.Create(
                 HlslBytecodeFailedWithWin32Exception,
-                structDeclarationSymbol,
-                structDeclarationSymbol,
+                diagnosticsInfo.TypeLocation?.ToLocation(),
+                diagnosticsInfo.TypeName,
                 win32Error.HResult,
                 win32Error.Message);
         }
         else if (info is HlslBytecodeInfo.CompilerError fxcError)
         {
             diagnostic = DiagnosticInfo.Create(
                 HlslBytecodeFailedWithCompilationException,
-                structDeclarationSymbol,
-                structDeclarationSymbol,
+                diagnosticsInfo.TypeLocation?.ToLocation(),
+                diagnosticsInfo.TypeName,
                 fxcError.Message);
         }
 
@@ -128,11 +199,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 +213,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?.ToLocation(),
+                diagnosticsInfo.TypeName));
         }
-        else if (hasRequiresDoublePrecisionSupportAttribute && !success.RequiresDoublePrecisionSupport)
+        else if (diagnosticsInfo.HasRequiresDoublePrecisionSupportAttribute && !success.RequiresDoublePrecisionSupport)
         {
             diagnostics.Add(DiagnosticInfo.Create(
                 UnnecessaryRequiresDoublePrecisionSupportAttribute,
-                attributeData!.GetLocation(),
-                structDeclarationSymbol));
+                (diagnosticsInfo.RequiresDoublePrecisionSupportAttributeLocation ?? diagnosticsInfo.TypeLocation)?.ToLocation(),
+                diagnosticsInfo.TypeName));
         }
     }
 

src/ComputeSharp.SourceGeneration/ComputeSharp.SourceGeneration.projitems

@@ -32,6 +32,7 @@
     <Compile Include="$(MSBuildThisFileDirectory)Helpers\ObjectPool{T}.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Models\DiagnosticInfo.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Models\HierarchyInfo.cs" />
+    <Compile Include="$(MSBuildThisFileDirectory)Models\LocationInfo.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Models\TypeInfo.cs" />
   </ItemGroup>
 </Project>

src/ComputeSharp.SourceGeneration/Models/LocationInfo.cs

@@ -0,0 +1,53 @@
+using System.Linq;
+using Microsoft.CodeAnalysis;
+using Microsoft.CodeAnalysis.Text;
+
+namespace ComputeSharp.SourceGeneration.Models;
+
+/// <summary>
+/// A model for a captured source location, to be used within equatable incremental models.
+/// The location is captured by value (ie. with no <see cref="SyntaxTree"/> references), so
+/// that models with a captured location will correctly compare as equal across unrelated
+/// edits (and so that they will never keep alive (or leak) any stale compilation objects).
+/// </summary>
+/// <param name="FilePath">The path of the source file for the referenced location.</param>
+/// <param name="TextSpan">The span for the referenced location.</param>
+/// <param name="LineSpan">The line span for the referenced location.</param>
+internal sealed record LocationInfo(string FilePath, TextSpan TextSpan, LinePositionSpan LineSpan)
+{
+    /// <summary>
+    /// Creates a new <see cref="LocationInfo"/> instance from an input <see cref="Location"/> value.
+    /// </summary>
+    /// <param name="location">The <see cref="Location"/> value to capture, if available.</param>
+    /// <returns>A <see cref="LocationInfo"/> instance for <paramref name="location"/>, if a source location was available.</returns>
+    public static LocationInfo? From(Location? location)
+    {
+        if (location is not { SourceTree: not null })
+        {
+            return null;
+        }
+
+        FileLinePositionSpan lineSpan = location.GetLineSpan();
+
+        return new LocationInfo(lineSpan.Path, location.SourceSpan, lineSpan.Span);
+    }
+
+    /// <summary>
+    /// Creates a new <see cref="LocationInfo"/> instance from an input <see cref="ISymbol"/> value.
+    /// </summary>
+    /// <param name="symbol">The <see cref="ISymbol"/> instance to capture the location for.</param>
+    /// <returns>A <see cref="LocationInfo"/> instance for <paramref name="symbol"/>, if a source location was available.</returns>
+    public static LocationInfo? From(ISymbol symbol)
+    {
+        return From(symbol.Locations.FirstOrDefault());
+    }
+
+    /// <summary>
+    /// Creates a new <see cref="Location"/> instance with the state from this model.
+    /// </summary>
+    /// <returns>A new <see cref="Location"/> instance with the state from this model.</returns>
+    public Location ToLocation()
+    {
+        return Location.Create(FilePath, TextSpan, LineSpan);
+    }
+}