diff options
Diffstat (limited to 'SC/Exp.hs')
| -rw-r--r-- | SC/Exp.hs | 62 |
1 files changed, 51 insertions, 11 deletions
@@ -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 |