fir-0.1.0.0: src/CodeGen/Binary.hs
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PackageImports #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-|
Module: CodeGen.Binary
Emitting SPIR-V binary data.
-}
module CodeGen.Binary
( putModule
, instruction
, structuralInstruction
, deferredInstructions
)
where
-- base
import Control.Monad
( when )
import Data.List
( sortOn)
import Data.Foldable
( traverse_)
import Data.Maybe
( isJust )
import Data.Word
( Word32 )
import qualified Data.Bits as Bits
-- binary
import qualified Data.Binary as Binary
import qualified Data.Binary.Put as Binary
-- containers
import Data.Map.Strict
( Map )
import qualified Data.Map.Strict as Map
import Data.Set
( Set )
-- lens
import Control.Lens
( view )
-- mtl
import Control.Monad.Except
( ExceptT, throwError )
-- text-short
import Data.Text.Short
( ShortText )
import qualified Data.Text.Short as ShortText
( pack )
-- transformers
import Control.Monad.Trans.Class
( lift )
-- fir
import CodeGen.IDs
( extInstID )
import CodeGen.Instruction
( Args(..), toArgs
, ID(..), TyID(..), pattern MkTyID
, Variable, variableID, variableDeclaration
, Instruction(..)
)
import CodeGen.Monad
( CGMonad, note, liftPut )
import CodeGen.State
( CGContext(..), CGState(moduleState), ModuleState(..)
, _emittingCode
, currentBlockID, leaveBlock
, currentFunctionContext
)
import Data.Binary.Class.Put
( Put(put, wordCount) )
import Data.Containers.Traversals
( traverseWithKey_, traverseSet_ )
import FIR.ProgramState
( FunctionContext(TopLevel) )
import qualified SPIRV.Capability as SPIRV
import qualified SPIRV.Decoration as SPIRV
import qualified SPIRV.ExecutionMode as SPIRV
import qualified SPIRV.Extension as SPIRV
( Extension, ExtInst, extInstName )
import qualified SPIRV.Operation as SPIRV.Op
import qualified SPIRV.PrimTy as SPIRV
import qualified SPIRV.Stage as SPIRV
import qualified SPIRV.Version as SPIRV
----------------------------------------------------------------------------
-- emitting a SPIR-V module
-- some instruction need to be floated to the top
putModule :: CGContext -> CGState -> ExceptT ShortText Binary.PutM ()
putModule
CGContext { .. }
cgState
| emittingCode = do
lift $ do putHeader spirvVersion ( idNumber currentID )
putCapabilities neededCapabilities
putExtensions neededExtensions
putExtendedInstructions knownExtInsts
putMemoryModel backend
let knownBindingIDs = map ( \ (name, (ident, _)) -> (name, ident) ) ( Map.toList knownBindings )
usedGlobalIDs = map ( fmap variableID ) ( Map.toList usedGlobals )
builtinIDs = map ( fmap variableID ) ( Map.elems builtinVariables )
putEntryPoints entryPoints interfaces
putExecutionModes entryPoints userEntryPoints
lift $ do
putKnownStringLits knownStringLits
putBindingAnnotations knownBindingIDs
putBindingAnnotations usedGlobalIDs
putBindingAnnotations builtinIDs
putNames names
putDecorations decorations
putMemberDecorations memberDecorations
-- Type and constant declarations need to be interleaved.
-- For instance, an array type needs to have
-- its length (a constant) defined earlier.
putTypesAndConstants knownTypes knownConstants
putUndefineds knownUndefineds
putGlobals
( Map.elems usedGlobals
++ map snd ( Map.elems builtinVariables )
)
| otherwise = pure ()
where
ModuleState { .. } = moduleState cgState
----------------------------------------------------------------------------
-- individual binary instructions
-- | Emit an instruction in the current block (if there is one).
--
-- A block termination instruction ends the current block.
instruction :: Instruction -> CGMonad ()
instruction inst = do
mbBlock <- currentBlockID
case mbBlock of
Nothing -> do
context <- currentFunctionContext
case context of
TopLevel ->
throwError ( "codeGen: instruction outside of a function: " <> ShortText.pack ( show ( operation inst ) ) )
_ -> pure ()
Just _ -> do
emitInstruction inst
when ( SPIRV.Op.isBlockTerminator ( operation inst ) )
leaveBlock
-- | Emit an instruction which doesn't belong to a block, such as a label or @OpFunction@.
structuralInstruction :: Instruction -> CGMonad ()
structuralInstruction = emitInstruction
-- | Emit instructions in the current block, without inspecting them during
-- code generation (for instructions that are only known after generating
-- subsequent code).
--
-- Doesn't support extended instructions.
deferredInstructions :: [ Instruction ] -> CGMonad ()
deferredInstructions insts = do
emitting <- view _emittingCode
mbBlock <- currentBlockID
when ( emitting && isJust mbBlock )
( liftPut ( traverse_ putInstruction insts ) )
emitInstruction :: Instruction -> CGMonad ()
emitInstruction inst = do
emitting <- view _emittingCode
when emitting
case operation inst of
SPIRV.Op.ExtCode ext _ -> do
extID <- extInstID ext
liftPut $ putExtendedInstruction extID inst
_ -> liftPut $ putInstruction inst
putInstruction :: Instruction -> Binary.Put
putInstruction Instruction { operation = op, resTy = opResTy, resID = opResID, args = opArgs }
= case op of
SPIRV.Op.Code opCode ->
let n :: Word32
n = 1 -- OpCode and word count (first byte)
+ maybe 0 (const 1) opResTy -- result type (if present)
+ maybe 0 (const 1) opResID -- ID (if present)
+ wordCount opArgs
in do put @Word32 ( Bits.shift n 16 + fromIntegral opCode)
traverse_ put opResTy
traverse_ put opResID
put opArgs
_ -> error "putInstruction: unexpected extended instruction"
putExtendedInstruction :: ID -> Instruction -> Binary.Put
putExtendedInstruction extID Instruction { operation = op, resTy = opResTy, resID = opResID, args = opArgs }
= case op of
SPIRV.Op.ExtCode _ extOpCode ->
putInstruction
Instruction
{ operation = SPIRV.Op.ExtInst
, resTy = opResTy
, resID = opResID
, args = Arg extID
$ Arg extOpCode
opArgs
}
_ -> error "putExtendedInstruction: unexpected non-extended instruction"
putHeader :: SPIRV.Version -> Word32 -> Binary.Put
putHeader ver bound
= do
put SPIRV.magicNo
put ver
put libraryMagicNo
put bound
put ( 0 :: Word32 )
where
libraryMagicNo :: Word32
libraryMagicNo = 0x21524946 -- FIR!
putCapabilities :: Set SPIRV.Capability -> Binary.Put
putCapabilities = traverseSet_ putCap
where
putCap :: SPIRV.Capability -> Binary.Put
putCap cap
= putInstruction
Instruction
{ operation = SPIRV.Op.Capability
, resTy = Nothing
, resID = Nothing
, args = Arg cap
EndArgs
}
putExtensions :: Set SPIRV.Extension -> Binary.Put
putExtensions = traverseSet_ putExt
where
putExt :: SPIRV.Extension -> Binary.Put
putExt ext
= putInstruction
Instruction
{ operation = SPIRV.Op.Extension
, resTy = Nothing
, resID = Nothing
, args = Arg ext
EndArgs
}
putExtendedInstructions :: Map SPIRV.ExtInst ID -> Binary.Put
putExtendedInstructions
= traverseWithKey_ \ extInst extID ->
putInstruction
Instruction
{ operation = SPIRV.Op.ExtInstImport
, resTy = Nothing
, resID = Just extID
, args = Arg ( SPIRV.extInstName extInst ) EndArgs
}
putMemoryModel :: SPIRV.Backend -> Binary.Put
putMemoryModel bk
= putInstruction
Instruction
{ operation = SPIRV.Op.MemoryModel
, resTy = Nothing
, resID = Nothing
, args = Arg addressingModel
$ Arg memoryModel
EndArgs
}
where
addressingModel :: Word32
addressingModel = case bk of
SPIRV.Vulkan -> 0 -- logical addressing
SPIRV.OpenCL -> 2 -- physical 64-bit addressing
memoryModel :: Word32
memoryModel = case bk of
SPIRV.Vulkan -> 1 -- GLSL450 memory model
SPIRV.OpenCL -> 2 -- OpenCL memory model
putEntryPoint :: SPIRV.ExecutionModel -> ShortText -> ID -> Map ShortText ID -> Binary.Put
putEntryPoint model modelName entryPointID interface
= putInstruction
Instruction
{ operation = SPIRV.Op.EntryPoint
-- slight kludge to account for unusual parameters for OpEntryPoint
-- instead of result type, resTy field holds the ExecutionModel value
, resTy = Just . MkTyID $ SPIRV.executionModelID model
, resID = Just entryPointID
, args = Arg modelName
$ toArgs interface -- 'Map ShortText ID' has the appropriate traversable instance
}
putEntryPoints
:: Map (ShortText, SPIRV.ExecutionModel) ID
-> Map (ShortText, SPIRV.ExecutionModel) (Map ShortText ID)
-> ExceptT ShortText Binary.PutM ()
putEntryPoints entryPointIDs
= traverseWithKey_
( \(modelName, model) interface -> do
entryPointID
<- note
( "putEntryPoints: " <> ShortText.pack (show model)
<> " entry point named \"" <> modelName
<> "\" not bound to any ID."
)
( Map.lookup (modelName, model) entryPointIDs )
lift ( putEntryPoint model modelName entryPointID interface )
)
putModelExecutionModes :: ID -> SPIRV.ExecutionModes -> Binary.Put
putModelExecutionModes modelID
= traverse_
( \case
SPIRV.MaxPatchVertices {} -> pure () -- custom execution mode that doesn't exist in SPIR-V
mode -> putInstruction
Instruction
{ operation = SPIRV.Op.ExecutionMode
, resTy = Nothing
, resID = Nothing
, args = Arg modelID
$ Arg mode EndArgs
}
)
putExecutionModes
:: Map (ShortText, SPIRV.ExecutionModel) ID
-> Map (ShortText, SPIRV.ExecutionModel) SPIRV.ExecutionModes
-> ExceptT ShortText Binary.PutM ()
putExecutionModes entryPointIDs
= traverseWithKey_
( \(modelName, model) executionModes -> do
entryPointID
<- note
( "putExecutionModes: " <> ShortText.pack (show model)
<> " entry point named \"" <> modelName
<> "\" not bound to any ID."
)
( Map.lookup (modelName, model) entryPointIDs )
lift ( putModelExecutionModes entryPointID executionModes )
)
putKnownStringLits :: Map ShortText ID -> Binary.Put
putKnownStringLits
= traverseWithKey_
( \ lit ident -> putInstruction
Instruction
{ operation = SPIRV.Op.String
, resTy = Nothing
, resID = Just ident
, args = Arg lit EndArgs
}
)
putBindingAnnotations :: [ ( ShortText, ID ) ] -> Binary.Put
putBindingAnnotations
= traverse_
( \ ( name, ident ) -> putInstruction
Instruction
{ operation = SPIRV.Op.Name
, resTy = Nothing
, resID = Nothing
, args = Arg ident
$ Arg name EndArgs
}
)
putNames :: Set ( ID, Either ShortText (Word32, ShortText) ) -> Binary.Put
putNames = traverse_
( \case
( ident, Left name )
-> putInstruction
Instruction
{ operation = SPIRV.Op.Name
, resTy = Nothing
, resID = Nothing
, args = Arg ident
$ Arg name EndArgs
}
( ident, Right (index,name) )
-> putInstruction
Instruction
{ operation = SPIRV.Op.MemberName
, resTy = Nothing
, resID = Nothing
, args = Arg ident
$ Arg index
$ Arg name EndArgs
}
)
putDecorations :: Map ID SPIRV.Decorations -> Binary.Put
putDecorations
= traverseWithKey_
( \ decoratee ->
traverse_
( \ dec ->
putInstruction
Instruction
{ operation = SPIRV.Op.Decorate
, resTy = Nothing
, resID = Nothing
, args = Arg decoratee
$ Arg dec EndArgs
}
)
)
putMemberDecorations :: Map (TyID, Word32) SPIRV.Decorations -> Binary.Put
putMemberDecorations
= traverseWithKey_
( \ (structTyID, index) ->
traverse_
( \dec ->
putInstruction
Instruction
{ operation = SPIRV.Op.MemberDecorate
, resTy = Nothing
, resID = Nothing
, args = Arg structTyID
$ Arg index
$ Arg dec EndArgs
}
)
)
putTypesAndConstants
:: Map types Instruction
-> Map constants Instruction
-> Binary.Put
putTypesAndConstants ts cs
= traverse_ putInstruction
( sortOn resID $ Map.elems ts ++ Map.elems cs )
putUndefineds :: Map SPIRV.PrimTy (ID, TyID) -> Binary.Put
putUndefineds = traverse_
( \ ( undefID, undefTyID ) ->
putInstruction
Instruction
{ operation = SPIRV.Op.Undef
, resTy = Just undefTyID
, resID = Just undefID
, args = EndArgs
}
)
putGlobals :: [ Variable ] -> Binary.Put
putGlobals = traverse_ ( putInstruction . variableDeclaration )