Compute.hs

{-# LANGUAGE OverloadedLists #-}

-- XXX: TypeError in Compatible generates unused constraint argument
{-# OPTIONS_GHC -Wno-redundant-constraints #-}

module Engine.Vulkan.Pipeline.Compute
  ( Config(..)
  , Configure

  , Stages(..)
  , stageNames
  , stageFlagBits
  , StageCode
  , StageSpirv
  , StageReflect

  , Pipeline(..)
  , allocate
  , create
  , destroy

  , bind
  , Compute
  ) where

import RIO

import Data.Kind (Type)
import Data.List qualified as List
import Data.Tagged (Tagged(..))
import Data.Vector qualified as Vector
import GHC.Generics (Generic1)
import GHC.Stack (callStack, getCallStack, srcLocModule, withFrozenCallStack)
import UnliftIO.Resource (MonadResource, ReleaseKey)
import UnliftIO.Resource qualified as Resource
import Vulkan.Core10 qualified as Vk
import Vulkan.Core12.Promoted_From_VK_EXT_descriptor_indexing qualified as Vk12
import Vulkan.CStruct.Extends (SomeStruct(..), pattern (:&), pattern (::&))
import Vulkan.Utils.Debug qualified as Debug
import Vulkan.Zero (Zero(..))

import Engine.SpirV.Reflect (Reflect)
import Engine.Vulkan.DescSets (Bound(..), Compatible)
import Engine.Vulkan.Pipeline (Pipeline(..), destroy)
import Engine.Vulkan.Pipeline.Stages (StageInfo(..))
import Engine.Vulkan.Shader qualified as Shader
import Engine.Vulkan.Types (HasVulkan(..), MonadVulkan, DsBindings, getPipelineCache)
import Render.Code (Code)
import Resource.Collection (Generically1(..))

data Config (dsl :: [Type]) spec = Config
  { cComputeCode        :: ByteString
  , cDescLayouts        :: Tagged dsl [DsBindings]
  , cPushConstantRanges :: Vector Vk.PushConstantRange
  , cSpecialization     :: spec
  }

data Compute

type family Configure pipeline spec where
  Configure (Pipeline dsl Compute Compute) spec = Config dsl spec

newtype Stages a = Stages
  { comp :: a -- ^ compute
  }
  deriving (Eq, Ord, Show, Functor, Foldable, Traversable, Generic1)
  deriving Applicative via (Generically1 Stages)

instance StageInfo Stages where
  stageNames = Stages
    { comp = "comp"
    }

  stageFlagBits = Stages
    { comp = Vk.SHADER_STAGE_COMPUTE_BIT
    }

type StageCode = Stages (Maybe Code)
type StageSpirv = Stages (Maybe ByteString)
type StageReflect = Reflect Stages

allocate
  :: ( MonadVulkan env m
     , MonadResource m
     , HasCallStack
     , Shader.Specialization spec
     )
  => Config dsl spec
  -> m (ReleaseKey, Pipeline dsl Compute Compute)
allocate config = withFrozenCallStack do
  ctx <- ask
  Resource.allocate
    (create ctx config)
    (destroy ctx)

create
  :: ( HasVulkan ctx
     , MonadUnliftIO m
     , Shader.Specialization spec
     )
  => ctx
  -> Config dsl spec
  -> m (Pipeline dsl Compute Compute)
create context Config{..} = do
  -- XXX: copypasta from Pipeline.create
  let
    originModule =
      fromString . List.intercalate "|" $
        map (srcLocModule . snd) (getCallStack callStack)

  dsLayouts <- Vector.forM (Vector.fromList $ unTagged cDescLayouts) \bindsFlags -> do
    let
      (binds, flags) = List.unzip bindsFlags

      setCI =
        zero
          { Vk.bindings = Vector.fromList binds
          }
        ::& zero
          { Vk12.bindingFlags = Vector.fromList flags
          }
        :& ()

    Vk.createDescriptorSetLayout device setCI Nothing

  -- TODO: get from outside
  layout <- Vk.createPipelineLayout device (layoutCI dsLayouts) Nothing
  Debug.nameObject device layout originModule

  -- Compute stuff begins...

  shader <- Shader.withSpecialization cSpecialization $
    Shader.create context Stages
      { comp = Just cComputeCode
      }

  let
    cis = Vector.singleton . SomeStruct $
      pipelineCI (Shader.sPipelineStages shader) layout

  Vk.createComputePipelines device cache cis Nothing >>= \case
    (Vk.SUCCESS, pipelines) ->
      case pipelines of
        [one] -> do
          Shader.destroy context shader
          Debug.nameObject device one originModule
          pure Pipeline
            { pipeline     = one
            , pLayout      = Tagged layout
            , pDescLayouts = Tagged dsLayouts
            }
        _ ->
          error "assert: exactly one pipeline requested"
    (err, _) ->
      throwString $ "createComputePipelines: " <> show err

  where
    device = getDevice context
    cache = getPipelineCache context

    layoutCI dsLayouts = Vk.PipelineLayoutCreateInfo
      { flags              = zero
      , setLayouts         = dsLayouts
      , pushConstantRanges = cPushConstantRanges
      }

    pipelineCI stages layout = zero
      { Vk.layout             = layout
      , Vk.stage              = stage
      , Vk.basePipelineHandle = zero
      }
      where
        stage = case stages of
          [one]   -> one
          _assert -> error "compute code has one stage"

bind
  :: ( Compatible pipeLayout boundLayout
     , MonadIO m
     )
  => Vk.CommandBuffer
  -> Pipeline pipeLayout Compute Compute
  -> Bound boundLayout Compute Compute m ()
  -> Bound boundLayout noVertices noInstances m ()
bind cb Pipeline{pipeline} (Bound attrAction) = do
  Bound $ Vk.cmdBindPipeline cb Vk.PIPELINE_BIND_POINT_COMPUTE pipeline
  Bound attrAction