{-# LANGUAGE GADTs #-} {-# LANGUAGE LambdaCase #-} {-# LANGUAGE PolyKinds #-} {-# LANGUAGE TypeApplications #-} module Compile (compile) where import Control.Monad.Trans.State.Strict import Data.Bifunctor (first, second) import Data.Foldable (toList) import Data.Functor.Const import qualified Data.Functor.Product as Product import Data.Functor.Product (Product) import Data.List (intersperse, intercalate) import qualified Data.Map.Strict as Map import qualified Data.Set as Set import Data.Set (Set) import Data.Some import qualified Data.Vector as V import Foreign import Array import AST import AST.Pretty (ppTy) import Compile.Exec import Data import Interpreter.Rep -- In shape and index arrays, the innermost dimension is on the right (last index). -- TODO: array lifetimes in C? compile :: SList STy env -> Ex env t -> IO (SList Value env -> IO (Rep t)) compile = \env expr -> do lib <- buildKernel (compileToString env expr) ["kernel"] let arg_metrics = reverse (unSList metricsSTy env) (arg_offsets, result_offset) = computeStructOffsets arg_metrics result_type = typeOf expr result_size = sizeofSTy result_type return $ \val -> do allocaBytes (result_offset + result_size) $ \ptr -> do let args = zip (reverse (unSList Some (slistZip env val))) arg_offsets serialiseArguments args ptr $ do callKernelFun "kernel" lib ptr deserialise result_type ptr result_offset where serialiseArguments :: [(Some (Product STy Value), Int)] -> Ptr () -> IO r -> IO r serialiseArguments ((Some (Product.Pair t (Value arg)), off) : args) ptr k = serialise t arg ptr off $ serialiseArguments args ptr k serialiseArguments _ _ k = k data StructDecl = StructDecl String -- ^ name String -- ^ contents String -- ^ comment deriving (Show) data Stmt = SVarDecl Bool String String CExpr -- ^ const, type, variable name, right-hand side | SVarDeclUninit String String -- ^ type, variable name (no initialiser) | SAsg String CExpr -- ^ variable name, right-hand side | SBlock [Stmt] | SIf CExpr [Stmt] [Stmt] | SVerbatim String -- ^ no implicit ';', just printed as-is deriving (Show) data CExpr = CELit String -- ^ inserted as-is, assumed no parentheses needed | CEStruct String [(String, CExpr)] -- ^ struct construction literal: `(name){.field=expr}` | CEProj CExpr String -- ^ field projection: expr.field | CECall String [CExpr] -- ^ function(arg1, ..., argn) | CEBinop CExpr String CExpr -- ^ expr + expr | CEIf CExpr CExpr CExpr -- ^ expr ? expr : expr deriving (Show) printStructDecl :: StructDecl -> ShowS printStructDecl (StructDecl name contents comment) = showString "typedef struct { " . showString contents . showString " } " . showString name . showString ("; // " ++ comment) printStmt :: Int -> Stmt -> ShowS printStmt indent = \case SVarDecl cnst typ name rhs -> showString ((if cnst then "const " else "") ++ typ ++ " " ++ name ++ " = ") . printCExpr 0 rhs . showString ";" SVarDeclUninit typ name -> showString (typ ++ " " ++ name ++ ";") SAsg name rhs -> showString (name ++ " = ") . printCExpr 0 rhs . showString ";" SBlock stmts -> showString "{" . compose [showString ("\n" ++ replicate (2*indent+2) ' ') . printStmt (indent+1) stmt | stmt <- stmts] . showString ("\n" ++ replicate (2*indent) ' ' ++ "}") SIf cond b1 b2 -> showString "if (" . printCExpr 0 cond . showString ") " . printStmt indent (SBlock b1) . showString " else " . printStmt indent (SBlock b2) SVerbatim s -> showString s -- d values: -- * 0: top level -- * 1: in 1st or 2nd component of a ternary operator (technically same as top level, but readability) -- * 2-...: various operators (see precTable) -- * 98: inside unknown operator -- * 99: left of a field projection -- Unlisted operators are conservatively written with full parentheses. printCExpr :: Int -> CExpr -> ShowS printCExpr d = \case CELit s -> showString s CEStruct name pairs -> showParen (d >= 99) $ showString ("(" ++ name ++ "){") . compose (intersperse (showString ", ") [showString ("." ++ n ++ " = ") . printCExpr 0 e | (n, e) <- pairs]) . showString "}" CEProj e name -> printCExpr 99 e . showString ("." ++ name) CECall n es -> showString (n ++ "(") . compose (intersperse (showString ", ") (map (printCExpr 0) es)) . showString ")" CEBinop e1 n e2 -> let mprec = Map.lookup n precTable p = maybe (-1) fst mprec -- precedence of this operator (d1, d2) = maybe (98, 98) snd mprec -- precedences for the arguments in showParen (d > p) $ printCExpr d1 e1 . showString (" " ++ n ++ " ") . printCExpr d2 e2 CEIf e1 e2 e3 -> showParen (d > 0) $ printCExpr 1 e1 . showString " ? " . printCExpr 1 e2 . showString " : " . printCExpr 0 e3 where precTable = Map.fromList [("||", (2, (2, 2))) ,("&&", (3, (3, 3))) ,("==", (4, (5, 5))) ,("!=", (4, (5, 5))) ,("<", (5, (6, 6))) ,(">", (5, (6, 6))) ,("<=", (5, (6, 6))) ,(">=", (5, (6, 6))) ,("+", (6, (6, 6))) ,("-", (6, (6, 7))) ,("*", (7, (7, 7))) ,("/", (7, (7, 8))) ,("%", (7, (7, 8)))] repTy :: Ty -> String repTy (TScal st) = case st of TI32 -> "int32_t" TI64 -> "int64_t" TF32 -> "float" TF64 -> "double" TBool -> "uint8_t" repTy t = genStructName t repSTy :: STy t -> String repSTy = repTy . unSTy genStructName :: Ty -> String genStructName = \t -> "ty_" ++ gen t where -- all tags start with a letter, so the array mangling is unambiguous. gen :: Ty -> String gen TNil = "n" gen (TPair a b) = 'P' : gen a ++ gen b gen (TEither a b) = 'E' : gen a ++ gen b gen (TMaybe t) = 'M' : gen t gen (TArr n t) = "A" ++ show (fromNat n) ++ gen t gen (TScal st) = case st of TI32 -> "i" TI64 -> "j" TF32 -> "f" TF64 -> "d" TBool -> "b" gen (TAccum t) = 'C' : gen t genStruct :: String -> Ty -> Maybe StructDecl genStruct name topty = case topty of TNil -> Just $ StructDecl name "" com TPair a b -> Just $ StructDecl name (repTy a ++ " a; " ++ repTy b ++ " b;") com TEither a b -> -- 0 -> a, 1 -> b Just $ StructDecl name ("uint8_t tag; union { " ++ repTy a ++ " a; " ++ repTy b ++ " b; };") com TMaybe t -> -- 0 -> nothing, 1 -> just Just $ StructDecl name ("uint8_t tag; " ++ repTy t ++ " a;") com TArr n t -> Just $ StructDecl name ("size_t sh[" ++ show (fromNat n) ++ "]; " ++ repTy t ++ " *a;") com TScal _ -> Nothing TAccum t -> Just $ StructDecl name (repTy t ++ " a;") com where com = ppTy 0 topty -- State: (already-generated (skippable) struct names, the structs in declaration order) genStructs :: Ty -> State (Set String, Bag StructDecl) () genStructs ty = do let name = genStructName ty seen <- gets ((name `Set.member`) . fst) case (if seen then Nothing else genStruct name ty) of Nothing -> pure () Just decl -> do -- already mark this struct as generated now, so we don't generate it twice modify (first (Set.insert name)) case ty of TNil -> pure () TPair a b -> genStructs a >> genStructs b TEither a b -> genStructs a >> genStructs b TMaybe t -> genStructs t TArr _ t -> genStructs t TScal _ -> pure () TAccum t -> genStructs t modify (second (<> pure decl)) genAllStructs :: Foldable t => t Ty -> [StructDecl] genAllStructs tys = toList . snd $ execState (mapM_ genStructs tys) (mempty, mempty) data CompState = CompState { csStructs :: Set Ty , csStmts :: Bag Stmt , csNextId :: Int } deriving (Show) type CompM a = State CompState a genId :: CompM Int genId = state $ \s -> (csNextId s, s { csNextId = csNextId s + 1 }) genName :: CompM String genName = ('x' :) . show <$> genId emit :: Stmt -> CompM () emit stmt = modify $ \s -> s { csStmts = csStmts s <> pure stmt } scope :: CompM a -> CompM (a, [Stmt]) scope m = do stmts <- state $ \s -> (csStmts s, s { csStmts = mempty }) res <- m innerStmts <- state $ \s -> (csStmts s, s { csStmts = stmts }) return (res, toList innerStmts) emitStruct :: STy t -> CompM String emitStruct ty = do let ty' = unSTy ty modify $ \s -> s { csStructs = Set.insert ty' (csStructs s) } return (genStructName ty') nameEnv :: SList f env -> SList (Const String) env nameEnv = flip evalState (0::Int) . slistMapA (\_ -> state $ \i -> (Const ("arg" ++ show i), i + 1)) compileToString :: SList STy env -> Ex env t -> String compileToString env expr = let args = nameEnv env (res, s) = runState (compile' args expr) (CompState mempty mempty 1) structs = genAllStructs (csStructs s <> Set.fromList (unSList unSTy env)) (arg_pairs, arg_metrics) = unzip $ reverse (unSList (\(Product.Pair t (Const n)) -> ((n, repSTy t), metricsSTy t)) (slistZip env args)) (arg_offsets, result_offset') = computeStructOffsets arg_metrics result_offset = align (alignmentSTy (typeOf expr)) result_offset' in ($ "") $ compose [showString "#include \n" ,showString "#include \n\n" ,compose $ map (\sd -> printStructDecl sd . showString "\n") structs ,showString "\n" ,showString $ "static " ++ repSTy (typeOf expr) ++ " typed_kernel(" ++ intercalate ", " (reverse (unSList (\(Product.Pair t (Const n)) -> repSTy t ++ " " ++ n) (slistZip env args))) ++ ") {\n" ,compose $ map (\st -> showString " " . printStmt 1 st . showString "\n") (toList (csStmts s)) ,showString (" return ") . printCExpr 0 res . showString ";\n}\n\n" ,showString "void kernel(void *data) {\n" -- Some code here assumes that we're on a 64-bit system, so let's check that ,showString " if (sizeof(void*) != 8) { abort(); }\n" ,showString $ " *(" ++ repSTy (typeOf expr) ++ "*)(data + " ++ show result_offset ++ ") = typed_kernel(" ++ concat (map (\((arg, typ), off, idx) -> "\n *(" ++ typ ++ "*)(data + " ++ show off ++ ")" ++ (if idx < length arg_pairs - 1 then "," else "") ++ " // " ++ arg) (zip3 arg_pairs arg_offsets [0::Int ..])) ++ "\n );\n" ,showString "}\n"] -- | Takes list of metrics (alignment, sizeof). -- Returns (offsets, size of struct). computeStructOffsets :: [(Int, Int)] -> ([Int], Int) computeStructOffsets = go 0 0 where go off maxal [(al, sz)] = ([off], align (max maxal al) (off + sz)) go off maxal ((al, sz) : pairs@((al2,_):_)) = first (off :) $ go (align al2 (off + sz)) (max maxal al) pairs go _ _ [] = ([], 0) -- | Assumes that this is called at the correct alignment. serialise :: STy t -> Rep t -> Ptr () -> Int -> IO r -> IO r serialise topty topval ptr off k = -- TODO: this code is quadratic in the depth of the type because of the alignment/sizeOf calls case (topty, topval) of (STNil, ()) -> k (STPair a b, (x, y)) -> serialise a x ptr off $ serialise b y ptr (align (alignmentSTy b) (off + sizeofSTy a)) k (STEither a _, Left x) -> do pokeByteOff ptr off (0 :: Word8) -- alignment of (a + b) is alignment of (union {a b}) serialise a x ptr (off + alignmentSTy topty) k (STEither _ b, Right y) -> do pokeByteOff ptr off (1 :: Word8) serialise b y ptr (off + alignmentSTy topty) k (STMaybe _, Nothing) -> do pokeByteOff ptr off (0 :: Word8) k (STMaybe t, Just x) -> do pokeByteOff ptr off (1 :: Word8) serialise t x ptr (off + alignmentSTy t) k (STArr n t, Array sh vec) -> do pokeShape ptr off n sh let off1 = off + 8 * fromSNat n eltsz = sizeofSTy t allocaBytes (shapeSize sh * sizeofSTy t) $ \arrptr -> let loop i | i == shapeSize sh = k | otherwise = serialise t (vec V.! i) arrptr (off1 + i * eltsz) $ loop (i+1) in loop 0 (STScal sty, x) -> case sty of STI32 -> pokeByteOff ptr off (x :: Int32) >> k STI64 -> pokeByteOff ptr off (x :: Int64) >> k STF32 -> pokeByteOff ptr off (x :: Float) >> k STF64 -> pokeByteOff ptr off (x :: Double) >> k STBool -> pokeByteOff ptr off (fromIntegral (fromEnum x) :: Word8) >> k (STAccum{}, _) -> error "Cannot serialise accumulators" -- | Assumes that this is called at the correct alignment. deserialise :: STy t -> Ptr () -> Int -> IO (Rep t) deserialise topty ptr off = -- TODO: this code is quadratic in the depth of the type because of the alignment/sizeOf calls case topty of STNil -> return () STPair a b -> do x <- deserialise a ptr off y <- deserialise b ptr (align (alignmentSTy b) (off + sizeofSTy a)) return (x, y) STEither a b -> do tag <- peekByteOff @Word8 ptr off if tag == 0 -- alignment of (a + b) is alignment of (union {a b}) then Left <$> deserialise a ptr (off + alignmentSTy topty) else Right <$> deserialise b ptr (off + alignmentSTy topty) STMaybe t -> do tag <- peekByteOff @Word8 ptr off if tag == 0 then return Nothing else Just <$> deserialise t ptr (off + alignmentSTy t) STArr n t -> do sh <- peekShape ptr off n let off1 = off + 8 * fromSNat n eltsz = sizeofSTy t Array sh <$> V.generateM (shapeSize sh) (\i -> deserialise t ptr (off1 + i * eltsz)) STScal sty -> case sty of STI32 -> peekByteOff @Int32 ptr off STI64 -> peekByteOff @Int64 ptr off STF32 -> peekByteOff @Float ptr off STF64 -> peekByteOff @Double ptr off STBool -> toEnum . fromIntegral <$> peekByteOff @Word8 ptr off STAccum{} -> error "Cannot serialise accumulators" align :: Int -> Int -> Int align a off = (off + a - 1) `div` a * a alignmentSTy :: STy t -> Int alignmentSTy = fst . metricsSTy sizeofSTy :: STy t -> Int sizeofSTy = snd . metricsSTy -- | Returns (alignment, sizeof) metricsSTy :: STy t -> (Int, Int) metricsSTy STNil = (1, 0) metricsSTy (STPair a b) = let (a1, s1) = metricsSTy a (a2, s2) = metricsSTy b in (max a1 a2, align (max a1 a2) (s1 + s2)) metricsSTy (STEither a b) = let (a1, s1) = metricsSTy a (a2, s2) = metricsSTy b in (max a1 a2, max a1 a2 + max s1 s2) -- the union after the tag byte is aligned metricsSTy (STMaybe t) = let (a, s) = metricsSTy t in (a, a + s) -- the union after the tag byte is aligned metricsSTy (STArr n _) = (8, fromSNat n * 8 + 8) metricsSTy (STScal sty) = case sty of STI32 -> (4, 4) STI64 -> (8, 8) STF32 -> (4, 4) STF64 -> (8, 8) STBool -> (1, 1) -- compiled to uint8_t metricsSTy (STAccum t) = metricsSTy t pokeShape :: Ptr () -> Int -> SNat n -> Shape n -> IO () pokeShape ptr off = go . fromSNat where go :: Int -> Shape n -> IO () go rank = \case ShNil -> return () sh `ShCons` n -> do pokeByteOff ptr (off + (rank - 1) * 8) (fromIntegral n :: Int64) go (rank - 1) sh peekShape :: Ptr () -> Int -> SNat n -> IO (Shape n) peekShape ptr off = \case SZ -> return ShNil SS n -> ShCons <$> peekShape ptr off n <*> (fromIntegral <$> peekByteOff @Int64 ptr (off + (fromSNat n) * 8)) compile' :: SList (Const String) env -> Ex env t -> CompM CExpr compile' env = \case EVar _ _ i -> return $ CELit (getConst (slistIdx env i)) ELet _ rhs body -> do e <- compile' env rhs var <- genName emit $ SVarDecl True (repSTy (typeOf rhs)) var e compile' (Const var `SCons` env) body EPair _ a b -> do name <- emitStruct (STPair (typeOf a) (typeOf b)) e1 <- compile' env a e2 <- compile' env b return $ CEStruct name [("a", e1), ("b", e2)] EFst _ e -> CEProj <$> compile' env e <*> pure "a" ESnd _ e -> CEProj <$> compile' env e <*> pure "b" ENil _ -> do name <- emitStruct STNil return $ CEStruct name [] EInl _ t e -> do name <- emitStruct (STEither (typeOf e) t) e1 <- compile' env e return $ CEStruct name [("tag", CELit "0"), ("a", e1)] EInr _ t e -> do name <- emitStruct (STEither t (typeOf e)) e2 <- compile' env e return $ CEStruct name [("tag", CELit "1"), ("b", e2)] ECase _ (EOp _ OIf e) a b -> do e1 <- compile' env e (e2, stmts2) <- scope $ compile' (Const undefined `SCons` env) a -- don't access that nil, stupid you (e3, stmts3) <- scope $ compile' (Const undefined `SCons` env) b retvar <- genName emit $ SVarDeclUninit (repSTy (typeOf a)) retvar emit $ SIf e1 (stmts2 <> pure (SAsg retvar e2)) (stmts3 <> pure (SAsg retvar e3)) return (CELit retvar) ECase _ e a b -> do let STEither t1 t2 = typeOf e e1 <- compile' env e var <- genName fieldvar <- genName (e2, stmts2) <- scope $ compile' (Const fieldvar `SCons` env) a (e3, stmts3) <- scope $ compile' (Const fieldvar `SCons` env) b retvar <- genName emit $ SVarDeclUninit (repSTy (typeOf a)) retvar emit $ SBlock (pure (SVarDecl True (repSTy (typeOf e)) var e1) <> pure (SIf (CEBinop (CEProj (CELit var) "tag") "==" (CELit "0")) (pure (SVarDecl True (repSTy t1) fieldvar (CEProj (CELit var) "a")) <> stmts2 <> pure (SAsg retvar e2)) (pure (SVarDecl True (repSTy t2) fieldvar (CEProj (CELit var) "b")) <> stmts3 <> pure (SAsg retvar e3)))) return (CELit retvar) ENothing _ t -> do name <- emitStruct (STMaybe t) return $ CEStruct name [("tag", CELit "0")] EJust _ e -> do name <- emitStruct (STMaybe (typeOf e)) e1 <- compile' env e return $ CEStruct name [("tag", CELit "1"), ("a", e1)] EMaybe _ a b e -> do e1 <- compile' env e var <- genName fieldvar <- genName (e2, stmts2) <- scope $ compile' env a (e3, stmts3) <- scope $ compile' (Const fieldvar `SCons` env) b retvar <- genName emit $ SVarDeclUninit (repSTy (typeOf a)) retvar emit $ SBlock (pure (SVarDecl True (repSTy (typeOf e)) var e1) <> pure (SIf (CEBinop (CEProj (CELit var) "tag") "==" (CELit "0")) (stmts2 <> pure (SAsg retvar e2)) (pure (SVarDecl True (repSTy (typeOf b)) fieldvar (CEProj (CELit var) "a")) <> stmts3 <> pure (SAsg retvar e3)))) return (CELit retvar) -- EConstArr _ n t arr -> do -- name <- emitStruct (STArr n (STScal t)) -- error "TODO" -- EBuild _ n a b -> error "TODO" -- genStruct (STArr n (typeOf b)) <> EBuild ext n (compile' a) (compile' b) -- EFold1Inner _ a b c -> error "TODO" -- EFold1Inner ext (compile' a) (compile' b) (compile' c) -- ESum1Inner _ e -> error "TODO" -- ESum1Inner ext (compile' e) -- EUnit _ e -> error "TODO" -- EUnit ext (compile' e) -- EReplicate1Inner _ a b -> error "TODO" -- EReplicate1Inner ext (compile' a) (compile' b) -- EMaximum1Inner _ e -> error "TODO" -- EMaximum1Inner ext (compile' e) -- EMinimum1Inner _ e -> error "TODO" -- EMinimum1Inner ext (compile' e) EConst _ t x -> case t of STI32 -> return $ CELit $ "(int32_t)" ++ show x STI64 -> return $ CELit $ "(int64_t)" ++ show x STF32 -> return $ CELit $ show x ++ "f" STF64 -> return $ CELit $ show x STBool -> return $ CELit $ if x then "1" else "0" -- EIdx0 _ e -> error "TODO" -- EIdx0 ext (compile' e) -- EIdx1 _ a b -> error "TODO" -- EIdx1 ext (compile' a) (compile' b) -- EIdx _ a b -> error "TODO" -- EIdx ext (compile' a) (compile' b) -- EShape _ e -> error "TODO" -- EShape ext (compile' e) EOp _ op (EPair _ e1 e2) -> do e1' <- compile' env e1 e2' <- compile' env e2 compileOpPair op e1' e2' EOp _ op e -> do e' <- compile' env e compileOpGeneral op e' -- ECustom _ t1 t2 t3 a b c e1 e2 -> error "TODO" -- ECustom ext t1 t2 t3 (compile' a) (compile' b) (compile' c) (compile' e1) (compile' e2) -- EWith _ a b -> error "TODO" -- EWith (compile' a) (compile' b) -- EAccum _ n a b e -> error "TODO" -- EAccum n (compile' a) (compile' b) (compile' e) EError _ t s -> do name <- emitStruct t -- using 'show' here is wrong, but it's good enough for me. emit $ SVerbatim $ "fprintf(stderr, \"ERROR: %s\\n\", " ++ show s ++ "); exit(1);" return $ CEStruct name [] EZero{} -> error "Compile: monoid operations should have been eliminated" EPlus{} -> error "Compile: monoid operations should have been eliminated" EOneHot{} -> error "Compile: monoid operations should have been eliminated" _ -> error "Compile: not implemented" compileOpGeneral :: SOp a b -> CExpr -> CompM CExpr compileOpGeneral op e1 = do let unary cop = return @(State CompState) $ CECall cop [e1] let binary cop = do name <- genName emit $ SVarDecl True (repSTy (opt1 op)) name e1 return $ CEBinop (CEProj (CELit name) "a") cop (CEProj (CELit name) "b") case op of OAdd _ -> binary "+" OMul _ -> binary "*" ONeg _ -> unary "-" OLt _ -> binary "<" OLe _ -> binary "<=" OEq _ -> binary "==" ONot -> unary "!" OAnd -> binary "&&" OOr -> binary "||" OIf -> do name <- emitStruct (STEither STNil STNil) _ <- emitStruct STNil return $ CEIf e1 (CEStruct name [("tag", CELit "0")]) (CEStruct name [("tag", CELit "1")]) ORound64 -> unary "(int64_t)round" -- ew OToFl64 -> unary "(double)" ORecip _ -> return $ CEBinop (CELit "1.0") "/" e1 OExp STF32 -> unary "expf" OExp STF64 -> unary "exp" OLog STF32 -> unary "logf" OLog STF64 -> unary "log" OIDiv _ -> binary "/" compileOpPair :: SOp a b -> CExpr -> CExpr -> CompM CExpr compileOpPair op e1 e2 = do let binary cop = return @(State CompState) $ CEBinop e1 cop e2 case op of OAdd _ -> binary "+" OMul _ -> binary "*" OLt _ -> binary "<" OLe _ -> binary "<=" OEq _ -> binary "==" OAnd -> binary "&&" OOr -> binary "||" OIDiv _ -> binary "/" _ -> error "compileOpPair: got unary operator" compose :: Foldable t => t (a -> a) -> a -> a compose = foldr (.) id