Compilation to categories

knossos.cabal

@@ -12,6 +12,7 @@ executable ksc
                  CSE
                  Cgen
                  KMonad
+                 Ksc.CatLang
                  Ksc.Futhark
                  Ksc.Pipeline
                  Lang

src/ksc/AD.hs

@@ -119,6 +119,8 @@ gradBuild TupleAD s n ti body
      grad_body = mkLet (gradTVar TupleAD s ti)
                        (Tuple [Var ti, lmZero s (Var ti)]) $
                  gradE TupleAD s body
+gradBuild SplitAD _ _ _ _ = error "gradBuild:SplitAD"  -- Doesn't make sense
+
 ---------------
 gradFold :: ADPlan -> Shape -> TVar -> TExpr -> TExpr -> TExpr -> TExpr
 gradFold BasicAD s ti body acc v =
@@ -156,6 +158,8 @@ gradFold TupleAD s _ti _body acc _v =
   lmDummyFold (TypeTuple [t_acc, TypeLM (typeof s) t_acc])
   where t_acc = typeof acc
 
+gradFold SplitAD _ _ _ _ _ = error "gradFold:SplitAD"  -- Doesn't make sense
+
 ---------------
 gradCall :: ADPlan -> Shape -> TFun -> TExpr -> TExpr
 gradCall BasicAD s f args
@@ -194,6 +198,8 @@ gradCall TupleAD s f args
         grad_arg = gradE TupleAD s arg
         arg      = args
 
+gradCall SplitAD _ _ _ = error "gradCall:SplitAD"  -- Doesn't make sense
+
 ----------------------
 gradLet :: HasCallStack => ADPlan -> Shape -> TVar -> TExpr -> TExpr -> TExpr
 gradLet BasicAD s v e1 e2
@@ -212,6 +218,8 @@ gradLet TupleAD s v e1 e2
     gradE TupleAD s e2
   where dv = gradTVar TupleAD s v
 
+gradLet SplitAD _ _ _ _ = error "gradLet:SplitAD"  -- Doesn't make sense
+
 {- Note [Shadowing after grad]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 We transform   let x = <rhs>
@@ -234,7 +242,7 @@ lmVCat_AD :: ADPlan -> [TExpr] -> TExpr
 lmVCat_AD BasicAD ms = lmVCat ms
 lmVCat_AD TupleAD ms = Tuple [ Tuple  (map pFst ms)
                              , lmVCat (map pSnd ms) ]
-
+lmVCat_AD SplitAD _ = error "lmVCat_AD:SplitAD"  -- Doesn't make sense
 
 
 ---------------------------------

src/ksc/Annotate.hs

@@ -4,7 +4,7 @@
 	     ScopedTypeVariables, TypeApplications #-}
 
 module Annotate (
-  annotDecls, lintDefs
+  annotDecls, lintDefs, callResultTy_maybe, callResultTy
   ) where
 
 import Lang
@@ -214,42 +214,47 @@ tcVar var mb_ty
 -- callResultTy_maybe is given a (global) function and the type of its
 -- argument, and returns the type of its result.
 --
--- It has special cases for a bunch opuilt-in functions with polymorphic
+-- It has special cases for a bunch built-in functions with polymorphic
 -- types; that is, where the result type is a function of the argument types
 -- Otherwise it just looks in the global symbol table.
 callResultTy_maybe :: GblSymTab -> Fun -> Type
                    -> Either SDoc Type
-callResultTy_maybe env fun args
+callResultTy_maybe env fun arg_ty
   | is_user_fun fun
-  = userCallResultTy_maybe fun env args
+  = userCallResultTy_maybe fun env arg_ty
   | otherwise
-  = primCallResultTy_maybe fun (typeof args)
+  = primCallResultTy_maybe fun arg_ty
   where
     is_user_fun = isUserFun . funIdOfFun
 
+callResultTy :: GblSymTab -> Fun -> Type -> Type
+callResultTy gst fun arg_ty
+ = case callResultTy_maybe gst fun arg_ty of
+     Left err     -> pprTrace "callResultType:" err TypeUnknown
+     Right res_ty -> res_ty
+
 -----------------------------------------------
 --     The typecheck monad
 -----------------------------------------------
 
 userCallResultTy_maybe :: HasCallStack => Fun -> GblSymTab
                        -> Type -> Either SDoc Type
-userCallResultTy_maybe fn env args
-  = case lookupGblST (fn, typeof args) env of
-      Just def -> userCallResultTy_help def args
+userCallResultTy_maybe fn env arg_ty
+  = case lookupGblST (fn, arg_ty) env of
+      Just def -> userCallResultTy_help def arg_ty
       Nothing  -> Left (text "Not in scope: userCall:" <+> ppr fn $$ ppr env)
 
 userCallResultTy_help :: HasCallStack
                       => TDef -> Type -> Either SDoc Type
 userCallResultTy_help (Def { def_fun  = fn
                            , def_res_ty = ret_ty
                            , def_args = params })
-                      args
-  = case check_args bndr_tys arg_tys of
+                      arg_ty
+  = case check_args bndr_ty arg_ty of
       Just err -> Left err
       Nothing  -> Right ret_ty
   where
-    bndr_tys   = tVarType params
-    arg_tys    = typeof args
+    bndr_ty   = tVarType params
 
     check_args :: Type -> Type -> Maybe SDoc
     -- Return (Just err) if there's a wrong-ness
@@ -402,18 +407,18 @@ lookupLclTc v
            Just ty -> return ty }
 
 lookupGblTc :: Fun -> TypedExpr -> TcM Type
-lookupGblTc fun args
+lookupGblTc fun arg
   = do { st <- getSymTabTc
-       ; case callResultTy_maybe (gblST st) fun (typeof args) of
+       ; case callResultTy_maybe (gblST st) fun (typeof arg) of
            Left err -> do { addErr $ hang err 2 (mk_extra st)
                           ; return TypeUnknown }
            Right res_ty -> return res_ty }
   where
     mk_extra st
       = vcat [ text "In a call of:" <+> ppr fun <+> parens (text (show fun))
-             , text " Arg types:" <+> ppr (typeof args)
-             , text " Args:"      <+> ppr (exprOf args)
-             , text "ST lookup:"  <+> ppr (Map.lookup (fun, typeof args) (gblST st))
+             , text " Arg types:" <+> ppr (typeof arg)
+             , text " Args:"      <+> ppr (exprOf arg)
+             , text "ST lookup:"  <+> ppr (Map.lookup (fun, typeof arg) (gblST st))
              -- This is very verbose, and obscures error messages, but can be useful for typos.
              -- Perhaps think about printing it only for failed lookup of userfun
              -- , text "ST keys:" <+> gblDoc st

src/ksc/Cgen.hs

@@ -199,7 +199,7 @@ params_withPackedParams param = case typeof param of
   TypeTuple tys ->
     let params  = zipWith mkParam [1..] tys
         mkParam i ty = TVar ty (Simple name)
-          where name = nameOfVar (tVarVar param) ++ "arg" ++ show i
+          where name = tVarName param ++ "arg" ++ show i
         packParams = Let param (Tuple (map Var params))
     in (params, ensureDon'tReuseParams params . packParams)
   _             -> ([param], id)
@@ -525,6 +525,8 @@ cgenUserFun (f, ty) = case f of
   DrvFun   s (AD BasicAD Rev) -> "rev$" ++ cgenFunId (s, ty)
   DrvFun   s (AD TupleAD Fwd) -> "fwdt$" ++ cgenFunId (s, ty)
   DrvFun   s (AD TupleAD Rev) -> "revt$" ++ cgenFunId (s, ty)
+  DrvFun   s (AD SplitAD Fwd) -> "revsf$" ++ cgenFunId (s, ty)
+  DrvFun   s (AD SplitAD Rev) -> "revsr$" ++ cgenFunId (s, ty)
 
 cgenAnyFun :: HasCallStack => (TFun, Type) -> CType -> String
 cgenAnyFun (tf, ty) cftype = case tf of