Skip to content

Commit

Permalink
fix route ends
Browse files Browse the repository at this point in the history
  • Loading branch information
vidsinghal committed Nov 11, 2024
1 parent 315dbac commit 2c525ea
Show file tree
Hide file tree
Showing 4 changed files with 113 additions and 44 deletions.
64 changes: 64 additions & 0 deletions gibbon-compiler/src/Gibbon/L2/Syntax.hs
Original file line number Diff line number Diff line change
Expand Up @@ -275,6 +275,27 @@ instance (Out l, Show l, Typeable (E2 l (UrTy l))) => Typeable (E2Ext l (UrTy l)
SSPush{} -> ProdTy []
SSPop{} -> ProdTy []

gRecoverTypeLoc ddfs env2 ex =
case ex of
LetRegionE _r _ _ bod -> gRecoverTypeLoc ddfs env2 bod
LetParRegionE _r _ _ bod -> gRecoverTypeLoc ddfs env2 bod
LetLocE _l _rhs bod -> gRecoverTypeLoc ddfs env2 bod
StartOfPkdCursor{} -> CursorTy
TagCursor{} -> CursorTy
RetE _loc var -> case M.lookup (singleLocVar var) (vEnv env2) of
Just ty -> ty
Nothing -> error $ "gRecoverTypeLoc: unbound variable " ++ sdoc var
FromEndE _loc -> error "Shouldn't enconter FromEndE in tail position"
BoundsCheck{} -> error "Shouldn't enconter BoundsCheck in tail position"
IndirectionE tycon _ _ (to,_) _ -> PackedTy tycon to
AddFixed{} -> error "Shouldn't enconter AddFixed in tail position"
GetCilkWorkerNum -> IntTy
LetAvail _ bod -> gRecoverTypeLoc ddfs env2 bod
AllocateTagHere{} -> ProdTy []
AllocateScalarsHere{} -> ProdTy []
SSPush{} -> ProdTy []
SSPop{} -> ProdTy []

instance (Typeable (E2Ext l d),
Expression (E2Ext l d),
Flattenable (E2 l d))
Expand Down Expand Up @@ -553,6 +574,49 @@ instance Typeable (PreExp E2Ext LocVar (UrTy LocVar)) where
in gRecoverType ddfs env2' e


gRecoverTypeLoc ddfs env2 ex =
case ex of
VarE v -> M.findWithDefault (error $ "Cannot find type of variable " ++ show v ++ " in " ++ show (vEnv env2)) (singleLocVar v) (vEnv env2)
LitE _ -> IntTy
CharE{} -> CharTy
FloatE{} -> FloatTy
LitSymE _ -> SymTy
AppE v locs _ -> let fnty = fEnv env2 # (singleLocVar v)
outty = arrOut fnty
mp = M.fromList $ zip (allLocVars fnty) locs
in substLoc mp outty

PrimAppE (DictInsertP ty) ((VarE v):_) -> SymDictTy (Just v) $ stripTyLocs ty
PrimAppE (DictEmptyP ty) ((VarE v):_) -> SymDictTy (Just v) $ stripTyLocs ty
PrimAppE p _ -> primRetTy p

LetE (v,_,t,_) e -> gRecoverTypeLoc ddfs (extendVEnvLocVar (singleLocVar v) t env2) e
IfE _ e _ -> gRecoverTypeLoc ddfs env2 e
MkProdE es -> ProdTy $ L.map (gRecoverTypeLoc ddfs env2) es
DataConE loc c _ -> PackedTy (getTyOfDataCon ddfs c) loc
TimeIt e _ _ -> gRecoverTypeLoc ddfs env2 e
MapE _ e -> gRecoverTypeLoc ddfs env2 e
FoldE _ _ e -> gRecoverTypeLoc ddfs env2 e
Ext ext -> gRecoverTypeLoc ddfs env2 ext
ProjE i e ->
case gRecoverTypeLoc ddfs env2 e of
(ProdTy tys) -> tys !! i
oth -> error$ "typeExp: Cannot project fields from this type: "++show oth
++"\nExpression:\n "++ sdoc ex
++"\nEnvironment:\n "++sdoc (vEnv env2)
SpawnE v locs _ -> let fnty = fEnv env2 # (singleLocVar v)
outty = arrOut fnty
mp = M.fromList $ zip (allLocVars fnty) locs
in substLoc mp outty
SyncE -> voidTy
WithArenaE _v e -> gRecoverTypeLoc ddfs env2 e
CaseE _ mp ->
let (c,vlocs,e) = head mp
(vars,locs) = unzip vlocs
env2' = extendPatternMatchEnvLocVar c ddfs vars locs env2
in gRecoverTypeLoc ddfs env2' e


--------------------------------------------------------------------------------
-- Do this manually to get prettier formatting: (Issue #90)

Expand Down
1 change: 1 addition & 0 deletions gibbon-compiler/src/Gibbon/Language/Syntax.hs
Original file line number Diff line number Diff line change
Expand Up @@ -662,6 +662,7 @@ type HasSimplifiableExt e l d = ( Show l, Out l, Show d, Out d
-- bind it with a LetE.
class Expression e => Typeable e where
gRecoverType :: DDefs (TyOf e) -> Env2 Var (TyOf e) -> e -> TyOf e
gRecoverTypeLoc :: DDefs (TyOf e) -> Env2 LocVar (TyOf e) -> e -> TyOf e

-- | Generic substitution over expressions.
class Expression e => Substitutable e where
Expand Down
9 changes: 4 additions & 5 deletions gibbon-compiler/src/Gibbon/Passes/ParAlloc.hs
Original file line number Diff line number Diff line change
Expand Up @@ -167,9 +167,8 @@ parAllocExp ddefs fundefs env2 reg_env after_env mb_parent_id pending_binds spaw
env2' = extendVEnvLocVar (Single v) ty env2

vars = gFreeVars (substLocInExp after_env rhs) `S.difference` (M.keysSet fundefs)
used = (allFreeVars (substLocInExp after_env rhs)) `S.difference` (M.keysSet fundefs)
used' = S.map Single used

used = (allFreeVars (substLocInExp after_env rhs)) `S.difference` (S.map singleLocVar (M.keysSet fundefs))

-- Swallow this binding, and add v to 'spawned'
if not (S.disjoint vars spawned)
then do
Expand All @@ -178,13 +177,13 @@ parAllocExp ddefs fundefs env2 reg_env after_env mb_parent_id pending_binds spaw
spawned' = S.insert v spawned
parAllocExp ddefs fundefs env2' reg_env' after_env mb_parent_id pending_binds'' spawned' boundlocs region_on_spawn bod
-- Swallow this binding, and but don't add v to 'spawned'
else if not (S.isSubsetOf used' boundlocs)
else if not (S.isSubsetOf used boundlocs)
then do
rhs' <- go rhs
let pending_binds'' = PVar (v, locs, ty', rhs') : pending_binds'
parAllocExp ddefs fundefs env2' reg_env' after_env mb_parent_id pending_binds'' spawned boundlocs region_on_spawn bod
-- Emit this binding as usual
else if S.disjoint vars spawned && S.isSubsetOf used' boundlocs
else if S.disjoint vars spawned && S.isSubsetOf used boundlocs
then do
let boundlocs' = S.insert (Single v) boundlocs `S.union` (S.fromList locs)
LetE <$> (v,locs,ty',) <$> go rhs
Expand Down
83 changes: 44 additions & 39 deletions gibbon-compiler/src/Gibbon/Passes/RouteEnds.hs
Original file line number Diff line number Diff line change
Expand Up @@ -205,7 +205,7 @@ routeEnds prg@Prog{ddefs,fundefs,mainExp} = do
-- Then process the actual function bodies using the new fundefs structure:
fds'' <- mapM (fd fundefs') fds'
let fundefs'' = M.fromList $ L.map (\f -> (funName f,f)) fds''
env2 = progToEnv prg
env2 = progToEnv' prg
-- Handle the main expression (if it exists):
mainExp' <- case mainExp of
Nothing -> return Nothing
Expand Down Expand Up @@ -235,18 +235,19 @@ routeEnds prg@Prog{ddefs,fundefs,mainExp} = do
do let (ArrowTy2 locin tyins eff _tyout _locout _isPar) = funTy
handleLoc (LRM l _r _m) ls = if S.member (Traverse l) eff then l:ls else ls
retlocs = L.foldr handleLoc [] locin
funArgs' = L.map singleLocVar funArgs
lenv = L.foldr
(\(a,t) acc -> case t of
PackedTy _ loc -> M.insert a loc acc
_ -> acc)
M.empty (zip funArgs tyins)
M.empty (zip funArgs' tyins)
pakdLocs = concatMap
(\t -> case t of
PackedTy _ loc -> [(loc, t)]
ProdTy ys -> pakdLocs ys
_ -> [])
initVEnv = M.fromList $ pakdLocs tyins ++ zip funArgs tyins
env2 = Env2 initVEnv (initFunEnv fundefs)
initVEnv = M.fromList $ pakdLocs tyins ++ zip funArgs' tyins
env2 = Env2 initVEnv (initFunEnv' fundefs)
funBody' <- bindReturns funBody
funBody'' <- exp fns retlocs emptyRel lenv M.empty env2 funBody'
return FunDef{funName,funTy,funArgs,funBody=funBody'',funMeta}
Expand All @@ -260,8 +261,8 @@ routeEnds prg@Prog{ddefs,fundefs,mainExp} = do
-- 4. a map of var to location
-- 5. a map from location to location after it
-- 6. the expression to process
exp :: FunDefs2 -> [LocVar] -> EndOfRel -> M.Map Var LocVar ->
M.Map LocVar LocVar -> Env2 Ty2 -> Exp2 -> PassM Exp2
exp :: FunDefs2 -> [LocVar] -> EndOfRel -> M.Map LocVar LocVar ->
M.Map LocVar LocVar -> Env2 LocVar Ty2 -> Exp2 -> PassM Exp2
exp fns retlocs eor lenv afterenv env2 e =
case e of

Expand Down Expand Up @@ -291,21 +292,21 @@ routeEnds prg@Prog{ddefs,fundefs,mainExp} = do
-- the end witnesses returned from the function.
LetE (v,_ls,ty,(AppE f lsin e1)) e2 -> do
let lenv' = case ty of
PackedTy _n l -> M.insert v l lenv
PackedTy _n l -> M.insert (singleLocVar v) l lenv
_ -> lenv

(outlocs,newls,eor') <- doBoundApp f lsin
e2' <- exp fns retlocs eor' lenv' afterenv (extendVEnv v ty env2) e2
e2' <- exp fns retlocs eor' lenv' afterenv (extendVEnvLocVar (singleLocVar v) ty env2) e2
return $ LetE (v,outlocs,ty, AppE f lsin e1)
(wrapBody e2' newls)

-- Exactly like AppE.
LetE (v,_ls,ty,(SpawnE f lsin e1)) e2 -> do
let lenv' = case ty of
PackedTy _n l -> M.insert v l lenv
PackedTy _n l -> M.insert (singleLocVar v) l lenv
_ -> lenv
(outlocs,newls,eor') <- doBoundApp f lsin
e2' <- exp fns retlocs eor' lenv' afterenv (extendVEnv v ty env2) e2
e2' <- exp fns retlocs eor' lenv' afterenv (extendVEnvLocVar (singleLocVar v) ty env2) e2
return $ LetE (v,outlocs,ty, SpawnE f lsin e1)
(wrapBody e2' newls)

Expand All @@ -319,18 +320,19 @@ routeEnds prg@Prog{ddefs,fundefs,mainExp} = do
forM brs $ \(dc, vls, e) ->
case vls of
[] ->
case (M.lookup x lenv) of
case (M.lookup (singleLocVar x) lenv) of
Just l1 -> do
l2 <- gensym "jump"
let eor' = mkEnd l1 l2 eor
e' = Ext $ LetLocE l2 (AfterConstantLE 1 l1) e
e'' <- exp fns retlocs eor' lenv (M.insert l1 l2 lenv) env2 e'
let l2loc = singleLocVar l2
eor' = mkEnd l1 l2loc eor
e' = Ext $ LetLocE l2loc (AfterConstantLE 1 l1) e
e'' <- exp fns retlocs eor' lenv (M.insert l1 l2loc lenv) env2 e'
return (dc, vls, e'')
Nothing -> error $ "Failed to find " ++ sdoc x ++ " in " ++ sdoc lenv
_ -> do
let need = snd $ last vls
argtys = lookupDataCon ddefs dc
lx = case M.lookup x lenv of
lx = case M.lookup (singleLocVar x) lenv of
Nothing -> error $ "Failed to find " ++ (show x)
Just l -> l
-- we know lx and need have the same end, since
Expand All @@ -347,14 +349,17 @@ routeEnds prg@Prog{ddefs,fundefs,mainExp} = do
-- bind a location to after it
handleLoc (eor,e) (l1,ty) = do
l2 <- gensym "jump"
let eor' = mkEnd l1 l2 eor
let l2loc = singleLocVar l2
eor' = mkEnd l1 l2loc eor
(Just jump) = L1.sizeOfTy ty
e' = Ext $ LetLocE l2 (AfterConstantLE jump l1) e
e' = Ext $ LetLocE l2loc (AfterConstantLE jump l1) e
return (eor', e')
vars = L.map fst vls
varsToLocs = L.map singleLocVar vars
locs = L.map snd vls
env2' = extendsVEnv (M.fromList (zip locs argtys ++ zip vars argtys)) env2
lenv' = M.union lenv $ M.fromList vls
vls' = L.map (\(v, l) -> (singleLocVar v, l)) vls
env2' = extendsVEnvLocVar (M.fromList (zip locs argtys ++ zip varsToLocs argtys)) env2
lenv' = M.union lenv $ M.fromList vls'
(eor'',e') <- foldM handleLoc (eor',e) $ zip (L.map snd vls) argtys
e'' <- exp fns retlocs eor'' lenv' afterenv' env2' e'
return (dc, vls, e'')
Expand All @@ -363,18 +368,18 @@ routeEnds prg@Prog{ddefs,fundefs,mainExp} = do


CaseE complex brs -> do
let ty = gRecoverType ddefs env2 complex
let ty = gRecoverTypeLoc ddefs env2 complex
v <- gensym "flt_RE"
let ex = L1.mkLets [(v,[],ty,complex)] (CaseE (VarE v) brs)
exp fns retlocs eor lenv afterenv (extendVEnv v ty env2) ex
exp fns retlocs eor lenv afterenv (extendVEnvLocVar (singleLocVar v) ty env2) ex


-- This shouldn't happen, but as a convenience we can ANF-ify this AppE
-- by gensyming a new variable, sticking the AppE in a LetE, and recuring.
-- Question: should this fail instead? I'm not sure.
AppE v args arg -> do
v' <- gensym "tailapp"
let ty = gRecoverType ddefs env2 e
let ty = gRecoverTypeLoc ddefs env2 e
e' = LetE (v',[], ty, AppE v args arg) (VarE v')
go (e')

Expand All @@ -397,41 +402,41 @@ routeEnds prg@Prog{ddefs,fundefs,mainExp} = do
-- Processing the RHS here would cause an infinite loop.

LetE (v,ls,ty@(PackedTy _ loc),e1@DataConE{}) e2 -> do
e2' <- exp fns retlocs eor (M.insert v loc lenv) afterenv (extendVEnv v ty env2) e2
e2' <- exp fns retlocs eor (M.insert (singleLocVar v) loc lenv) afterenv (extendVEnvLocVar (singleLocVar v) ty env2) e2
return $ LetE (v,ls,ty,e1) e2'

LetE (v,ls,ty@(PackedTy _ loc),e1@(PrimAppE (ReadPackedFile{}) [])) e2 -> do
e2' <- exp fns retlocs eor (M.insert v loc lenv) afterenv (extendVEnv v ty env2) e2
e2' <- exp fns retlocs eor (M.insert (singleLocVar v) loc lenv) afterenv (extendVEnvLocVar (singleLocVar v) ty env2) e2
return $ LetE (v,ls,ty,e1) e2'

LetE (v,ls,ty,e1@ProjE{}) e2 -> do
let lenv' = case ty of
PackedTy _ loc -> M.insert v loc lenv
PackedTy _ loc -> M.insert (singleLocVar v) loc lenv
_ -> lenv
e2' <- exp fns retlocs eor lenv' afterenv (extendVEnv v ty env2) e2
e2' <- exp fns retlocs eor lenv' afterenv (extendVEnvLocVar (singleLocVar v) ty env2) e2
return $ LetE (v,ls,ty,e1) e2'

LetE (v,ls,ty,e1@MkProdE{}) e2 -> do
LetE (v,ls,ty,e1) <$> exp fns retlocs eor lenv afterenv (extendVEnv v ty env2) e2
LetE (v,ls,ty,e1) <$> exp fns retlocs eor lenv afterenv (extendVEnvLocVar (singleLocVar v) ty env2) e2

LetE (v,ls,ty,e1@(PrimAppE (DictLookupP _) _)) e2 -> do
LetE (v,ls,ty,e1) <$> exp fns retlocs eor lenv afterenv (extendVEnv v ty env2) e2
LetE (v,ls,ty,e1) <$> exp fns retlocs eor lenv afterenv (extendVEnvLocVar (singleLocVar v) ty env2) e2

--

LetE (v,ls,ty@(PackedTy n l),e1) e2 -> do
e1' <- go e1
e2' <- exp fns retlocs eor (M.insert v l lenv) afterenv (extendVEnv v ty env2) e2
e2' <- exp fns retlocs eor (M.insert (singleLocVar v) l lenv) afterenv (extendVEnvLocVar (singleLocVar v) ty env2) e2
return $ LetE (v,ls,PackedTy n l,e1') e2'

LetE (v,ls,ty,e1@TimeIt{}) e2 -> do
e1' <- go e1
e2' <- exp fns retlocs eor lenv afterenv (extendVEnv v ty env2) e2
e2' <- exp fns retlocs eor lenv afterenv (extendVEnvLocVar (singleLocVar v) ty env2) e2
return $ LetE (v,ls,ty,e1') e2'

-- Most boring LetE case, just recur on body
LetE (v,ls,ty,rhs) bod -> do
bod' <- exp fns retlocs eor lenv afterenv (extendVEnv v ty env2) bod
bod' <- exp fns retlocs eor lenv afterenv (extendVEnvLocVar (singleLocVar v) ty env2) bod
return $ LetE (v,ls,ty,rhs) bod'

IfE e1 e2 e3 -> do
Expand All @@ -440,28 +445,28 @@ routeEnds prg@Prog{ddefs,fundefs,mainExp} = do
return $ IfE e1 e2' e3'

MkProdE ls -> do
let tys = L.map (gRecoverType ddefs env2) ls
let tys = L.map (gRecoverTypeLoc ddefs env2) ls
prodty = ProdTy tys
v <- gensym "flt_RE"
let ex = L1.mkLets [(v,[],prodty,(MkProdE ls))] (VarE v)
exp fns retlocs eor lenv afterenv (extendVEnv v prodty env2) ex
exp fns retlocs eor lenv afterenv (extendVEnvLocVar (singleLocVar v) prodty env2) ex

ProjE{} -> do
v <- gensym "flt_RE"
let ty = gRecoverType ddefs env2 e
let ty = gRecoverTypeLoc ddefs env2 e
lenv' = case ty of
PackedTy _ loc -> M.insert v loc lenv
PackedTy _ loc -> M.insert (singleLocVar v) loc lenv
_ -> lenv
ex = L1.mkLets [(v,[],ty,e)] (VarE v)
exp fns retlocs eor lenv' afterenv (extendVEnv v ty env2) ex
exp fns retlocs eor lenv' afterenv (extendVEnvLocVar (singleLocVar v) ty env2) ex

-- Could fail here, but try to fix the broken program
DataConE loc dc es -> do
v' <- gensym "taildc"
let ty = PackedTy (getTyOfDataCon ddefs dc) loc
e' = LetE (v',[],ty, DataConE loc dc es)
(VarE v')
exp fns retlocs eor (M.insert v' loc lenv) afterenv (extendVEnv v' ty env2) (e')
exp fns retlocs eor (M.insert (singleLocVar v') loc lenv) afterenv (extendVEnvLocVar (singleLocVar v') ty env2) (e')

LitE i -> return (LitE i)
CharE i -> return (CharE i)
Expand Down Expand Up @@ -547,7 +552,7 @@ routeEnds prg@Prog{ddefs,fundefs,mainExp} = do
wrapBody e [] = e

-- Process a let bound fn app.
doBoundApp :: Var -> [LocVar] -> PassM ([LocVar], [(LocVar, Var)], EndOfRel)
doBoundApp :: Var -> [LocVar] -> PassM ([LocVar], [(LocVar, LocVar)], EndOfRel)
doBoundApp f lsin = do
let fty = funtype f
rets = S.fromList $ locRets fty
Expand All @@ -559,7 +564,7 @@ routeEnds prg@Prog{ddefs,fundefs,mainExp} = do
-- For each traversed location, gensym a new variable for its end,
-- and generate a list of (location, endof location) pairs.
let handleTravList lst (_l,False) = return lst
handleTravList lst (l,True) = gensym "endof" >>= \l' -> return $ (l,l'):lst
handleTravList lst (l,True) = gensym "endof" >>= \l' -> return $ (l, (singleLocVar l')):lst

-- Walk through our pairs of (location, endof location) and update the
-- endof relation.
Expand Down

0 comments on commit 2c525ea

Please sign in to comment.