fir-0.1.0.0: src/FIR/Layout.hs
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise -fconstraint-solver-iterations=8 #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fconstraint-solver-iterations=10 #-} -- for some reason this needs to be bumped with GHC 9.14
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-|
Module: FIR.Layout
This module is concerned with the memory layout of types for use with Vulkan.
There are three different situations one needs to account for:
- uniform buffers must form a struct or array of structs,
laid out according to the Extended alignment rules,
- push constants and storage buffers must form a struct or array of structs,
laid out according to the Base alignment rules,
- vertex input data is specified with Location and Component information.
This module automatically performs layout of types according to the above rules.
See also "FIR.Validation.Layout", which performs validation of user-supplied layouts.
In the first two situations above, the shader that makes use of such structs
must provide explicit Offset and Stride decorations. These offsets
are computed by this module to avoid the user needing to manually specify them.
In the last situation, per-vertex/per-instance data is laid out in locations,
for use by the vertex shader. See "FIR.Validation.Layout".
For further reference, see:
- SPIR-V specification, 2.16.2 "Validation Rules for Shader Capabilities", bullet point 3,
- SPIR-V specification, 2.18.1 "Memory Layout",
- Vulkan specification, 14.1.4 "Location Assignment", and 14.1.5 "Component Assignment",
- Vulkan specification, 14.5.2 "Descriptor Set Interface", and 14.5.4 "Offset and Stride Assignment".
-}
module FIR.Layout
( Layout(..), Components
, Poke(..), pokeArrayOff, pokeArray
, MemberOffset, memberOffset
, variableLayout
, arrayStride, memberLayoutDecorations
, roundUp, nextAligned
)
where
-- base
import Control.Monad
( foldM_ )
import Data.Int
( Int8, Int16, Int32, Int64 )
import Data.Kind
( Type )
import Data.Maybe
( isJust )
import Data.Proxy
( Proxy(..) )
import Data.Type.Equality
( (:~:)(Refl) )
import Data.Word
( Word8, Word16, Word32, Word64 )
import Foreign.Ptr
( Ptr, castPtr, plusPtr )
import qualified Foreign.Storable as Storable
( poke )
import GHC.TypeLits
( Symbol, SomeSymbol(..)
, TypeError, ErrorMessage(..)
, someSymbolVal
)
import GHC.TypeNats
( Nat, KnownNat, SomeNat(..)
, type (<=), type (<=?), type (+), type (*)
, Div
, someNatVal
)
import Numeric.Natural
( Natural )
import Unsafe.Coerce
( unsafeCoerce )
-- containers
import qualified Data.Set as Set
( empty, fromList, insert )
-- half
import Numeric.Half
( Half, fromHalf )
-- mtl
import Control.Monad.Except
( MonadError, throwError )
-- text-short
import Data.Text.Short
( ShortText )
import qualified Data.Text.Short as ShortText
( pack )
-- fir
import Data.Constraint.All
( AllDict(ConsDict)
, All(allDict)
)
import Data.Type.Known
( Known, knownValue )
import Data.Type.Map
( (:->)((:->)), Value )
import Data.Type.Nat
( NextPositivePowerOf2, RoundUp )
import Data.Type.POrd
( POrd(Max) )
import FIR.Prim.Array
( Array(..) )
import FIR.Prim.RayTracing
( AccelerationStructure(..) )
import FIR.Prim.Struct
( Struct(..)
, LocationSlot(LocationSlot)
, FieldKind(..)
, StructFieldKind(fieldKind)
)
import FIR.Prim.Types
( ScalarTy, PrimTyMap(..), SPrimTyMap(..)
, SPrimTy(..)
, primTySing
)
import qualified FIR.Prim.Types as Prim
( PrimTy )
import Math.Algebra.Class
( AdditiveGroup, Semiring )
import Math.Linear
( V, M(unM) )
import qualified SPIRV.Decoration as SPIRV
import qualified SPIRV.Layout as SPIRV
import qualified SPIRV.PrimTy as SPIRV
import SPIRV.PrimTy
( PrimTy(..) ) -- field accessors for named field puns
import qualified SPIRV.ScalarTy as SPIRV
import qualified SPIRV.Storage as SPIRV
( StorageClass )
import qualified SPIRV.Storage as Storage
--------------------------------------------------------------------------------------------
-- | Specification of alignment requirements.
data Layout
= Base -- ^ Layout used for storage buffers and push constants.
| Extended -- ^ Layout used for uniform buffers.
| Locations -- ^ Layout used for vertex binding input data.
deriving stock ( Eq, Show )
-- | Type-class for data that can be stored in memory,
-- following the specified alignment requirements.
class ( KnownNat (SizeOf lay a)
, KnownNat (Alignment lay a)
) => Poke a (lay :: Layout) where
type SizeOf lay a :: Nat
sizeOf :: Word32
sizeOf = knownValue @(SizeOf lay a)
type Alignment lay a :: Nat
alignment :: Word32
alignment = knownValue @(Alignment lay a)
poke :: Ptr a -> a -> IO ()
type Components a = ( SizeOf Locations a `Div` 4 :: Nat )
type family SumRoundedUpSizes (lay :: Layout) (as :: [fld :-> Type]) (acc :: Nat) :: Nat where
SumRoundedUpSizes _ '[] acc = acc
SumRoundedUpSizes lay ( ( _ ':-> a ) ': as ) acc = SumRoundedUpSizes lay as ( ( acc `RoundUp` Alignment lay a ) + SizeOf lay a )
-- | The offset, in bytes, of a member of a structure, in a layout.
type family MemberOffset (lay :: Layout) (k :: Symbol) (struct :: Type) :: Nat where
MemberOffset lay k ( Struct as ) = MemberOffsetFrom lay k as 0
type family MemberOffsetFrom (lay :: Layout) (k :: Symbol) (as :: [Symbol :-> Type]) (acc :: Nat) :: Nat where
MemberOffsetFrom lay k ( ( k ':-> a ) ': _ ) acc =
acc `RoundUp` Alignment lay a
MemberOffsetFrom lay k ( ( _ ':-> a ) ': as ) acc =
MemberOffsetFrom lay k as ( ( acc `RoundUp` Alignment lay a ) + SizeOf lay a )
MemberOffsetFrom _ k '[] _ =
TypeError ( Text "No structure member named " :<>: ShowType k :<>: Text "." )
-- | 'MemberOffset', as a value.
memberOffset
:: forall lay k struct. KnownNat ( MemberOffset lay k struct ) => Word32
memberOffset = knownValue @( MemberOffset lay k struct )
type family MaximumAlignment (lay :: Layout) (as :: [fld :-> Type]) :: Nat where
MaximumAlignment _ '[] = 1
MaximumAlignment lay ( ( _ ':-> a ) ': as ) = Max ( Alignment lay a ) ( MaximumAlignment lay as )
roundUp :: Integral a => a -> a -> a
roundUp n r = (n + r) - ( 1 + ( (n + r - 1) `mod` r ) )
-- | Write elements to an array, whose start is given by a pointer.
pokeArray :: forall a lay. ( Poke a lay, Poke (Array 1 a) lay ) => Ptr a -> [a] -> IO ()
pokeArray ptr = pokeArrayOff @a @lay ptr 0
-- | Write elements to an array, starting from the given index offset.
pokeArrayOff :: forall a lay. ( Poke a lay, Poke (Array 1 a) lay ) => Ptr a -> Int -> [a] -> IO ()
pokeArrayOff ptr off vals0
| off < 0 = pure ()
| otherwise = go ( ptr `plusPtr` ( sz * off ) ) vals0
where
sz :: Int
sz = fromIntegral $ nextAligned (sizeOf @a @lay) (alignment @(Array 1 a) @lay)
go :: Ptr a -> [a] -> IO ()
go _ [] = pure ()
go ptr' (val:vals) = do
poke @a @lay ptr' val
go ( ptr' `plusPtr` sz ) vals
nextAligned :: Integral a => a -> a -> a
nextAligned size ali
= case size `mod` ali of
0 -> size
x -> size + ali - x
instance Poke Word8 lay where
type SizeOf lay Word8 = 4
type Alignment lay Word8 = 4
poke ptr = Storable.poke (castPtr @_ @Word32 ptr) . fromIntegral
instance Poke Word16 lay where
type SizeOf lay Word16 = 4
type Alignment lay Word16 = 4
poke ptr = Storable.poke (castPtr @_ @Word32 ptr) . fromIntegral
instance Poke Word32 lay where
type SizeOf lay Word32 = 4
type Alignment lay Word32 = 4
poke = Storable.poke
instance Poke Word64 lay where
type SizeOf lay Word64 = 8
type Alignment lay Word64 = 8
poke = Storable.poke
instance Poke Int8 lay where
type SizeOf lay Int8 = 4
type Alignment lay Int8 = 4
poke ptr = Storable.poke (castPtr @_ @Int32 ptr) . fromIntegral
instance Poke Int16 lay where
type SizeOf lay Int16 = 4
type Alignment lay Int16 = 4
poke ptr = Storable.poke (castPtr @_ @Int32 ptr) . fromIntegral
instance Poke Int32 lay where
type SizeOf lay Int32 = 4
type Alignment lay Int32 = 4
poke = Storable.poke
instance Poke Int64 lay where
type SizeOf lay Int64 = 4
type Alignment lay Int64 = 4
poke = Storable.poke
instance Poke Half lay where
type SizeOf lay Half = 4
type Alignment lay Half = 4
poke ptr = Storable.poke (castPtr @_ @Float ptr) . fromHalf
instance Poke Float lay where
type SizeOf lay Float = 4
type Alignment lay Float = 4
poke = Storable.poke
instance Poke Double lay where
type SizeOf lay Double = 8
type Alignment lay Double = 8
poke = Storable.poke
instance TypeError ( Text "Cannot store Boolean types." ) => Poke Bool lay where
type SizeOf lay Bool = 4
type Alignment lay Bool = 4
poke = error "unreachable"
-- Acceleration structures are just represented as Word64 address
instance Poke AccelerationStructure lay where
type SizeOf lay AccelerationStructure = 8
type Alignment lay AccelerationStructure = 8
poke ptr ( AccelerationStructureFromWord64 word64 ) = poke @Word64 @lay ( castPtr ptr ) word64
instance ( Poke a lay
, ScalarTy a
, KnownNat n, 1 <= n
, KnownNat (SizeOf lay (V n a))
, KnownNat (Alignment lay (V n a))
) => Poke (V n a) lay where
type SizeOf lay (V n a) = n * SizeOf lay a
type Alignment lay (V n a)
= SizeOf lay a * NextPositivePowerOf2 n
poke ptr v
= foldM_
( \accPtr elt -> poke @a @lay accPtr elt *> pure (accPtr `plusPtr` off) )
( castPtr ptr )
v
where
off :: Int
off = fromIntegral ( sizeOf @a @lay )
-- A matrix is laid out like an array of its columns.
-- As the size of a column is no larger than its alignment,
-- the stride between columns is the alignment of the matrix.
instance (Poke (V m a) Base, ScalarTy a, KnownNat n, KnownNat m, 1 <= n, 1 <= m)
=> Poke (M m n a) Base where
type SizeOf Base (M m n a)
= n * Alignment Base (M m n a)
type Alignment Base (M m n a) = Alignment Base (V m a)
poke ptr mat
= foldM_
( \accPtr col -> poke @(V m a) @Base accPtr col *> pure (accPtr `plusPtr` colStride) )
( castPtr ptr )
( unM mat )
where
colStride = fromIntegral $ alignment @(M m n a) @Base
instance ( Poke (V m a) Extended, ScalarTy a, KnownNat n, KnownNat m, 1 <= n, 1 <= m
, KnownNat (SizeOf Extended (M m n a))
, KnownNat (Alignment Extended (M m n a))
)
=> Poke (M m n a) Extended where
type SizeOf Extended (M m n a)
= n * Alignment Extended (M m n a)
type Alignment Extended (M m n a) = Alignment Extended (V m a) `RoundUp` 16
poke ptr mat
= foldM_
( \accPtr col -> poke @(V m a) @Extended accPtr col *> pure (accPtr `plusPtr` colStride) )
( castPtr ptr )
( unM mat )
where
colStride = fromIntegral $ alignment @(M m n a) @Extended
instance ( Poke (V m a) Locations
, ScalarTy a
, KnownNat n, KnownNat m, 1 <= n, 1 <= m
, KnownNat (SizeOf Locations (M m n a))
)
=> Poke (M m n a) Locations where
type SizeOf Locations (M m n a)
= n * ( SizeOf Locations (V m a) `RoundUp` 16 )
type Alignment Locations (M m n a) = 16
poke ptr mat
= foldM_
( \accPtr col -> poke @(V m a) @Locations accPtr col *> pure (accPtr `plusPtr` colSize) )
( castPtr ptr )
( unM mat )
where
colSize :: Int
colSize = fromIntegral ( sizeOf @(V m a) @Locations ) `roundUp` 16
instance ( Poke a Base, Prim.PrimTy a, KnownNat n, 1 <= n
, KnownNat ( SizeOf Base ( Array n a ) )
)
=> Poke (Array n a) Base where
type SizeOf Base (Array n a) = n * ( SizeOf Base a `RoundUp` Alignment Base (Array n a) )
type Alignment Base (Array n a) = Alignment Base a
poke ptr (MkArray arr)
= foldM_
( \accPtr elt -> poke @a @Base accPtr elt *> pure (accPtr `plusPtr` off) )
( castPtr ptr )
arr
where
off :: Int
off = fromIntegral $ nextAligned (sizeOf @a @Base) (alignment @(Array n a) @Base)
instance ( Poke a Extended
, Prim.PrimTy a
, KnownNat n, 1 <= n
, KnownNat ( SizeOf Extended ( Array n a ) )
, KnownNat ( Alignment Extended ( Array n a ) )
)
=> Poke (Array n a) Extended where
type SizeOf Extended (Array n a)
= n * ( SizeOf Extended a `RoundUp` Alignment Extended (Array n a) )
type Alignment Extended (Array n a)
= Alignment Extended a `RoundUp` 16
poke ptr (MkArray arr)
= foldM_
( \accPtr elt -> poke @a @Extended accPtr elt *> pure (accPtr `plusPtr` off) )
( castPtr ptr )
arr
where
off :: Int
off = fromIntegral $ nextAligned (sizeOf @a @Extended) (alignment @(Array n a) @Extended)
instance (Poke a Locations, Prim.PrimTy a, KnownNat n, 1 <= n)
=> Poke (Array n a) Locations where
type SizeOf Locations (Array n a)
= n * SizeOf Locations a
type Alignment Locations (Array n a)
= Alignment Locations a
poke ptr (MkArray arr)
= foldM_
( \accPtr elt -> poke @a @Locations accPtr elt *> pure (accPtr `plusPtr` off) )
( castPtr ptr )
arr
where
off :: Int
off = fromIntegral $ nextAligned (sizeOf @a @Locations) (alignment @(Array n a) @Locations `roundUp` 16)
class ( Poke (Value x) lay ) => Pokeable (lay :: Layout) (x :: fld :-> Type) where
instance ( Poke (Value x) lay ) => Pokeable (lay :: Layout) (x :: fld :-> Type) where
structPoke :: forall (fld :: Type) (lay :: Layout) (as :: [fld :-> Type]).
( PrimTyMap as, All (Pokeable lay) as )
=> Word32 -> Ptr (Struct as) -> Struct as -> IO ()
structPoke currOff ptr struct = case primTyMapSing @_ @as of
SNil -> pure ()
scons@SCons -> case scons of
( _ :: SPrimTyMap ((k ':-> b) ': bs) ) ->
case (struct, allDict @(Pokeable lay) @as) of
( b :& bs, ConsDict ) -> do
let
off :: Word32
off = nextAligned currOff ( alignment @b @lay )
poke @b @lay ( castPtr ptr `plusPtr` ( fromIntegral off ) ) b
structPoke @fld @lay @bs ( off + sizeOf @b @lay ) ( castPtr ptr ) bs
instance ( PrimTyMap as
, All (Pokeable Base) as
, KnownNat (SizeOf Base (Struct as))
, KnownNat (Alignment Base (Struct as))
) => Poke (Struct as) Base where
type SizeOf Base (Struct as)
= SumRoundedUpSizes Base as 0 `RoundUp` Alignment Base (Struct as)
type Alignment Base (Struct as)
= MaximumAlignment Base as
poke = structPoke @_ @Base @as 0
instance ( PrimTyMap as
, All (Pokeable Extended) as
, KnownNat (SizeOf Extended (Struct as))
, KnownNat (Alignment Extended (Struct as))
) => Poke (Struct as) Extended where
type SizeOf Extended (Struct as)
= SumRoundedUpSizes Extended as 0 `RoundUp` Alignment Extended (Struct as)
type Alignment Extended (Struct as)
= MaximumAlignment Extended as `RoundUp` 16
poke = structPoke @_ @Extended @as 0
instance
( TypeError
( Text "Structure contains insufficient layout information."
:$$: Text ""
:$$: Text "To use the 'Locations' layout, location/component information needs to be supplied, such as:"
:$$: Text " Struct '[ Slot 0 0 ':-> V 2 Float, Slot 0 2 ':-> V 2 Float ],"
:$$: Text "where 'Slot l c' refers to the interface slot in location 'l' with component 'c',"
:$$: Text "with 0 <= l < maxVertexInputAttributes (at least 16), 0 <= c < 4."
)
)
=> Poke (Struct (as :: [Symbol :-> Type])) Locations
where
type SizeOf Locations (Struct (as :: [Symbol :-> Type]))
= 16
type Alignment Locations (Struct (as :: [Symbol :-> Type]))
= 16
poke = error "unreachable"
-- TODO: refactor this using 'Data.Constraint.All'
data PokeSlotsDict (as :: [LocationSlot Nat :-> Type]) :: Type where
NoSlotsPoke :: PokeSlotsDict '[]
PokeSlotDict :: ( Poke a Locations
, Known (LocationSlot Nat) slot
, PokeSlots as
, KnownNat (SizeOf Locations (Struct as))
, KnownNat (Alignment Locations (Struct as))
)
=> PokeSlotsDict ( (slot ':-> a) ': as )
class PokeSlots (as :: [LocationSlot Nat :-> Type]) where
pokeSlotsDict :: PokeSlotsDict as
instance PokeSlots '[] where
pokeSlotsDict = NoSlotsPoke
instance ( Poke a Locations
, Known (LocationSlot Nat) loc
, PokeSlots as
, KnownNat (SizeOf Locations (Struct as))
, KnownNat (Alignment Locations (Struct as))
)
=> PokeSlots ( (loc ':-> a) ': as )
where
pokeSlotsDict = PokeSlotDict
instance forall (as :: [LocationSlot Nat :-> Type])
. ( PrimTyMap as
, PokeSlots as
, KnownNat (SizeOf Locations (Struct as))
, KnownNat (Alignment Locations (Struct as))
) => Poke (Struct as) Locations where
type SizeOf Locations (Struct as)
= SizeOfLocations as
type Alignment Locations (Struct as)
= 16
poke ptr struct = case primTyMapSing @_ @as of
SNil -> pure ()
scons@SCons -> case scons of
( _ :: SPrimTyMap ((loc ':-> b) ': bs) ) ->
case (struct, pokeSlotsDict @as) of
( b :& bs, PokeSlotDict ) ->
let LocationSlot l c = knownValue @loc
off = fromIntegral $ 16 * l + 4 * c
in do
-- if we are poking another struct with location information,
-- we want to keep the pointer as is as opposed to adding an offset to it
case primTySing @b of
sStruct@SStruct
| ( _ :: SPrimTy (Struct (cs :: [fld :-> Type]) ) ) <- sStruct
, LocationField <- fieldKind @fld
-> poke @b @Locations (castPtr ptr ) b
_ -> poke @b @Locations (castPtr ptr `plusPtr` off) b
poke @(Struct bs) @Locations (castPtr ptr) bs
type family SizeOfLocations (as :: [LocationSlot Nat :-> Type]) :: Nat where
SizeOfLocations '[] = 0
SizeOfLocations ( ( 'LocationSlot l _ ':-> _ ) ': as )
= Max ( 16 * (l+1) ) ( SizeOfLocations as )
--------------------------------------------------------------------------------------------
-- Layout of SPIR-V types in memory.
-- | The layout rules imposed by a storage class, if any.
storageLayoutRules :: SPIRV.StorageClass -> Maybe SPIRV.LayoutRules
storageLayoutRules Storage.Uniform = Just SPIRV.ExtendedAlignment
storageLayoutRules Storage.StorageBuffer = Just SPIRV.BaseAlignment
storageLayoutRules Storage.PushConstant = Just SPIRV.BaseAlignment
storageLayoutRules (Storage.RayStorage Storage.ShaderRecordBuffer) = Just SPIRV.BaseAlignment
storageLayoutRules _ = Nothing
-- | The layout of a global variable of the given type, in the given storage class.
--
-- In storage classes with explicit layout, as well as for inputs and outputs,
-- structures and arrays of structures are interface blocks.
variableLayout :: SPIRV.BlockKind -> SPIRV.StorageClass -> SPIRV.PrimTy -> SPIRV.Layout
variableLayout kind storageClass ty
= SPIRV.Layout
{ SPIRV.layoutRules = rules
, SPIRV.blockKind
= if isInterface && isBlockShaped ty
then Just kind
else Nothing
}
where
rules :: Maybe SPIRV.LayoutRules
rules = storageLayoutRules storageClass
isInterface :: Bool
isInterface = isJust rules || storageClass `elem` [ Storage.Input, Storage.Output ]
isBlockShaped :: SPIRV.PrimTy -> Bool
isBlockShaped ( SPIRV.Struct {} ) = True
isBlockShaped ( SPIRV.Array _ ( SPIRV.Struct {} ) ) = True
isBlockShaped _ = False
-- | The stride of an array with the given element type.
arrayStride :: MonadError ShortText m => SPIRV.LayoutRules -> SPIRV.PrimTy -> m Word32
arrayStride rules eltTy = uncurry nextAligned <$> primTySizeAndAli rules ( SPIRV.Array 1 eltTy )
-- | The layout decorations of the members of a structure with the given member types.
memberLayoutDecorations
:: forall m
. MonadError ShortText m
=> SPIRV.LayoutRules -> [ SPIRV.PrimTy ] -> m [ SPIRV.Decorations ]
memberLayoutDecorations rules = go 0
where
go :: Word32 -> [ SPIRV.PrimTy ] -> m [ SPIRV.Decorations ]
go _ [] = pure []
go offset ( ty : tys ) = do
( size, ali ) <- primTySizeAndAli rules ( sizedTy ty )
matDecs <- matrixDecorations ( innermostElement ty )
let
memberOffset :: Word32
memberOffset = nextAligned offset ali
( Set.insert ( SPIRV.Offset memberOffset ) matDecs : ) <$> go ( memberOffset + size ) tys
-- A run-time array has no size; it is laid out like an array with a single element.
sizedTy :: SPIRV.PrimTy -> SPIRV.PrimTy
sizedTy ( SPIRV.RuntimeArray { eltTy } ) = SPIRV.Array 1 eltTy
sizedTy ty = ty
innermostElement :: SPIRV.PrimTy -> SPIRV.PrimTy
innermostElement ( SPIRV.Array { eltTy } ) = innermostElement eltTy
innermostElement ( SPIRV.RuntimeArray { eltTy } ) = innermostElement eltTy
innermostElement ty = ty
-- The layout of a matrix is specified on the structure member that contains it;
-- a matrix is laid out like an array of its columns.
matrixDecorations :: SPIRV.PrimTy -> m SPIRV.Decorations
matrixDecorations ( SPIRV.Matrix { rows, entryTy } ) = do
colStride <- arrayStride rules ( SPIRV.Vector rows ( SPIRV.Scalar entryTy ) )
pure ( Set.fromList [ SPIRV.MatrixStride colStride, SPIRV.ColMajor ] )
matrixDecorations _ = pure Set.empty
primTySizeAndAli :: MonadError ShortText m => SPIRV.LayoutRules -> SPIRV.PrimTy -> m ( Word32, Word32 )
primTySizeAndAli SPIRV.BaseAlignment ty
| Just ( LayoutablePrimTy ( _ :: Proxy ty ) ) <- layoutable @Base ty
= pure $ ( sizeOf @ty @Base, alignment @ty @Base )
primTySizeAndAli SPIRV.ExtendedAlignment ty
| Just ( LayoutablePrimTy ( _ :: Proxy ty ) ) <- layoutable @Extended ty
= pure $ ( sizeOf @ty @Extended, alignment @ty @Extended )
primTySizeAndAli rules ty = throwError
( "'primTySizeAndAli': cannot compute " <> ShortText.pack (show rules) <> " size & alignment for " <> ShortText.pack (show ty) )
--------------------------------------------------------------------------------------------
-- Some functions for promoting SPIR-V primitive types that can be laid-out.
-- This allows us to re-use the type families that compute layout information
-- at the value level.
data SLayout (lay :: Layout) where
SBase :: SLayout Base
SExtended :: SLayout Extended
SLocations :: SLayout Locations
class KnownLayout (lay :: Layout) where
layoutSing :: SLayout lay
instance KnownLayout Base where
layoutSing = SBase
instance KnownLayout Extended where
layoutSing = SExtended
instance KnownLayout Locations where
layoutSing = SLocations
data LayoutablePrimTy ( lay :: Layout ) where
LayoutablePrimTy :: ( Prim.PrimTy a, Poke a lay ) => Proxy a -> LayoutablePrimTy lay
data LayoutableScalarTy ( lay :: Layout ) where
LayoutableScalarTy :: ( ScalarTy a, Poke a lay ) => Proxy a -> LayoutableScalarTy lay
data LayoutableMatScalarTy ( lay :: Layout ) where
LayoutableMatScalarTy :: ( ScalarTy a, AdditiveGroup a, Semiring a, Poke a lay ) => Proxy a -> LayoutableMatScalarTy lay
data LayoutableMembers (lay :: Layout) where
LayoutableMembers
:: ( PrimTyMap as, All (Pokeable lay) as
, KnownNat (SizeOf lay (Struct as))
, KnownNat (Alignment lay (Struct as))
)
=> Proxy as -> LayoutableMembers lay
layoutable :: forall ( lay :: Layout ). KnownLayout lay => SPIRV.PrimTy -> Maybe ( LayoutablePrimTy lay )
layoutable ( SPIRV.Scalar s )
| Just ( LayoutableScalarTy ( _ :: Proxy a ) ) <- layoutableScalar @lay s
= Just ( LayoutablePrimTy ( Proxy :: Proxy a ) )
layoutable ( SPIRV.Vector n ( SPIRV.Scalar s ) )
| Just ( LayoutableScalarTy ( _ :: Proxy a ) ) <- layoutableScalar @lay s
, True <- 1 <= n
, SomeNat ( _ :: Proxy n ) <- someNatVal ( fromIntegral n )
, Refl <- ( unsafeCoerce Refl :: ( 1 <=? n ) :~: True )
, let
npp2 :: Natural
npp2
| n == 1 = 1
| otherwise = 2 ^ ( floor ( logBase 2 ( realToFrac ( n - 1 ) :: Float ) ) + 1 :: Int )
, SomeNat ( _ :: Proxy npp2 ) <- someNatVal npp2
, Refl <- ( unsafeCoerce Refl :: NextPositivePowerOf2 n :~: npp2 )
= Just $ LayoutablePrimTy ( Proxy :: Proxy ( V n a ) )
layoutable ( SPIRV.Matrix m n a )
| Just ( LayoutableMatScalarTy ( _ :: Proxy a ) ) <- layoutableMatScalar @lay a
, SomeNat ( _ :: Proxy m ) <- someNatVal ( fromIntegral m )
, SomeNat ( _ :: Proxy n ) <- someNatVal ( fromIntegral n )
, True <- 1 <= n
, Refl <- ( unsafeCoerce Refl :: ( 1 <=? n ) :~: True )
, True <- 1 <= m
, Refl <- ( unsafeCoerce Refl :: ( 1 <=? m ) :~: True )
, let
npp2 :: Natural
npp2
| m == 1 = 1
| otherwise = 2 ^ ( floor ( logBase 2 ( realToFrac ( m - 1 ) :: Float ) ) + 1 :: Int )
, SomeNat ( _ :: Proxy npp2 ) <- someNatVal npp2
, Refl <- ( unsafeCoerce Refl :: NextPositivePowerOf2 m :~: npp2 )
= case layoutSing @lay of
SBase -> Just $ LayoutablePrimTy ( Proxy :: Proxy ( M m n a ) )
SExtended
| let
colStride :: Word32
colStride = alignment @(V m a) @Extended `roundUp` 16
, SomeNat ( _ :: Proxy ali ) <- someNatVal ( fromIntegral colStride )
, Refl <- ( unsafeCoerce Refl :: ali :~: Alignment Extended ( M m n a ) )
, SomeNat ( _ :: Proxy sz ) <- someNatVal ( fromIntegral $ n * colStride )
, Refl <- ( unsafeCoerce Refl :: sz :~: SizeOf Extended ( M m n a ) )
-> Just $ LayoutablePrimTy ( Proxy :: Proxy ( M m n a ) )
_ -> Nothing
layoutable ( SPIRV.Array { size, eltTy } )
| Just ( LayoutablePrimTy ( _ :: Proxy a ) ) <- layoutable @lay eltTy
, True <- 1 <= size
, SomeNat ( _ :: Proxy lg ) <- someNatVal ( fromIntegral size )
, Refl <- ( unsafeCoerce Refl :: ( 1 <=? lg ) :~: True )
= case layoutSing @lay of
SBase
| SomeNat ( _ :: Proxy sz ) <- someNatVal ( fromIntegral $ size * ( nextAligned ( sizeOf @a @Base ) ( alignment @a @Base ) ) )
, Refl <- ( unsafeCoerce Refl :: sz :~: SizeOf Base ( Array lg a ) )
-> Just $ LayoutablePrimTy ( Proxy :: Proxy ( Array lg a ) )
SExtended
| SomeNat ( _ :: Proxy sz ) <- someNatVal ( fromIntegral $ size * ( nextAligned ( sizeOf @a @Extended ) ( alignment @a @Extended `roundUp` 16 ) ) )
, Refl <- ( unsafeCoerce Refl :: sz :~: SizeOf Extended ( Array lg a ) )
, SomeNat ( _ :: Proxy ali ) <- someNatVal ( fromIntegral $ alignment @a @Extended `roundUp` 16 )
, Refl <- ( unsafeCoerce Refl :: ali :~: (Alignment Extended a `RoundUp` 16) )
-> Just $ LayoutablePrimTy ( Proxy :: Proxy ( Array lg a ) )
_ -> Nothing
layoutable ( SPIRV.Struct { eltTys } ) = layoutableStruct @lay eltTys
layoutable _ = Nothing
layoutableScalar :: forall ( lay :: Layout ). SPIRV.ScalarTy -> Maybe ( LayoutableScalarTy lay )
layoutableScalar ( SPIRV.Integer SPIRV.Unsigned SPIRV.W8 ) = Just $ LayoutableScalarTy ( Proxy :: Proxy Word8 )
layoutableScalar ( SPIRV.Integer SPIRV.Unsigned SPIRV.W16 ) = Just $ LayoutableScalarTy ( Proxy :: Proxy Word16 )
layoutableScalar ( SPIRV.Integer SPIRV.Unsigned SPIRV.W32 ) = Just $ LayoutableScalarTy ( Proxy :: Proxy Word32 )
layoutableScalar ( SPIRV.Integer SPIRV.Unsigned SPIRV.W64 ) = Just $ LayoutableScalarTy ( Proxy :: Proxy Word64 )
layoutableScalar ( SPIRV.Integer SPIRV.Signed SPIRV.W8 ) = Just $ LayoutableScalarTy ( Proxy :: Proxy Int8 )
layoutableScalar ( SPIRV.Integer SPIRV.Signed SPIRV.W16 ) = Just $ LayoutableScalarTy ( Proxy :: Proxy Int16 )
layoutableScalar ( SPIRV.Integer SPIRV.Signed SPIRV.W32 ) = Just $ LayoutableScalarTy ( Proxy :: Proxy Int32 )
layoutableScalar ( SPIRV.Integer SPIRV.Signed SPIRV.W64 ) = Just $ LayoutableScalarTy ( Proxy :: Proxy Int64 )
layoutableScalar ( SPIRV.Floating SPIRV.W16 ) = Just $ LayoutableScalarTy ( Proxy :: Proxy Half )
layoutableScalar ( SPIRV.Floating SPIRV.W32 ) = Just $ LayoutableScalarTy ( Proxy :: Proxy Float )
layoutableScalar ( SPIRV.Floating SPIRV.W64 ) = Just $ LayoutableScalarTy ( Proxy :: Proxy Double )
layoutableScalar _ = Nothing
layoutableMatScalar :: forall ( lay :: Layout ). SPIRV.ScalarTy -> Maybe ( LayoutableMatScalarTy lay )
layoutableMatScalar ( SPIRV.Integer SPIRV.Signed SPIRV.W8 ) = Just $ LayoutableMatScalarTy ( Proxy :: Proxy Int8 )
layoutableMatScalar ( SPIRV.Integer SPIRV.Signed SPIRV.W16 ) = Just $ LayoutableMatScalarTy ( Proxy :: Proxy Int16 )
layoutableMatScalar ( SPIRV.Integer SPIRV.Signed SPIRV.W32 ) = Just $ LayoutableMatScalarTy ( Proxy :: Proxy Int32 )
layoutableMatScalar ( SPIRV.Integer SPIRV.Signed SPIRV.W64 ) = Just $ LayoutableMatScalarTy ( Proxy :: Proxy Int64 )
layoutableMatScalar ( SPIRV.Floating SPIRV.W16 ) = Just $ LayoutableMatScalarTy ( Proxy :: Proxy Half )
layoutableMatScalar ( SPIRV.Floating SPIRV.W32 ) = Just $ LayoutableMatScalarTy ( Proxy :: Proxy Float )
layoutableMatScalar ( SPIRV.Floating SPIRV.W64 ) = Just $ LayoutableMatScalarTy ( Proxy :: Proxy Double )
layoutableMatScalar _ = Nothing
layoutableStruct
:: forall ( lay :: Layout ) mbName
. KnownLayout lay
=> [ ( mbName, SPIRV.PrimTy ) ]
-> Maybe ( LayoutablePrimTy lay )
layoutableStruct elts = makeLayoutableStruct =<< layoutableStructMembers 0 elts
where
layoutableStructMembers :: Int -> [ ( mbName, SPIRV.PrimTy ) ] -> Maybe [ ( SomeSymbol, LayoutablePrimTy lay ) ]
layoutableStructMembers _ [] = Just []
layoutableStructMembers i ( (_,ty) : tys )
| Just layoutableTy <- layoutable @lay ty
, Just layoutableTys <- layoutableStructMembers (i+1) tys
, nm <- someSymbolVal ( show i )
= Just ( ( nm, layoutableTy ) : layoutableTys )
layoutableStructMembers _ _ = Nothing
makeLayoutableStruct :: [ ( SomeSymbol, LayoutablePrimTy lay ) ] -> Maybe ( LayoutablePrimTy lay )
makeLayoutableStruct members
| Just ( LayoutableMembers ( _ :: Proxy as ) ) <- makeLayoutableMembers members
= case layoutSing @lay of
SBase -> Just ( LayoutablePrimTy ( Proxy :: Proxy ( Struct as ) ) )
SExtended -> Just ( LayoutablePrimTy ( Proxy :: Proxy ( Struct as ) ) )
_ -> Nothing
makeLayoutableStruct _ = Nothing
makeLayoutableMembers :: [ ( SomeSymbol, LayoutablePrimTy lay ) ] -> Maybe ( LayoutableMembers lay )
makeLayoutableMembers [] = case layoutSing @lay of
SBase -> Just ( LayoutableMembers ( Proxy :: Proxy ('[] :: [ Symbol :-> Type ] ) ) )
SExtended -> Just ( LayoutableMembers ( Proxy :: Proxy ('[] :: [ Symbol :-> Type ] ) ) )
_ -> Nothing
makeLayoutableMembers members@( ( SomeSymbol ( _ :: Proxy nm ), LayoutablePrimTy ( _ :: Proxy a ) ) : as )
| Just ( LayoutableMembers ( _ :: Proxy ( as :: [ fld :-> Type ] ) ) ) <- makeLayoutableMembers as
, NamedField <- fieldKind @fld
, ( structSize, structAli ) <- structSizeAndAlignment @lay members
, SomeNat ( _ :: Proxy structSize ) <- someNatVal ( fromIntegral structSize )
, Refl <- ( unsafeCoerce Refl :: structSize :~: SizeOf lay ( Struct ( ( nm ':-> a ) ': as ) ) )
, SomeNat ( _ :: Proxy structAli ) <- someNatVal ( fromIntegral structAli )
, Refl <- ( unsafeCoerce Refl :: structAli :~: Alignment lay ( Struct ( ( nm ':-> a ) ': as ) ) )
= case layoutSing @lay of
SBase -> Just ( LayoutableMembers ( Proxy :: Proxy ( ( nm ':-> a ) ': as ) ) )
SExtended -> Just ( LayoutableMembers ( Proxy :: Proxy ( ( nm ':-> a ) ': as ) ) )
_ -> Nothing
makeLayoutableMembers _ = Nothing
-- | Size and alignment of a structure with the given members.
--
-- The size of a structure is rounded up to its alignment.
structSizeAndAlignment
:: forall ( lay :: Layout )
. KnownLayout lay
=> [ ( SomeSymbol, LayoutablePrimTy lay ) ] -> ( Word32, Word32 )
structSizeAndAlignment members = ( nextAligned ( foldl' memberEnd 0 members ) structAli, structAli )
where
structAli :: Word32
structAli = case layoutSing @lay of
SExtended -> maxAli `roundUp` 16
_ -> maxAli
maxAli :: Word32
maxAli = maximum ( 1 : map memberAlignment members )
memberAlignment :: ( SomeSymbol, LayoutablePrimTy lay ) -> Word32
memberAlignment ( _, LayoutablePrimTy ( _ :: Proxy b ) ) = alignment @b @lay
memberEnd :: Word32 -> ( SomeSymbol, LayoutablePrimTy lay ) -> Word32
memberEnd offset ( _, LayoutablePrimTy ( _ :: Proxy b ) ) = nextAligned offset ( alignment @b @lay ) + sizeOf @b @lay