Use Type Based Alias Analysis (TBAA) in LLVM backend (#5567)

TBAA allows us to specify a type hierachy in metadata with
the property that nodes on different branches don't alias.
This should somewhat improve the optimizations LLVM does that
rely on alias information.

3 files changed, 84 insertions, 34 deletions

diff --git a/compiler/llvmGen/LlvmCodeGen/CodeGen.hs b/compiler/llvmGen/LlvmCodeGen/CodeGen.hs
index 07ccbb1348..4309dcdae1 100644
--- a/compiler/llvmGen/LlvmCodeGen/CodeGen.hs
+++ b/compiler/llvmGen/LlvmCodeGen/CodeGen.hs
@@ -550,7 +550,7 @@ genStore env addr@(CmmMachOp (MO_Sub _) [
     = genStore_fast env addr r (negate $ fromInteger n) val
 
 -- generic case
-genStore env addr val = genStore_slow env addr val
+genStore env addr val = genStore_slow env addr val [top]
 
 -- | CmmStore operation
 -- This is a special case for storing to a global register pointer
@@ -558,8 +558,9 @@ genStore env addr val = genStore_slow env addr val
 genStore_fast :: LlvmEnv -> CmmExpr -> GlobalReg -> Int -> CmmExpr
               -> UniqSM StmtData
 genStore_fast env addr r n val
-  = let gr  = lmGlobalRegVar r
-        grt = (pLower . getVarType) gr
+  = let gr   = lmGlobalRegVar r
+        meta = [getTBAA r]
+        grt  = (pLower . getVarType) gr
         (ix,rem) = n `divMod` ((llvmWidthInBits . pLower) grt  `div` 8)
     in case isPointer grt && rem == 0 of
             True -> do
@@ -570,7 +571,7 @@ genStore_fast env addr r n val
                 case pLower grt == getVarType vval of
                      -- were fine
                      True  -> do
-                         let s3 = Store vval ptr
+                         let s3 = MetaStmt meta $ Store vval ptr
                          return (env',  stmts `snocOL` s1 `snocOL` s2
                                  `snocOL` s3, top)
 
@@ -578,19 +579,19 @@ genStore_fast env addr r n val
                      False -> do
                          let ty = (pLift . getVarType) vval
                          (ptr', s3) <- doExpr ty $ Cast LM_Bitcast ptr ty
-                         let s4 = Store vval ptr'
+                         let s4 = MetaStmt meta $ Store vval ptr'
                          return (env',  stmts `snocOL` s1 `snocOL` s2
                                  `snocOL` s3 `snocOL` s4, top)
 
             -- If its a bit type then we use the slow method since
             -- we can't avoid casting anyway.
-            False -> genStore_slow env addr val
+            False -> genStore_slow env addr val meta
 
 
 -- | CmmStore operation
 -- Generic case. Uses casts and pointer arithmetic if needed.
-genStore_slow :: LlvmEnv -> CmmExpr -> CmmExpr -> UniqSM StmtData
-genStore_slow env addr val = do
+genStore_slow :: LlvmEnv -> CmmExpr -> CmmExpr -> [MetaData] -> UniqSM StmtData
+genStore_slow env addr val meta = do
     (env1, vaddr, stmts1, top1) <- exprToVar env addr
     (env2, vval,  stmts2, top2) <- exprToVar env1 val
 
@@ -599,17 +600,17 @@ genStore_slow env addr val = do
         -- sometimes we need to cast an int to a pointer before storing
         LMPointer ty@(LMPointer _) | getVarType vval == llvmWord -> do
             (v, s1) <- doExpr ty $ Cast LM_Inttoptr vval ty
-            let s2 = Store v vaddr
+            let s2 = MetaStmt meta $ Store v vaddr
             return (env2, stmts `snocOL` s1 `snocOL` s2, top1 ++ top2)
 
         LMPointer _ -> do
-            let s1 = Store vval vaddr
+            let s1 = MetaStmt meta $ Store vval vaddr
             return (env2, stmts `snocOL` s1, top1 ++ top2)
 
         i@(LMInt _) | i == llvmWord -> do
             let vty = pLift $ getVarType vval
             (vptr, s1) <- doExpr vty $ Cast LM_Inttoptr vaddr vty
-            let s2 = Store vval vptr
+            let s2 = MetaStmt meta $ Store vval vptr
             return (env2, stmts `snocOL` s1 `snocOL` s2, top1 ++ top2)
 
         other ->
@@ -841,8 +842,8 @@ genMachOp env opt op e = genMachOp_slow env opt op e
 genMachOp_fast :: LlvmEnv -> EOption -> MachOp -> GlobalReg -> Int -> [CmmExpr]
                -> UniqSM ExprData
 genMachOp_fast env opt op r n e
-  = let gr  = lmGlobalRegVar r
-        grt = (pLower . getVarType) gr
+  = let gr   = lmGlobalRegVar r
+        grt  = (pLower . getVarType) gr
         (ix,rem) = n `divMod` ((llvmWidthInBits . pLower) grt  `div` 8)
     in case isPointer grt && rem == 0 of
             True -> do
@@ -1031,7 +1032,7 @@ genLoad env e@(CmmMachOp (MO_Sub _) [
     = genLoad_fast env e r (negate $ fromInteger n) ty
 
 -- generic case
-genLoad env e ty = genLoad_slow env e ty
+genLoad env e ty = genLoad_slow env e ty [top]
 
 -- | Handle CmmLoad expression.
 -- This is a special case for loading from a global register pointer
@@ -1039,9 +1040,10 @@ genLoad env e ty = genLoad_slow env e ty
 genLoad_fast :: LlvmEnv -> CmmExpr -> GlobalReg -> Int -> CmmType
                 -> UniqSM ExprData
 genLoad_fast env e r n ty =
-    let gr  = lmGlobalRegVar r
-        grt = (pLower . getVarType) gr
-        ty' = cmmToLlvmType ty
+    let gr   = lmGlobalRegVar r
+        meta = [getTBAA r]
+        grt  = (pLower . getVarType) gr
+        ty'  = cmmToLlvmType ty
         (ix,rem) = n `divMod` ((llvmWidthInBits . pLower) grt  `div` 8)
     in case isPointer grt && rem == 0 of
             True  -> do
@@ -1051,7 +1053,7 @@ genLoad_fast env e r n ty =
                 case grt == ty' of
                      -- were fine
                      True -> do
-                         (var, s3) <- doExpr ty' $ Load ptr
+                         (var, s3) <- doExpr ty' (MetaExpr meta $ Load ptr)
                          return (env, var, unitOL s1 `snocOL` s2 `snocOL` s3,
                                      [])
 
@@ -1059,29 +1061,31 @@ genLoad_fast env e r n ty =
                      False -> do
                          let pty = pLift ty'
                          (ptr', s3) <- doExpr pty $ Cast LM_Bitcast ptr pty
-                         (var, s4) <- doExpr ty' $ Load ptr'
+                         (var, s4) <- doExpr ty' (MetaExpr meta $ Load ptr')
                          return (env, var, unitOL s1 `snocOL` s2 `snocOL` s3
                                     `snocOL` s4, [])
 
             -- If its a bit type then we use the slow method since
             -- we can't avoid casting anyway.
-            False -> genLoad_slow env e ty
+            False -> genLoad_slow env e ty meta
 
 
 -- | Handle Cmm load expression.
 -- Generic case. Uses casts and pointer arithmetic if needed.
-genLoad_slow :: LlvmEnv -> CmmExpr -> CmmType -> UniqSM ExprData
-genLoad_slow env e ty = do
+genLoad_slow :: LlvmEnv -> CmmExpr -> CmmType -> [MetaData] -> UniqSM ExprData
+genLoad_slow env e ty meta = do
     (env', iptr, stmts, tops) <- exprToVar env e
     case getVarType iptr of
          LMPointer _ -> do
-                    (dvar, load) <- doExpr (cmmToLlvmType ty) $ Load iptr
+                    (dvar, load) <- doExpr (cmmToLlvmType ty)
+                                           (MetaExpr meta $ Load iptr)
                     return (env', dvar, stmts `snocOL` load, tops)
 
          i@(LMInt _) | i == llvmWord -> do
                     let pty = LMPointer $ cmmToLlvmType ty
                     (ptr, cast)  <- doExpr pty $ Cast LM_Inttoptr iptr pty
-                    (dvar, load) <- doExpr (cmmToLlvmType ty) $ Load ptr
+                    (dvar, load) <- doExpr (cmmToLlvmType ty)
+                                           (MetaExpr meta $ Load ptr)
                     return (env', dvar, stmts `snocOL` cast `snocOL` load, tops)
 
          other -> pprPanic "exprToVar: CmmLoad expression is not right type!"
@@ -1099,7 +1103,6 @@ genLoad_slow env e ty = do
 getCmmReg :: LlvmEnv -> CmmReg -> ExprData
 getCmmReg env r@(CmmLocal (LocalReg un _))
   = let exists = varLookup un env
-
         (newv, stmts) = allocReg r
         nenv = varInsert un (pLower $ getVarType newv) env
     in case exists of
@@ -1204,7 +1207,7 @@ funEpilogue Nothing = do
     return (vars, concatOL stmts)
   where
     loadExpr r = do
-        let reg = lmGlobalRegVar r
+        let reg  = lmGlobalRegVar r
         (v,s) <- doExpr (pLower $ getVarType reg) $ Load reg
         return (v, unitOL s)
 
@@ -1214,7 +1217,7 @@ funEpilogue (Just live) = do
     return (vars, concatOL stmts)
   where
     loadExpr r | r `elem` alwaysLive || r `elem` live = do
-        let reg = lmGlobalRegVar r
+        let reg  = lmGlobalRegVar r
         (v,s) <- doExpr (pLower $ getVarType reg) $ Load reg
         return (v, unitOL s)
     loadExpr r = do
diff --git a/compiler/llvmGen/LlvmCodeGen/Ppr.hs b/compiler/llvmGen/LlvmCodeGen/Ppr.hs
index c914bb2431..187d1ecf03 100644
--- a/compiler/llvmGen/LlvmCodeGen/Ppr.hs
+++ b/compiler/llvmGen/LlvmCodeGen/Ppr.hs
@@ -11,6 +11,7 @@ module LlvmCodeGen.Ppr (
 import Llvm
 import LlvmCodeGen.Base
 import LlvmCodeGen.Data
+import LlvmCodeGen.Regs
 
 import CLabel
 import OldCmm
@@ -25,6 +26,16 @@ import Unique
 -- * Top level
 --
 
+-- | Header code for LLVM modules
+pprLlvmHeader :: Doc
+pprLlvmHeader =
+    moduleLayout
+    $+$ text ""
+    $+$ ppLlvmFunctionDecls (map snd ghcInternalFunctions)
+    $+$ ppLlvmMetas stgTBAA
+    $+$ text ""
+
+
 -- | LLVM module layout description for the host target
 moduleLayout :: Doc
 moduleLayout =
@@ -64,11 +75,6 @@ moduleLayout =
 #endif
 
 
--- | Header code for LLVM modules
-pprLlvmHeader :: Doc
-pprLlvmHeader =
-    moduleLayout $+$ text "" $+$ ppLlvmFunctionDecls (map snd ghcInternalFunctions)
-
 -- | Pretty print LLVM data code
 pprLlvmData :: LlvmData -> Doc
 pprLlvmData (globals, types) =
diff --git a/compiler/llvmGen/LlvmCodeGen/Regs.hs b/compiler/llvmGen/LlvmCodeGen/Regs.hs
index ecce7a317b..55b2e0db80 100644
--- a/compiler/llvmGen/LlvmCodeGen/Regs.hs
+++ b/compiler/llvmGen/LlvmCodeGen/Regs.hs
@@ -3,7 +3,8 @@
 --
 
 module LlvmCodeGen.Regs (
-        lmGlobalRegArg, lmGlobalRegVar, alwaysLive
+        lmGlobalRegArg, lmGlobalRegVar, alwaysLive,
+        stgTBAA, top, base, stack, heap, rx, tbaa, getTBAA
     ) where
 
 #include "HsVersions.h"
@@ -11,8 +12,8 @@ module LlvmCodeGen.Regs (
 import Llvm
 
 import CmmExpr
-import Outputable ( panic )
 import FastString
+import Outputable ( panic )
 
 -- | Get the LlvmVar function variable storing the real register
 lmGlobalRegVar :: GlobalReg -> LlvmVar
@@ -49,6 +50,8 @@ lmGlobalReg suf reg
         DoubleReg 2    -> doubleGlobal $ "D2" ++ suf
         _other         -> panic $ "LlvmCodeGen.Reg: GlobalReg (" ++ (show reg)
                                 ++ ") not supported!"
+        -- LongReg, HpLim, CCSS, CurrentTSO, CurrentNusery, HpAlloc
+        -- EagerBlackholeInfo, GCEnter1, GCFun, BaseReg, PicBaseReg
     where
         wordGlobal   name = LMNLocalVar (fsLit name) llvmWord
         ptrGlobal    name = LMNLocalVar (fsLit name) llvmWordPtr
@@ -59,3 +62,41 @@ lmGlobalReg suf reg
 alwaysLive :: [GlobalReg]
 alwaysLive = [BaseReg, Sp, Hp, SpLim, HpLim, node]
 
+-- | STG Type Based Alias Analysis metadata
+stgTBAA :: [LlvmMeta]
+stgTBAA
+  = [ MetaUnamed topN   [MetaStr (fsLit "top")]
+    , MetaUnamed stackN [MetaStr (fsLit "stack"), MetaNode topN]
+    , MetaUnamed heapN  [MetaStr (fsLit "heap"),  MetaNode topN]
+    , MetaUnamed rxN    [MetaStr (fsLit "rx"),    MetaNode heapN]
+    , MetaUnamed baseN  [MetaStr (fsLit "base"),  MetaNode topN]
+    ]
+
+-- | Id values
+topN, stackN, heapN, rxN, baseN :: LlvmMetaUnamed
+topN   = LMMetaUnamed 0
+stackN = LMMetaUnamed 1
+heapN  = LMMetaUnamed 2
+rxN    = LMMetaUnamed 3
+baseN  = LMMetaUnamed 4
+
+-- | The various TBAA types
+top, heap, stack, rx, base :: MetaData
+top   = (tbaa, topN)
+heap  = (tbaa, heapN)
+stack = (tbaa, stackN)
+rx    = (tbaa, rxN)
+base  = (tbaa, baseN)
+
+-- | The TBAA metadata identifier
+tbaa :: LMString
+tbaa = fsLit "tbaa"
+
+-- | Get the correct TBAA metadata information for this register type
+getTBAA :: GlobalReg -> MetaData
+getTBAA BaseReg          = base
+getTBAA Sp               = stack
+getTBAA Hp               = heap
+getTBAA (VanillaReg _ _) = rx
+getTBAA _                = top
+