Skip to content

Commit 31a02c6

Browse files
committed
wip nested soa
1 parent 36302e2 commit 31a02c6

File tree

3 files changed

+96
-19
lines changed

3 files changed

+96
-19
lines changed

gibbon-compiler/src/Gibbon/Common.hs

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,8 @@ module Gibbon.Common
3333

3434
-- * Debugging/logging:
3535
, dbgLvl, dbgPrint, dbgPrintLn, dbgTrace, dbgTraceIt, minChatLvl
36-
, internalError, dumpIfSet, unwrapLocVar, singleLocVar, getDconLoc, getFieldLoc, freshCommonLoc, getAllFieldLocsSoA, varsInLocVar, varsInRegVar
36+
, internalError, dumpIfSet, unwrapLocVar, singleLocVar, getDconLoc, getFieldLoc, freshCommonLoc, getAllFieldLocsSoA
37+
, varsInLocVar, varsInRegVar, getAllFieldRegsSoA
3738
, appendNameToLocVar
3839

3940
-- * Establish conventions for the output of #lang gibbon:
@@ -580,6 +581,12 @@ getAllFieldLocsSoA loc = case loc of
580581
SoA _dcon fieldLocs -> fieldLocs
581582
Single _lc -> error "getFieldLocs : Did not expect a non SoA location!"
582583

584+
585+
getAllFieldRegsSoA :: RegVar -> [((DataCon, Int), RegVar)]
586+
getAllFieldRegsSoA loc = case loc of
587+
SoARv _dcon fieldLocs -> fieldLocs
588+
SingleR _lc -> error "getFieldLocs : Did not expect a non SoA location!"
589+
583590
freshSingleLocVar :: String -> PassM LocVar
584591
freshSingleLocVar m = do v <- gensym (toVar m)
585592
return $ Single v

gibbon-compiler/src/Gibbon/L3/Syntax.hs

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,9 @@ module Gibbon.L3.Syntax
2323
, getCursorizeTyFromRegVar''
2424
, getCursorizeTyFromRegVar'''
2525
, getIndexPositionOfSoALocVar
26+
, getIndexPositionOfSoARegVar
2627
, linearizeLocVar
28+
, linearizeRegVar
2729
, module Gibbon.Language
2830
)
2931
where
@@ -319,12 +321,31 @@ getIndexPositionOfSoALocVar flds loc = foldl (\(s, e, b) (_, fl) -> if b
319321
in (e, e + sz, seen)
320322
) (1, 1, False) flds
321323

324+
getIndexPositionOfSoARegVar :: [((DataCon, Int), RegVar)] -> RegVar -> (Int, Int, Bool)
325+
getIndexPositionOfSoARegVar flds loc = foldl (\(s, e, b) (_, fl) -> if b
326+
then
327+
(s, e, True)
328+
else
329+
let seen = if fl == loc then True else False
330+
in case fl of
331+
SingleR{} -> (e, e + 1, seen)
332+
SoARv{} -> let (CursorArrayTy sz) = getCursorizeTyFromRegVar fl
333+
in (e, e + sz, seen)
334+
) (1, 1, False) flds
335+
322336
linearizeLocVar :: LocVar -> [LocVar]
323337
linearizeLocVar loc = case loc of
324338
Single{} -> [loc]
325339
SoA dcloc flocs -> let flinear = concatMap (\(_, fl) -> linearizeLocVar fl) flocs
326340
in [singleLocVar dcloc] ++ flinear
327341

342+
343+
linearizeRegVar :: RegVar -> [RegVar]
344+
linearizeRegVar loc = case loc of
345+
SingleR{} -> [loc]
346+
SoARv dcloc flocs -> let flinear = concatMap (\(_, fl) -> linearizeRegVar fl) flocs
347+
in [dcloc] ++ flinear
348+
328349
getCursorizeTyFromLocVar :: LocVar -> Ty3
329350
getCursorizeTyFromLocVar lc = case lc of
330351
Single{} -> CursorTy

gibbon-compiler/src/Gibbon/Passes/Cursorize.hs

Lines changed: 67 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -152,29 +152,57 @@ mangle vars = toVar $ "mangle" ++ (L.foldr (\v acc -> acc ++ "_" ++ (fromVar v))
152152

153153
-- The LocVar here is the field location, which we need to generate code for.
154154
-- (Int, Int) is the start and end locations of that field.
155-
handleIndexingSoACursors :: (Int, Int) -> LocVar -> M.Map FreeVarsTy Var -> PassM (M.Map FreeVarsTy Var, [(Var, [()], Ty3, Exp3)])
156-
handleIndexingSoACursors (start, end) locvar var_env = do
155+
handleIndexingSoACursors :: (LocVar, Var) -> (Int, Int) -> LocVar -> M.Map FreeVarsTy Var -> PassM (M.Map FreeVarsTy Var, [(Var, [()], Ty3, Exp3)])
156+
handleIndexingSoACursors (arrLoc, arrName) (start, end) locvar var_env = do
157157
let par_var = case (M.lookup (fromLocVarToFreeVarsTy locvar) var_env) of
158158
Just v -> v
159159
Nothing -> case locvar of
160160
Single l -> l
161161
SoA{} -> error "Expected variable name for parent array!"
162-
case locvar of
162+
case arrLoc of
163163
Single{} -> do
164-
return (var_env, [(par_var, [], CursorTy, Ext $ IndexCursorArray par_var start)])
164+
return (var_env, [(arrName, [], CursorTy, Ext $ IndexCursorArray par_var start)])
165165
SoA{} -> do
166-
(bnds, var_env') <- foldlM (\(b, env) (i, l) -> do
166+
let linearized_locs = (linearizeLocVar locvar)
167+
(vars, bnds, var_env') <- foldlM (\(v, b, env) (i, l) -> do
167168
(lvar, fenv') <- case (M.lookup (fromLocVarToFreeVarsTy l) var_env) of
168169
Just v -> return (v, env)
169170
Nothing -> do
170171
new_var <- gensym "unpack"
171172
let env' = M.insert (fromLocVarToFreeVarsTy l) new_var env
172173
return (new_var, env')
173-
pure $ (b ++ [(lvar, [], CursorTy, Ext $ IndexCursorArray par_var i)], fenv')
174+
pure $ (v ++ [lvar], b ++ [(lvar, [], CursorTy, Ext $ IndexCursorArray par_var i)], fenv')
174175

175176

176-
) ([], var_env) (zip [start..end] (linearizeLocVar locvar))
177-
return (var_env, bnds)
177+
) ([], [], var_env) (zip [start..end] (take (end - start) (drop start linearized_locs)) )
178+
let make_cur_arr_let = [(arrName, [], getCursorizeTyFromLocVar arrLoc, Ext $ MakeCursorArray (length vars) vars)]
179+
return (var_env, bnds ++ make_cur_arr_let)
180+
181+
handleIndexingSoARegCursors :: (RegVar, Var) -> (Int, Int) -> RegVar -> M.Map FreeVarsTy Var -> PassM (M.Map FreeVarsTy Var, [(Var, [()], UrTy (), (PreExp E3Ext () (UrTy ())))])
182+
handleIndexingSoARegCursors (arrLoc, arrName) (start, end) locvar var_env = do
183+
let par_var = case (M.lookup (fromRegVarToFreeVarsTy locvar) var_env) of
184+
Just v -> v
185+
Nothing -> case locvar of
186+
SingleR l -> l
187+
SoARv{} -> error "Expected variable name for parent array!"
188+
case arrLoc of
189+
SingleR{} -> do
190+
return (var_env, [(arrName, [], CursorTy, Ext $ IndexCursorArray par_var start)])
191+
SoARv{} -> do
192+
let linearized_locs = (linearizeRegVar locvar)
193+
(vars, bnds, var_env') <- foldlM (\(v, b, env) (i, l) -> do
194+
(lvar, fenv') <- case (M.lookup (fromRegVarToFreeVarsTy l) var_env) of
195+
Just v -> return (v, env)
196+
Nothing -> do
197+
new_var <- gensym "unpack"
198+
let env' = M.insert (fromRegVarToFreeVarsTy l) new_var env
199+
return (new_var, env')
200+
pure $ (v ++ [lvar], b ++ [(lvar, [], CursorTy, Ext $ IndexCursorArray par_var i)], fenv')
201+
202+
203+
) ([], [], var_env) (zip [start..end] (take (end - start) (drop start linearized_locs)) )
204+
let make_cur_arr_let = [(arrName, [], getCursorizeTyFromRegVar''' arrLoc, Ext $ MakeCursorArray (length vars) vars)]
205+
return (var_env, bnds ++ make_cur_arr_let)
178206

179207

180208
cursorizeFunDef :: Bool -> DDefs Ty2 -> FunDefs2 -> FunDef2 -> PassM FunDef3
@@ -2562,8 +2590,13 @@ cursorizeAppE freeVarToVarEnv lenv ddfs fundefs denv tenv senv ex =
25622590
Just v -> v
25632591
Nothing -> error $ "cursorizeAppE: Did not find an end of region variable for the corresponding parent region.\n\n" ++ show f ++ "\n\n " ++ show r ++ "\n\n " ++ show acc
25642592
name <- gensym "cursor_reg_ptr"
2565-
let instrs = [LetE (name, [], getCursorizeTyFromRegVar r, Ext $ IndexCursorArray (name_par_reg) 1)]
2566-
return $ (M.insert loc_var name acc, acc' ++ instrs)
2593+
let (R par_reg_inner) = par_reg
2594+
-- Vidush: TODO, is this right?
2595+
let (start, end, _) = getIndexPositionOfSoARegVar (getAllFieldRegsSoA par_reg_inner) r
2596+
(_acc, instrs) <- handleIndexingSoARegCursors (r, name) (start, end) par_reg_inner acc
2597+
--let instrs = [LetE (name, [], getCursorizeTyFromRegVar r, Ext $ IndexCursorArray (name_par_reg) 1)]
2598+
let instrs' = map (\i -> LetE i) instrs
2599+
return $ (M.insert loc_var name _acc, acc' ++ instrs')
25672600
Nothing -> do
25682601
(dconReg_var, dcon_insts) <- case (M.lookup (fromRegVarToFreeVarsTy dconReg) acc) of
25692602
Just v -> return (v, [])
@@ -2583,15 +2616,29 @@ cursorizeAppE freeVarToVarEnv lenv ddfs fundefs denv tenv senv ex =
25832616
return (name_dcon, instrs)
25842617

25852618
-- Nothing -> error $ "cursorizeAppE: Did not find an end of region variable for the corresponding datacon region.\n\n" ++ show f ++ "\n\n " ++ show r ++ "\n\n " ++ show acc
2586-
let fieldReg_vars =
2587-
map
2588-
( \(key, field_reg) -> case (M.lookup (fromRegVarToFreeVarsTy field_reg) acc) of
2589-
Just v -> v
2590-
Nothing -> error "cursorizeAppE: Did not find an end of region variable for the corresponding field region.\n"
2619+
(fieldReg_vars, bnds) <-
2620+
foldlM
2621+
(\(vs, bds) (key, field_reg) -> do
2622+
v <- case (M.lookup (fromRegVarToFreeVarsTy field_reg) acc) of
2623+
Just vv -> return vv
2624+
Nothing -> error "cursorizeAppE: Did not find an end of region variable for the corresponding field region.\n"
2625+
case field_reg of
2626+
SingleR{} -> do
2627+
let (idx, _, _) = getIndexPositionOfSoARegVar fieldRegions field_reg
2628+
pure (vs ++ [v], bds ++ [(v, [], CursorTy, Ext $ IndexCursorArray v idx)])
2629+
SoARv{} -> do
2630+
let (start, end, _) = getIndexPositionOfSoARegVar fieldRegions field_reg
2631+
(nvars, bnds) <- foldlM (\(nv, bnd) i -> do
2632+
var_n <- gensym "unpack"
2633+
return (nv ++ [var_n], bnd ++ [(var_n, [], CursorTy, Ext $ IndexCursorArray v i)])
2634+
2635+
) ([], []) [start ..(end - 1)]
2636+
pure (vs ++ nvars, bds ++ bnds)
25912637
)
2638+
([], [])
25922639
fieldRegions
25932640
name <- gensym "cursor_reg_ptr"
2594-
let instrs = dcon_insts ++ [LetE (name, [], getCursorizeTyFromRegVar r, Ext $ MakeCursorArray (1 + length fieldReg_vars) ([dconReg_var] ++ fieldReg_vars))]
2641+
let instrs = dcon_insts ++ (map (\i -> LetE i) bnds) ++ [LetE (name, [], getCursorizeTyFromRegVar r, Ext $ MakeCursorArray (1 + length fieldReg_vars) ([dconReg_var] ++ fieldReg_vars))]
25952642
dbgTrace (minChatLvl) "Print Reg: " dbgTrace (minChatLvl) (sdoc (f, dconReg, fieldRegions)) dbgTrace (minChatLvl) "End soa Reg\n" return $ (M.insert loc_var name acc, acc' ++ instrs)
25962643
pure ret
25972644

@@ -3275,9 +3322,11 @@ unpackDataCon dcon_var freeVarToVarEnv lenv ddfs fundefs denv1 tenv1 senv isPack
32753322
field_var <- gensym $ toVar $ (fromVar "soa_field_") ++ (show idx_elem)
32763323
let acc3' = dbgTrace (minChatLvl) "print loc: " dbgTrace (minChatLvl) (sdoc (loc, scrut_loc)) dbgTrace (minChatLvl) "End cursorize print loc.\n" M.insert (fromLocVarToFreeVarsTy loc) field_var acc3
32773324
let field_cursor_ty = getCursorizeTyFromLocVar loc
3278-
let field_let = [(field_var, [], field_cursor_ty, Ext $ IndexCursorArray scrtCur (1 + idx_elem))]
3325+
let (start, end, _) = getIndexPositionOfSoALocVar (getAllFieldLocsSoA scrut_loc) loc
3326+
(acc3'', field_let) <- handleIndexingSoACursors (loc, field_var) (start, end) scrut_loc acc3'
3327+
--let field_let = [(field_var, [], field_cursor_ty, Ext $ IndexCursorArray scrtCur (1 + idx_elem))]
32793328
let curr_window = [((dcon', idx), field_var)]
3280-
return (acc1 ++ field_let, acc2 ++ curr_window, acc3')
3329+
return (acc1 ++ field_let, acc2 ++ curr_window, acc3'')
32813330
)
32823331
([], [], freeVarToVarEnv)
32833332
(getAllFieldLocsSoA scrut_loc)

0 commit comments

Comments
 (0)