vulkan-utils-spirv: New package

.github/workflows/ci.yml

@@ -263,6 +263,7 @@ jobs:
     - run: nix-build --arg compiler '${{ matrix.ghc }}' -A vulkan
     - run: nix-build --arg compiler '${{ matrix.ghc }}' -A VulkanMemoryAllocator
     - run: nix-build --arg compiler '${{ matrix.ghc }}' -A vulkan-utils
+    - run: nix-build --arg compiler '${{ matrix.ghc }}' -A vulkan-utils-spirv
     - run: nix-build --arg compiler '${{ matrix.ghc }}' -A openxr
     - run: nix-build --arg compiler '${{ matrix.ghc }}' ./nix/haskell-packages.nix --arg openxrNoVulkan true -A openxr
     - run: nix-build --arg compiler '${{ matrix.ghc }}' -A vulkan-examples
@@ -328,7 +329,8 @@ jobs:
         sudo apt-get install fd-find
     - name: Check scripts with shellcheck
       run: |
-        fdfind .sh$ . \
+        fdfind '\.sh$' . \
+          --type f \
           --exclude VulkanMemoryAllocator/VulkanMemoryAllocator \
           --exclude generate-new/Vulkan-Docs \
           --exclude generate-new/OpenXR-Docs \

cabal.project

@@ -3,6 +3,7 @@ packages:
   ./openxr
   ./VulkanMemoryAllocator
   ./utils
+  ./utils-spirv
   ./utils-init/vulkan-init-sdl2
   ./utils-init/vulkan-init-glfw
   ./examples

default.nix

@@ -21,7 +21,7 @@ let
       (pkgs.haskell.lib.dontCheck vulkan)
       (pkgs.haskell.lib.dontCheck VulkanMemoryAllocator)
     ] else
-      [ vulkan vulkan-utils vulkan-init-sdl2 vulkan-init-glfw VulkanMemoryAllocator vulkan-examples openxr ]
+      [ vulkan vulkan-utils vulkan-utils-spirv vulkan-init-sdl2 vulkan-init-glfw VulkanMemoryAllocator vulkan-examples openxr ]
       ++ pkgs.lib.optional (p.ghc.version == generator-ghc-version) generate-new;
 
 in if forShell then

examples/compute-reflect/Julia.hs

@@ -0,0 +1,101 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE DeriveAnyClass #-}
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE DerivingVia #-}
+{-# LANGUAGE OverloadedLists #-}
+{-# LANGUAGE TemplateHaskell #-}
+{-# LANGUAGE TypeFamilies #-}
+
+{-| Julia-set compute pipeline, its interface reflected from 'Julia.Shader.code':
+
+  * 'reflectShaderTypesBytes' generates the 'Params' push-constant record (with
+    a gl-block std430 'Storable');
+  * 'singleDescriptorSetLayoutInfo' builds the descriptor set layout for the
+    output SSBO;
+  * 'pushConstantRanges' builds the pipeline layout's push-constant range; and
+  * 'allocateSpecializationInfo' builds the pipeline's specialization info
+    (@maxIterations@, @escapeRadius@) from the reflected constant ids.
+
+Compare with the @compute@ example, which hand-writes all of this.
+-}
+module Julia
+  ( Params (..)
+  , Pipeline (..)
+  , allocatePipeline
+  , workgroup
+  ) where
+
+import Control.Monad.Trans.Resource (ResourceT, allocate)
+import qualified Data.Vector as V
+import Data.Word (Word32)
+import qualified Geomancy
+import Graphics.Gl.Block (Std430 (..))
+import Vulkan.CStruct.Extends (SomeStruct (..))
+import qualified Vulkan.Core10 as PipelineLayoutCreateInfo (PipelineLayoutCreateInfo (..))
+import qualified Vulkan.Core10 as Vk
+import Vulkan.Utils.Shader (shaderModuleStage)
+import Vulkan.Utils.SpirV.Descriptors (pushConstantRanges, singleDescriptorSetLayoutInfo)
+import Vulkan.Utils.SpirV.Reflect (reflectBytes)
+import Vulkan.Utils.SpirV.Specialization (allocateSpecializationInfo)
+import Vulkan.Utils.SpirV.TH (reflectShaderTypesBytes)
+import Vulkan.Zero (zero)
+
+import qualified Julia.Shader as Shader
+
+-- Generate the @Params@ push-constant record (and its std430 'Storable') from
+-- the same SPIR-V the runtime loads.
+reflectShaderTypesBytes Shader.code
+
+-- | Workgroup size on each axis (matches @local_size_x/y@ in the shader).
+workgroup :: Int
+workgroup = 16
+
+data Pipeline = Pipeline
+  { pipeline :: Vk.Pipeline
+  , pipelineLayout :: Vk.PipelineLayout
+  , descriptorSetLayout :: Vk.DescriptorSetLayout
+  }
+
+{- | The pipeline, specialized to the given iteration cap and escape radius.
+
+Descriptor set layout, push-constant range and specialization info all come
+from reflecting 'Shader.code'.
+-}
+allocatePipeline :: Vk.Device -> Word32 -> Float -> ResourceT IO Pipeline
+allocatePipeline dev maxIterations escapeRadius = do
+  -- Reflect the embedded module once; reuse it for the descriptor set layout,
+  -- the push-constant range and the specialization info.
+  reflected <- reflectBytes Shader.code
+
+  setLayoutInfo <- either fail pure (singleDescriptorSetLayoutInfo reflected)
+  (_, descriptorSetLayout) <- Vk.withDescriptorSetLayout dev setLayoutInfo Nothing allocate
+
+  mSpec <- allocateSpecializationInfo reflected (maxIterations, escapeRadius)
+  (_, shader) <- shaderModuleStage dev Vk.SHADER_STAGE_COMPUTE_BIT mSpec Shader.code
+  (_, pipelineLayout) <-
+    Vk.withPipelineLayout
+      dev
+      zero
+        { PipelineLayoutCreateInfo.setLayouts = [descriptorSetLayout]
+        , PipelineLayoutCreateInfo.pushConstantRanges =
+            V.fromList (pushConstantRanges reflected)
+        }
+      Nothing
+      allocate
+  let
+    pipelineCreateInfo :: Vk.ComputePipelineCreateInfo '[]
+    pipelineCreateInfo =
+      zero
+        { Vk.layout = pipelineLayout
+        , Vk.stage = shader
+        , Vk.basePipelineHandle = zero
+        }
+  (_, (_, [computePipeline])) <-
+    Vk.withComputePipelines dev zero [SomeStruct pipelineCreateInfo] Nothing allocate
+  pure
+    Pipeline
+      { pipeline = computePipeline
+      , pipelineLayout = pipelineLayout
+      , descriptorSetLayout = descriptorSetLayout
+      }

examples/compute-reflect/Julia/Shader.hs

@@ -0,0 +1,75 @@
+{-# LANGUAGE QuasiQuotes #-}
+
+{-| The Julia-set compute shader, compiled to SPIR-V at build time.
+
+Its interface is reflected by vulkan-utils-spirv in the "Julia" module:
+the @Params@ push-constant block, the output SSBO's descriptor set layout,
+and the specialization constants.
+-}
+module Julia.Shader
+  ( code
+  ) where
+
+import Data.ByteString (ByteString)
+import Vulkan.Utils.ShaderQQ.GLSL.Glslang (comp)
+
+code :: ByteString
+code =
+  [comp|
+    #version 450
+    #extension GL_ARB_separate_shader_objects : enable
+
+    layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in;
+
+    // Specialized at pipeline-creation time (see Vulkan.Utils.SpirV.Specialization).
+    layout(constant_id = 0) const uint  maxIterations = 1000u;
+    layout(constant_id = 1) const float escapeRadius = 2.0;
+
+    // Fields are ordered by non-increasing alignment (vec2/uvec2 = 8) so gl-block's
+    // Generic std430 layout matches the shader; see the guardrail in
+    // Vulkan.Utils.SpirV.Block.
+    layout(push_constant, std430) uniform Params {
+      vec2  center;      // c in the Julia iteration
+      uvec2 resolution;  // output size in pixels
+    } params;
+
+    layout(set = 0, binding = 0, std430) buffer OutputBuffer {
+      vec4 pixels[];
+    };
+
+    // From https://iquilezles.org/articles/palettes/
+    vec3 palette(float t) {
+      const vec3 a = vec3(0.5);
+      const vec3 b = vec3(0.5);
+      const vec3 c = vec3(1.0);
+      const vec3 d = vec3(0.0, 0.10, 0.20);
+      return a + b * cos(6.28318530718 * (c * t + d));
+    }
+
+    void main() {
+      uvec2 gid = gl_GlobalInvocationID.xy;
+      if (gid.x >= params.resolution.x || gid.y >= params.resolution.y) {
+        return;
+      }
+
+      vec2 uv =
+        ( vec2(gid) / vec2(params.resolution) * 2.0 - 1.0 ) * escapeRadius;
+
+      vec2 z = uv;
+      int i = 0;
+      for (; i < int(maxIterations); ++i) {
+        vec2 z2 = vec2(z.x * z.x - z.y * z.y, 2.0 * z.x * z.y);
+        z = z2 + params.center;
+        if (dot(z, z) > escapeRadius * escapeRadius) {
+          break;
+        }
+      }
+
+      uint idx = gid.y * params.resolution.x + gid.x;
+      if (i == int(maxIterations)) {
+        pixels[idx] = vec4(0.0, 0.0, 0.0, 1.0);
+      } else {
+        pixels[idx] = vec4(palette(float(i) / float(maxIterations)), 1.0);
+      }
+    }
+  |]

examples/compute-reflect/Main.hs

@@ -0,0 +1,34 @@
+{-| A compute example whose shader interface is derived from the quasiquoted
+shader in "Julia.Shader" at build time by @vulkan-utils-spirv@ — see "Julia"
+for the reflected pipeline and "Render" for the frame.
+-}
+module Main
+  ( main
+  )
+where
+
+import Control.Monad.Trans.Resource (runResourceT)
+import HeadlessBoot (HeadlessConfig (..), HeadlessVk (..), withHeadlessVk)
+import ImageReadback (savePng)
+import qualified Vulkan.Core10 as Vk
+import Vulkan.Utils.QueueAssignment (QueueFamilyIndex (..))
+import Vulkan.Utils.Queues (Queues (..))
+import Vulkan.Zero (zero)
+
+import qualified Render
+
+main :: IO ()
+main = runResourceT $ do
+  HeadlessVk{..} <-
+    withHeadlessVk
+      HeadlessConfig
+        { appName = "Haskell Vulkan compute-reflect example"
+        , instanceReqs = []
+        , deviceReqs = []
+        , vmaFlags = zero
+        }
+  let QueueFamilyIndex computeQueueFamilyIndex = fst (qCompute queues)
+
+  image <- Render.render allocator device computeQueueFamilyIndex
+  Vk.deviceWaitIdle device
+  savePng "julia-reflect.png" image

examples/compute-reflect/Render.hs

@@ -0,0 +1,141 @@
+{-# LANGUAGE OverloadedLists #-}
+{-# LANGUAGE OverloadedRecordDot #-}
+
+-- | The single headless frame: dispatch the Julia pipeline and read the image back.
+module Render
+  ( render
+  ) where
+
+import qualified Codec.Picture as JP
+import Control.Monad.IO.Class (liftIO)
+import Control.Monad.Trans.Resource (ResourceT, allocate)
+import Data.Word (Word32)
+import Foreign.Marshal.Array (peekArray)
+import Foreign.Marshal.Utils (with)
+import Foreign.Ptr (Ptr, castPtr, plusPtr)
+import Foreign.Storable (sizeOf)
+import Geomancy.UVec2 (uvec2)
+import Geomancy.Vec2 (vec2)
+import HeadlessBoot (submitAndWait)
+import ImageReadback (captureImageRGBA8)
+import Vulkan.CStruct.Utils (FixedArray, lowerArrayPtr)
+import qualified Vulkan.Core10 as CommandBufferBeginInfo (CommandBufferBeginInfo (..))
+import qualified Vulkan.Core10 as CommandPoolCreateInfo (CommandPoolCreateInfo (..))
+import qualified Vulkan.Core10 as Vk
+import Vulkan.Utils.Descriptors (bufferWrite)
+import Vulkan.Zero (zero)
+import qualified VulkanMemoryAllocator as VMA
+
+import Julia (Params (..))
+import qualified Julia
+
+render
+  :: VMA.Allocator
+  -> Vk.Device
+  -> Word32
+  -> ResourceT IO (JP.Image JP.PixelRGBA8)
+render allocator dev computeQueueFamilyIndex = do
+  let
+    width, height :: Int
+    width = 512
+    height = width
+
+    -- Push-constant block (reflected as 'Params').
+    params :: Params
+    params =
+      Params
+        { center = vec2 (-0.8) 0.156
+        , resolution = uvec2 (fromIntegral width) (fromIntegral height)
+        }
+
+    -- Specialization constants, in ascending constant_id order:
+    -- id 0 = maxIterations (uint), id 1 = escapeRadius (float).
+    maxIterations :: Word32
+    maxIterations = 1000
+    escapeRadius :: Float
+    escapeRadius = 2.0
+
+  -- Output storage buffer: one RGBA32F texel per pixel, mapped GPU_TO_CPU.
+  (_, (outBuffer, outAllocation, outInfo)) <-
+    VMA.withBuffer
+      allocator
+      zero
+        { Vk.size = fromIntegral $ width * height * 4 * sizeOf (0 :: Float)
+        , Vk.usage = Vk.BUFFER_USAGE_STORAGE_BUFFER_BIT
+        }
+      zero
+        { VMA.flags = VMA.ALLOCATION_CREATE_MAPPED_BIT
+        , VMA.usage = VMA.MEMORY_USAGE_GPU_TO_CPU
+        }
+      allocate
+
+  julia <- Julia.allocatePipeline dev maxIterations escapeRadius
+
+  (_, descriptorPool) <-
+    Vk.withDescriptorPool
+      dev
+      zero
+        { Vk.maxSets = 1
+        , Vk.poolSizes =
+            [ Vk.DescriptorPoolSize Vk.DESCRIPTOR_TYPE_STORAGE_BUFFER 1
+            ]
+        }
+      Nothing
+      allocate
+  [descriptorSet] <-
+    Vk.allocateDescriptorSets
+      dev
+      zero
+        { Vk.descriptorPool = descriptorPool
+        , Vk.setLayouts = [julia.descriptorSetLayout]
+        }
+
+  Vk.updateDescriptorSets
+    dev
+    [ bufferWrite descriptorSet 0 Vk.DESCRIPTOR_TYPE_STORAGE_BUFFER outBuffer
+    ]
+    []
+
+  (_, commandPool) <-
+    Vk.withCommandPool dev zero{CommandPoolCreateInfo.queueFamilyIndex = computeQueueFamilyIndex} Nothing allocate
+  (_, [cb]) <-
+    Vk.withCommandBuffers
+      dev
+      zero
+        { Vk.commandPool = commandPool
+        , Vk.level = Vk.COMMAND_BUFFER_LEVEL_PRIMARY
+        , Vk.commandBufferCount = 1
+        }
+      allocate
+
+  Vk.useCommandBuffer cb zero{CommandBufferBeginInfo.flags = Vk.COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT} do
+    Vk.cmdBindPipeline cb Vk.PIPELINE_BIND_POINT_COMPUTE julia.pipeline
+    Vk.cmdBindDescriptorSets cb Vk.PIPELINE_BIND_POINT_COMPUTE julia.pipelineLayout 0 [descriptorSet] []
+    -- Push the reflected 'Params' (std430) for this dispatch.
+    liftIO $ with params $ \pParams ->
+      Vk.cmdPushConstants
+        cb
+        julia.pipelineLayout
+        Vk.SHADER_STAGE_COMPUTE_BIT
+        0
+        (fromIntegral (sizeOf params))
+        (castPtr pParams)
+    Vk.cmdDispatch
+      cb
+      (ceiling (realToFrac width / realToFrac @_ @Float Julia.workgroup))
+      (ceiling (realToFrac height / realToFrac @_ @Float Julia.workgroup))
+      1
+
+  computeQueue <- Vk.getDeviceQueue dev computeQueueFamilyIndex 0
+  submitAndWait dev computeQueue cb "Timed out waiting for compute"
+
+  let
+    pixelAddr :: Int -> Int -> Ptr (FixedArray 4 Float)
+    pixelAddr x y =
+      plusPtr
+        (VMA.mappedData outInfo)
+        (((y * width) + x) * 4 * sizeOf (0 :: Float))
+  captureImageRGBA8 allocator outAllocation width height $ \x y -> do
+    let ptr = pixelAddr x y
+    [r, g, b, a] <- fmap (\f -> round (f * 255)) <$> peekArray 4 (lowerArrayPtr ptr)
+    pure $ JP.PixelRGBA8 r g b a