Lots of stuff; can compile simple single-generate program

1 files changed, 51 insertions, 11 deletions

diff --git a/SC/Exp.hs b/SC/Exp.hs
index 2bd2b37..5ddd4bf 100644
--- a/SC/Exp.hs
+++ b/SC/Exp.hs
@@ -27,9 +27,13 @@ data CompiledFun aenv t1 t2 =
             --   the given names.
             --   The arguments will refer to array variable names found in the
             --   original array environment.
-        [TypedAName]
+        [SomeArray]
             -- ^ Arrays that the constructed arguments use from the environment
 
+-- | The variable names corresponding to a single source-level array (before
+-- SoA conversion).
+data SomeArray = forall sh t. SomeArray (ShNames sh) (ANames t)
+
 -- | The function must be single-argument. Uncurry if necessary (e.g. for zipWith).
 compileFun :: AVarEnv aenv -> A.Fun aenv (t1 -> t2) -> SC (CompiledFun aenv t1 t2)
 compileFun aenv (A.Lam lhs (A.Body body)) = do
@@ -40,14 +44,19 @@ compileFun aenv (A.Lam lhs (A.Body body)) = do
     (usedA, res) <- compileExp' aenv env body
     (sts1, retexprs) <- toStExprs (A.expType body) res
     let sts2 = genoutstores outnames retexprs
+        arrayarguments =
+            concatMap (\(SomeArray shn ans) ->
+                          map (\(TypedName t n) -> (t, n)) (shnamesList shn)
+                            ++ map (\(TypedAName t n) -> (t, n)) (itupList ans))
+                      usedA
         arguments =
-            map (\(TypedAName t n) -> (t, n)) usedA
+            arrayarguments
             ++ map (\(TypedName t n) -> (t, n)) (itupList argnames)
             ++ map (\(TypedName t n) -> (t, n)) (itupList outnames)
     return $ CompiledFun
         (C.ProcDef funname arguments (sts1 ++ sts2))
         (\argexprs destnames ->
-            map (\(TypedAName _ n) -> C.EVar n) usedA
+            map (C.EVar . snd) arrayarguments
             ++ itupList argexprs
             ++ map (\(TypedName _ n) -> C.EPtrTo (C.EVar n)) (itupList destnames))
         usedA
@@ -63,7 +72,7 @@ compileExp :: AVarEnv aenv -> A.Exp aenv t -> SC (CompiledFun aenv () t)
 compileExp aenv expr = compileFun aenv (A.Lam (LeftHandSideWildcard TupRunit) (A.Body expr))
 
 compileExp' :: AVarEnv aenv -> VarEnv env -> A.OpenExp env aenv t
-            -> SC ([TypedAName], Either (Names t -> [C.Stmt]) ([C.Stmt], Exprs t))
+            -> SC ([SomeArray], Either (Names t -> [C.Stmt]) ([C.Stmt], Exprs t))
 compileExp' aenv env = \case
     A.Let lhs rhs body -> do
         (names, env') <- genVarsEnv lhs env
@@ -76,6 +85,9 @@ compileExp' aenv env = \case
     A.Evar (Var _ idx) ->
         return ([], Right ([], ITupSingle (C.EVar (veprj env idx))))
 
+    A.Nil ->
+        return ([], Right ([], ITupIgnore))
+
     A.Pair a b -> do
         (usedA1, res1) <- compileExp' aenv env a
         (usedA2, res2) <- compileExp' aenv env b
@@ -84,6 +96,10 @@ compileExp' aenv env = \case
             ITupIgnore -> []
             ITupSingle _ -> error "wat"))
 
+    A.Const ty x
+      | Just str <- showExpConst ty x
+      -> return ([], Right ([], ITupSingle (C.ELit str)))
+
     A.PrimApp (A.PrimAdd _) e -> binary aenv env "+" e
     A.PrimApp (A.PrimSub _) e -> binary aenv env "-" e
     A.PrimApp (A.PrimMul _) e -> binary aenv env "*" e
@@ -94,7 +110,7 @@ compileExp' aenv env = \case
         let (shnames, _) = aveprj aenv idx
             buildExprs :: ShNames sh -> Exprs sh
             buildExprs ShZ = ITupIgnore
-            buildExprs (ShS n names) = ITupPair (buildExprs names) (ITupSingle (C.EVar n))
+            buildExprs (ShS names n) = ITupPair (buildExprs names) (ITupSingle (C.EVar n))
         in return ([], Right ([], buildExprs shnames))
 
     A.ToIndex shr she idxe -> do
@@ -120,15 +136,15 @@ compileExp' aenv env = \case
         let sts0 = [C.SDecl (C.TInt C.B64) temp Nothing]
         (usedA1, sts1) <- fmap (`toStoring` ITupSingle (TypedName (C.TInt C.B64) temp))
                               <$> compileExp' aenv env e
-        let (_, anames) = aveprj aenv idx
-            usedA = itupList anames ++ usedA1
+        let (shnames, anames) = aveprj aenv idx
+            usedA = SomeArray shnames anames : usedA1
         return (usedA, Right (sts0 ++ sts1
                              ,itupmap (\(TypedAName _ name) -> C.EIndex name (C.EVar temp)) anames))
 
-    _ -> throw "Unsupported Exp constructor"
+    e -> throw $ "Unsupported Exp constructor: " ++ A.showExpOp e
   where
     binary :: AVarEnv aenv -> VarEnv env -> String -> A.OpenExp env aenv (a, b)
-           -> SC ([TypedAName], Either (Names t -> [C.Stmt]) ([C.Stmt], Exprs t))
+           -> SC ([SomeArray], Either (Names t -> [C.Stmt]) ([C.Stmt], Exprs t))
     binary aenv' env' op e' = do
         (usedA, res) <- compileExp' aenv' env' e'
         (sts, ITupPair (ITupSingle e1) (ITupSingle e2)) <-
@@ -138,8 +154,9 @@ compileExp' aenv env = \case
 toStExprs :: TypeR t -> Either (Names t -> [C.Stmt]) ([C.Stmt], Exprs t) -> SC ([C.Stmt], Exprs t)
 toStExprs ty (Left fun) = do
     names <- genVars ty
-    let sts1 = fun names
-    return (sts1, itupmap (\(TypedName _ n) -> C.EVar n) names)
+    let sts1 = [C.SDecl t n Nothing | TypedName t n <- itupList names]
+        sts2 = fun names
+    return (sts1 ++ sts2, itupmap (\(TypedName _ n) -> C.EVar n) names)
 toStExprs _ (Right pair) = return pair
 
 toStoring :: Either (Names t -> [C.Stmt]) ([C.Stmt], Exprs t) -> Names t -> [C.Stmt]
@@ -153,6 +170,29 @@ toStoring (Right (sts, exs)) = (sts ++) . flip go exs
     go (ITupPair ns1 ns2) (ITupPair es1 es2) = go ns1 es1 ++ go ns2 es2
     go (ITupPair _ _) _ = error "wat"
 
+showExpConst :: ScalarType t -> t -> Maybe String
+showExpConst = \case
+    SingleScalarType (NumSingleType (IntegralNumType it)) -> Just . goI it
+    SingleScalarType (NumSingleType (FloatingNumType ft)) -> goF ft
+    VectorScalarType _ -> const Nothing
+  where
+    goI :: IntegralType t -> t -> String
+    goI TypeInt = (++ "LL") . show
+    goI TypeInt8 = ("(int8_t)" ++) . show
+    goI TypeInt16 = ("(int16_t)" ++) . show
+    goI TypeInt32 = show
+    goI TypeInt64 = (++ "LL") . show
+    goI TypeWord = (++ "ULL") . show
+    goI TypeWord8 = ("(uint8_t)" ++) . show
+    goI TypeWord16 = ("(uint16_t)" ++) . show
+    goI TypeWord32 = (++ "U") . show
+    goI TypeWord64 = (++ "ULL") . show
+
+    goF :: FloatingType t -> t -> Maybe String
+    goF TypeHalf = const Nothing
+    goF TypeFloat = Just . (++ "f") . show
+    goF TypeDouble = Just . show
+
 genVarsEnv :: A.ELeftHandSide t env env' -> VarEnv env -> SC (Names t, VarEnv env')
 genVarsEnv (LeftHandSideWildcard _) env = return (ITupIgnore, env)
 genVarsEnv (LeftHandSideSingle ty) env = do