Change LocVar to loc in L2 IR for LetLocE, addtag and allocate scalar…

… for additional location metadata
iu-parfunc · Mar 5, 2024 · e7eab10 · e7eab10
1 parent 87693a1
commit e7eab10
Show file tree

Hide file tree

Showing 7 changed files with 52 additions and 47 deletions.
diff --git a/gibbon-compiler/src/Gibbon/Compiler.hs b/gibbon-compiler/src/Gibbon/Compiler.hs
@@ -799,7 +799,7 @@ Also see Note [Adding dummy traversals] and Note [Adding random access nodes].
               -- it adds regions to 'locs' in AppE and LetE which the
               -- typechecker doesn't know how to handle.
               l2' <- go "threadRegions"    threadRegions l2'
-              --l2' <- go "markTailCalls" markTailCalls l2'
+              l2' <- go "markTailCalls" markTailCalls l2'
               -- L2 -> L3
               -- TODO: Compose L3.TcM with (ReaderT Config)
               l3 <- go "cursorize"        cursorize     l2'

diff --git a/gibbon-compiler/src/Gibbon/L2/Syntax.hs b/gibbon-compiler/src/Gibbon/L2/Syntax.hs
@@ -144,7 +144,7 @@ instance Monoid RegionSize where
 data E2Ext loc dec
   = LetRegionE    Region RegionSize (Maybe RegionType) (E2 loc dec) -- ^ Allocate a new region.
   | LetParRegionE Region RegionSize (Maybe RegionType) (E2 loc dec) -- ^ Allocate a new region for parallel allocations.
-  | LetLocE LocVar (PreLocExp loc) (E2 loc dec) -- ^ Bind a new location.
+  | LetLocE loc (PreLocExp loc) (E2 loc dec) -- ^ Bind a new location.
   | RetE [loc] Var          -- ^ Return a value together with extra loc values.
   | FromEndE loc            -- ^ Bind a location from an EndOf location (for RouteEnds and after).
   | BoundsCheck Int -- Bytes required
@@ -168,8 +168,8 @@ data E2Ext loc dec
   | GetCilkWorkerNum
     -- ^ Translates to  __cilkrts_get_worker_number().
   | LetAvail [Var] (E2 loc dec) -- ^ These variables are available to use before the join point.
-  | AllocateTagHere LocVar TyCon
-  | AllocateScalarsHere LocVar
+  | AllocateTagHere loc TyCon
+  | AllocateScalarsHere loc
     -- ^ A marker which tells subsequent a compiler pass where to
     -- move the tag and scalar field allocations so that they happen
     -- before any of the subsequent packed fields.

diff --git a/gibbon-compiler/src/Gibbon/NewL2/FromOldL2.hs b/gibbon-compiler/src/Gibbon/NewL2/FromOldL2.hs
@@ -164,7 +164,7 @@ fromOldL2Exp ddefs fundefs locenv env2 ex =
           let rhs' = fmap (locenv #) rhs
               locarg = toLocArg loc rhs locenv
           bod' <- go (M.insert loc locarg locenv) env2 bod
-          pure $ Ext $ LetLocE loc rhs' bod'
+          pure $ Ext $ LetLocE locarg rhs' bod'
 
         RetE locs v -> do
           let locargs = map (locenv #) locs
@@ -202,10 +202,14 @@ fromOldL2Exp ddefs fundefs locenv env2 ex =
           pure $ Ext $ LetAvail avail rhs'
 
         AllocateTagHere loc tycon -> do
-          -- let locarg = locenv # loc
-          pure $ Ext $ AllocateTagHere loc tycon
+          case (M.lookup loc locenv) of 
+                Nothing -> error $ "Could not find a locarg for location variable: " ++ sdoc loc
+                Just locarg -> pure $ Ext $ AllocateTagHere locarg tycon
 
-        AllocateScalarsHere loc -> pure $ Ext $ AllocateScalarsHere loc
+        AllocateScalarsHere loc -> do
+            case (M.lookup loc locenv) of 
+                    Nothing -> error $ "Could not find a locarg for location variable: " ++ sdoc loc
+                    Just locarg -> pure $ Ext $ AllocateScalarsHere locarg
 
         SSPush mode loc end_loc tycon -> do
           pure $ Ext $ SSPush mode loc end_loc tycon
@@ -341,7 +345,7 @@ toOldL2Exp ex =
         LetLocE loc rhs bod -> do
           let rhs' = fmap New.toLocVar rhs
           bod' <- go bod
-          pure $ Ext $ LetLocE loc rhs' bod'
+          pure $ Ext $ LetLocE (New.toLocVar loc) rhs' bod'
 
         StartOfPkdCursor cur -> do
           pure $ Ext $ StartOfPkdCursor cur
@@ -377,9 +381,9 @@ toOldL2Exp ex =
           pure $ Ext $ LetAvail avail rhs'
 
         AllocateTagHere loc tycon -> do
-          pure $ Ext $ AllocateTagHere loc tycon
+          pure $ Ext $ AllocateTagHere (New.toLocVar loc) tycon
 
-        AllocateScalarsHere loc -> pure $ Ext $ AllocateScalarsHere loc
+        AllocateScalarsHere loc -> pure $ Ext $ AllocateScalarsHere (New.toLocVar loc)
 
         SSPush mode loc end_loc tycon -> do
           pure $ Ext $ SSPush mode loc end_loc tycon

diff --git a/gibbon-compiler/src/Gibbon/NewL2/Syntax.hs b/gibbon-compiler/src/Gibbon/NewL2/Syntax.hs
@@ -392,7 +392,7 @@ depList = L.map (\(a,b) -> (a,a,b)) . M.toList . go M.empty
               Old.LetParRegionE r _ _ rhs ->
                 go (M.insertWith (++) (Old.regionToVar r) (S.toList $ allFreeVars rhs) acc) rhs
               Old.LetLocE loc phs rhs  ->
-                go (M.insertWith (++) loc (dep phs ++ (S.toList $ allFreeVars rhs)) acc) rhs
+                go (M.insertWith (++) (toLocVar loc) (dep phs ++ (S.toList $ allFreeVars rhs)) acc) rhs
               Old.RetE{}         -> acc
               Old.FromEndE{}     -> acc
               Old.BoundsCheck{}  -> acc
@@ -440,7 +440,7 @@ allFreeVars ex =
       case ext of
         Old.LetRegionE r _ _ bod -> S.delete (Old.regionToVar r) (allFreeVars bod)
         Old.LetParRegionE r _ _ bod -> S.delete (Old.regionToVar r) (allFreeVars bod)
-        Old.LetLocE loc locexp bod -> S.delete loc (allFreeVars bod `S.union` gFreeVars locexp)
+        Old.LetLocE loc locexp bod -> S.delete (toLocVar loc) (allFreeVars bod `S.union` gFreeVars locexp)
         Old.StartOfPkdCursor v -> S.singleton v
         Old.TagCursor a b-> S.fromList [a,b]
         Old.RetE locs v     -> S.insert v (S.fromList (map toLocVar locs))
@@ -450,8 +450,8 @@ allFreeVars ex =
         Old.AddFixed v _    -> S.singleton v
         Old.GetCilkWorkerNum-> S.empty
         Old.LetAvail vs bod -> S.fromList vs `S.union` gFreeVars bod
-        Old.AllocateTagHere loc _ -> S.singleton loc
-        Old.AllocateScalarsHere loc -> S.singleton loc
+        Old.AllocateTagHere loc _ -> S.singleton (toLocVar loc)
+        Old.AllocateScalarsHere loc -> S.singleton (toLocVar loc)
         Old.SSPush _ a b _ -> S.fromList [a,b]
         Old.SSPop _ a b -> S.fromList [a,b]
     _ -> gFreeVars ex

diff --git a/gibbon-compiler/src/Gibbon/Passes/Cursorize.hs b/gibbon-compiler/src/Gibbon/Passes/Cursorize.hs
@@ -362,12 +362,12 @@ cursorizeExp ddfs fundefs denv tenv senv ex =
         -- is expressed in terms of corresponding cursor operations.
         -- See `cursorizeLocExp`
         LetLocE loc rhs bod -> do
-          let (rhs_either, lTy) = cursorizeLocExp denv tenv senv loc rhs
+          let (rhs_either, lTy) = cursorizeLocExp denv tenv senv (toLocVar loc) rhs
               lTy' = case lTy of 
                         MkTy2{unTy2} -> case unTy2 of 
                                               CursorTy -> CursorTy
                                               MutableCursorTy -> MutableCursorTy
-              (bnds,tenv') = case M.lookup loc denv of
+              (bnds,tenv') = case M.lookup (toLocVar loc) denv of
                                Nothing -> ([],tenv)
                                Just vs -> let extended = M.fromList [ (v,MkTy2 CursorTy) | (v,_,CursorTy,_) <- vs]
                                           in (vs, M.union extended tenv)
@@ -389,14 +389,14 @@ cursorizeExp ddfs fundefs denv tenv senv ex =
               let tenv''' = M.union tenv' tenv''
               case rhs of
                 FromEndLE{} ->
-                  if isBound loc tenv
-                  then cursorizeExp ddfs fundefs denv (M.insert loc (MkTy2 CursorTy) tenv''') senv' bod
+                  if isBound (toLocVar loc) tenv
+                  then cursorizeExp ddfs fundefs denv (M.insert (toLocVar loc) (MkTy2 CursorTy) tenv''') senv' bod
                   -- Discharge bindings that were waiting on 'loc'.
-                  else mkLets (bnds' ++ [(loc,[],lTy',rhs')] ++ bnds) <$>
-                         cursorizeExp ddfs fundefs denv (M.insert loc (MkTy2 CursorTy) tenv''') senv' bod
+                  else mkLets (bnds' ++ [((toLocVar loc),[],lTy',rhs')] ++ bnds) <$>
+                         cursorizeExp ddfs fundefs denv (M.insert (toLocVar loc) (MkTy2 CursorTy) tenv''') senv' bod
                 -- Discharge bindings that were waiting on 'loc'.
-                _ -> mkLets (bnds' ++ [(loc,[],lTy',rhs')] ++ bnds) <$>
-                       cursorizeExp ddfs fundefs denv (M.insert loc (MkTy2 CursorTy) tenv''') senv bod
+                _ -> mkLets (bnds' ++ [((toLocVar loc),[],lTy',rhs')] ++ bnds) <$>
+                       cursorizeExp ddfs fundefs denv (M.insert (toLocVar loc) (MkTy2 CursorTy) tenv''') senv bod
             Left denv' -> cursorizeExp ddfs fundefs denv' tenv' senv bod
 
         -- Exactly same as cursorizePackedExp
@@ -416,9 +416,9 @@ cursorizeExp ddfs fundefs denv tenv senv ex =
 
         LetAvail vs bod  -> Ext <$> L3.LetAvail vs <$> go bod
 
-        AllocateTagHere v tycon -> pure $ Ext $ L3.AllocateTagHere v tycon
+        AllocateTagHere v tycon -> pure $ Ext $ L3.AllocateTagHere (toLocVar v) tycon
 
-        AllocateScalarsHere v -> pure $ Ext $ L3.AllocateScalarsHere v
+        AllocateScalarsHere v -> pure $ Ext $ L3.AllocateScalarsHere (toLocVar v)
 
         SSPush a b c d -> pure $ Ext $ L3.SSPush a b c d
         SSPop a b c -> pure $ Ext $ L3.SSPop a b c
@@ -628,12 +628,12 @@ cursorizePackedExp ddfs fundefs denv tenv senv ex =
         -- is expressed in terms of corresponding cursor operations.
         -- See `cursorizeLocExp`
         LetLocE loc rhs bod -> do
-          let (rhs_either, lTy) = cursorizeLocExp denv tenv senv loc rhs
+          let (rhs_either, lTy) = cursorizeLocExp denv tenv senv (toLocVar loc) rhs
               lTy' = case lTy of 
                         MkTy2{unTy2} -> case unTy2 of 
                                               CursorTy -> CursorTy
                                               MutableCursorTy -> MutableCursorTy 
-              (bnds,tenv') = case M.lookup loc denv of
+              (bnds,tenv') = case M.lookup (toLocVar loc) denv of
                                Nothing -> ([],tenv)
                                Just vs -> let extended = M.fromList [ (v, MkTy2 CursorTy) | (v,_,CursorTy,_) <- vs]
                                           in (vs, M.union extended tenv)
@@ -642,14 +642,14 @@ cursorizePackedExp ddfs fundefs denv tenv senv ex =
               let tenv''' = M.union tenv' tenv''
               case rhs of
                 FromEndLE{} ->
-                  if isBound loc tenv
-                  then go (M.insert loc (MkTy2 CursorTy) tenv''') senv' bod
+                  if isBound (toLocVar loc) tenv
+                  then go (M.insert (toLocVar loc) (MkTy2 CursorTy) tenv''') senv' bod
                     -- Discharge bindings that were waiting on 'loc'.
-                  else onDi (mkLets (bnds' ++ [(loc,[],lTy',rhs')] ++ bnds)) <$>
-                         go (M.insert loc (MkTy2 CursorTy) tenv') senv' bod
+                  else onDi (mkLets (bnds' ++ [((toLocVar loc),[],lTy',rhs')] ++ bnds)) <$>
+                         go (M.insert (toLocVar loc) (MkTy2 CursorTy) tenv') senv' bod
                 -- Discharge bindings that were waiting on 'loc'.
-                _ -> onDi (mkLets (bnds' ++ [(loc,[],lTy',rhs')] ++ bnds)) <$>
-                       go (M.insert loc (MkTy2 CursorTy) tenv''') senv' bod
+                _ -> onDi (mkLets (bnds' ++ [((toLocVar loc),[],lTy',rhs')] ++ bnds)) <$>
+                       go (M.insert (toLocVar loc) (MkTy2 CursorTy) tenv''') senv' bod
             Left denv' -> onDi (mkLets bnds) <$>
                             cursorizePackedExp ddfs fundefs denv' tenv' senv bod
 
@@ -703,9 +703,9 @@ cursorizePackedExp ddfs fundefs denv tenv senv ex =
         LetAvail vs bod  -> do
           onDi (Ext . L3.LetAvail vs) <$> go tenv senv bod
 
-        AllocateTagHere v tycon -> pure <$> dl <$> Ext $ L3.AllocateTagHere v tycon
+        AllocateTagHere v tycon -> pure <$> dl <$> Ext $ L3.AllocateTagHere (toLocVar v) tycon
 
-        AllocateScalarsHere v -> pure <$> dl <$> Ext $ L3.AllocateScalarsHere v
+        AllocateScalarsHere v -> pure <$> dl <$> Ext $ L3.AllocateScalarsHere (toLocVar v)
 
         SSPush a b c d -> pure <$> dl <$> Ext $ L3.SSPush a b c d
         SSPop a b c -> pure <$> dl <$> Ext $ L3.SSPop a b c

diff --git a/gibbon-compiler/src/Gibbon/Passes/MarkTailCalls.hs b/gibbon-compiler/src/Gibbon/Passes/MarkTailCalls.hs
@@ -205,7 +205,7 @@ markTailCallsFnBody funName env exp2 = case exp2 of
                 let locInExp = freeLoc locexp
                     env' = case locInExp of
                         Nothing -> env
-                        Just l -> M.insert l (S.singleton loc, False) env
+                        Just l -> M.insert l (S.singleton (toLocVar loc), False) env
                     (bod', env'') = markTailCallsFnBody funName env' bod
                     locexp' = case locInExp of
                         Nothing -> locexp
@@ -488,7 +488,8 @@ copyOutputMutableBeforeCallsAndReplace exp = case exp of
                                         new_loc <- gensym "loc"
                                         let locexp = AfterConstantLE 0 l
                                         let map' = M.insert ll new_loc map
-                                        return $ (lst ++ [NewL2.LetLocE new_loc locexp (VarE new_loc)], map')
+                                        let new_loc_arg = Loc (LREM new_loc a b m') -- TODO: should this be outputMutable or just Output
+                                        return $ (lst ++ [NewL2.LetLocE new_loc_arg locexp (VarE new_loc)], map')
                                 )
                                 ([], M.empty)
                                 outputMutableLocs

diff --git a/gibbon-compiler/src/Gibbon/Passes/ThreadRegions.hs b/gibbon-compiler/src/Gibbon/Passes/ThreadRegions.hs
@@ -416,8 +416,8 @@ threadRegionsExp ddefs fundefs fnLocArgs renv env2 lfenv rlocs_env wlocs_env pkd
 
     LetE (v,locs,ty,rhs@(Ext (AllocateTagHere x x_tycon))) bod -> do
       let -- x_tycon = (wlocs_env # x)
-          rlocs_env' = M.insert x x_tycon rlocs_env
-          wlocs_env' = M.delete x wlocs_env
+          rlocs_env' = M.insert (NewL2.toLocVar x) x_tycon rlocs_env
+          wlocs_env' = M.delete (NewL2.toLocVar x) wlocs_env
       (LetE (v,locs,ty,rhs)) <$>
         threadRegionsExp ddefs fundefs fnLocArgs renv (extendVEnv v ty env2) lfenv rlocs_env' wlocs_env' pkd_env region_locs ran_env indirs redirs bod
 
@@ -442,14 +442,14 @@ threadRegionsExp ddefs fundefs fnLocArgs renv env2 lfenv rlocs_env wlocs_env pkd
                       AfterConstantLE _ lc   -> renv # (toLocVar lc)
                       AfterVariableLE _ lc _ -> renv # (toLocVar lc)
                       FromEndLE lc           -> renv # (toLocVar lc)
-              wlocs_env' = M.insert loc hole_tycon wlocs_env
+              wlocs_env' = M.insert (toLocVar loc) hole_tycon wlocs_env
               region_locs1 = case rhs of
-                               AfterConstantLE{} -> M.adjust (\locs -> locs ++ [loc]) reg region_locs
-                               AfterVariableLE{} -> M.adjust (\locs -> locs ++ [loc]) reg region_locs
-                               StartOfRegionLE{} -> M.insert reg [loc] region_locs
+                               AfterConstantLE{} -> M.adjust (\locs -> locs ++ [(toLocVar loc)]) reg region_locs
+                               AfterVariableLE{} -> M.adjust (\locs -> locs ++ [(toLocVar loc)]) reg region_locs
+                               StartOfRegionLE{} -> M.insert reg [(toLocVar loc)] region_locs
                                _ -> region_locs
           Ext <$> LetLocE loc rhs <$>
-            threadRegionsExp ddefs fundefs fnLocArgs (M.insert loc reg renv) env2 lfenv rlocs_env wlocs_env' pkd_env region_locs1 ran_env indirs redirs bod
+            threadRegionsExp ddefs fundefs fnLocArgs (M.insert (toLocVar loc) reg renv) env2 lfenv rlocs_env wlocs_env' pkd_env region_locs1 ran_env indirs redirs bod
 
         RetE locs v -> do
           let ty = lookupVEnv v env2
@@ -723,7 +723,7 @@ allFreeVars_sans_datacon_args ex =
       case ext of
         LetRegionE r _sz _ty bod -> S.delete (regionToVar r) (allFreeVars_sans_datacon_args bod)
         LetParRegionE r _sz _ty bod -> S.delete (regionToVar r) (allFreeVars_sans_datacon_args bod)
-        LetLocE loc locexp bod -> S.delete loc (allFreeVars_sans_datacon_args bod `S.union` gFreeVars locexp)
+        LetLocE loc locexp bod -> S.delete (toLocVar loc) (allFreeVars_sans_datacon_args bod `S.union` gFreeVars locexp)
         StartOfPkdCursor cur   -> S.singleton cur
         TagCursor a b-> S.fromList [a,b]
         RetE locs v     -> S.insert v (S.fromList (map toLocVar locs))
@@ -733,8 +733,8 @@ allFreeVars_sans_datacon_args ex =
         AddFixed v _    -> S.singleton v
         GetCilkWorkerNum-> S.empty
         LetAvail vs bod -> S.fromList vs `S.union` gFreeVars bod
-        AllocateTagHere loc _ -> S.singleton loc
-        AllocateScalarsHere loc -> S.singleton loc
+        AllocateTagHere loc _ -> S.singleton (toLocVar loc)
+        AllocateScalarsHere loc -> S.singleton (toLocVar loc)
         SSPush _ a b _ -> S.fromList [a,b]
         SSPop _ a b -> S.fromList [a,b]
     _ -> gFreeVars ex