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

  , bind
  , Compute
  ) where

import RIO

import Data.Kind (Type)
import Data.Tagged (Tagged(..))
import Data.Vector qualified as Vector
import GHC.Generics (Generic1)
import GHC.Stack (withFrozenCallStack)
import UnliftIO.Resource (MonadResource, ReleaseKey)
import Vulkan.Core10 qualified as Vk
import Vulkan.CStruct.Extends (SomeStruct(..))
import Vulkan.Zero (Zero(..))

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

data Config (dsl :: [Type]) spec = Config
  { cComputeCode        :: ByteString
  , cDescLayouts        :: Tagged dsl [DsLayoutBindings]
  , 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 $
    Pipeline.allocateWith $ create config

create
  :: ( MonadVulkan env io
     , Shader.Specialization spec
     , HasCallStack
     )
  => Config dsl spec
  -> io (Pipeline dsl Compute Compute)
create Config{..} = withFrozenCallStack do
  -- TODO: get from outside ?
  dsLayouts <- Layout.create $ Vector.fromList (unTagged cDescLayouts)

  -- TODO: get from outside ??
  pipelineLayout <- Layout.forPipeline
    dsLayouts
    cPushConstantRanges
  Named.objectOrigin pipelineLayout

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

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

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

  where
    cache = getPipelineCache undefined

    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