LLVM: refactor and comment register padding code (#17920)

2 files changed, 97 insertions, 97 deletions

diff --git a/compiler/GHC/CmmToLlvm/Base.hs b/compiler/GHC/CmmToLlvm/Base.hs
index 38b9b8e582..7f60d660cb 100644
--- a/compiler/GHC/CmmToLlvm/Base.hs
+++ b/compiler/GHC/CmmToLlvm/Base.hs
@@ -42,12 +42,14 @@ module GHC.CmmToLlvm.Base (
 #include "ghcautoconf.h"
 
 import GHC.Prelude
+import GHC.Utils.Panic
 
 import GHC.Llvm
 import GHC.CmmToLlvm.Regs
 
 import GHC.Cmm.CLabel
-import GHC.Platform.Regs ( activeStgRegs )
+import GHC.Cmm.Ppr.Expr ()
+import GHC.Platform.Regs ( activeStgRegs, globalRegMaybe )
 import GHC.Driver.Session
 import GHC.Data.FastString
 import GHC.Cmm              hiding ( succ )
@@ -65,7 +67,8 @@ import qualified GHC.Data.Stream as Stream
 import Data.Maybe (fromJust)
 import Control.Monad (ap)
 import Data.Char (isDigit)
-import Data.List (sort, groupBy, intercalate)
+import Data.List (sortBy, groupBy, intercalate)
+import Data.Ord (comparing)
 import qualified Data.List.NonEmpty as NE
 
 -- ----------------------------------------------------------------------------
@@ -157,8 +160,10 @@ llvmFunArgs :: Platform -> LiveGlobalRegs -> [LlvmVar]
 llvmFunArgs platform live =
     map (lmGlobalRegArg platform) (filter isPassed allRegs)
     where allRegs = activeStgRegs platform
-          paddedLive = map (\(_,r) -> r) $ padLiveArgs platform live
-          isLive r = r `elem` alwaysLive || r `elem` paddedLive
+          paddingRegs = padLiveArgs platform live
+          isLive r = r `elem` alwaysLive
+                     || r `elem` live
+                     || r `elem` paddingRegs
           isPassed r = not (isFPR r) || isLive r
 
 
@@ -170,91 +175,76 @@ isFPR (YmmReg _)    = True
 isFPR (ZmmReg _)    = True
 isFPR _             = False
 
-sameFPRClass :: GlobalReg -> GlobalReg -> Bool
-sameFPRClass (FloatReg _)  (FloatReg _) = True
-sameFPRClass (DoubleReg _) (DoubleReg _) = True
-sameFPRClass (XmmReg _)    (XmmReg _) = True
-sameFPRClass (YmmReg _)    (YmmReg _) = True
-sameFPRClass (ZmmReg _)    (ZmmReg _) = True
-sameFPRClass _             _          = False
-
-normalizeFPRNum :: GlobalReg -> GlobalReg
-normalizeFPRNum (FloatReg _)  = FloatReg 1
-normalizeFPRNum (DoubleReg _) = DoubleReg 1
-normalizeFPRNum (XmmReg _)    = XmmReg 1
-normalizeFPRNum (YmmReg _)    = YmmReg 1
-normalizeFPRNum (ZmmReg _)    = ZmmReg 1
-normalizeFPRNum _ = error "normalizeFPRNum expected only FPR regs"
-
-getFPRCtor :: GlobalReg -> Int -> GlobalReg
-getFPRCtor (FloatReg _)  = FloatReg
-getFPRCtor (DoubleReg _) = DoubleReg
-getFPRCtor (XmmReg _)    = XmmReg
-getFPRCtor (YmmReg _)    = YmmReg
-getFPRCtor (ZmmReg _)    = ZmmReg
-getFPRCtor _ = error "getFPRCtor expected only FPR regs"
-
-fprRegNum :: GlobalReg -> Int
-fprRegNum (FloatReg i)  = i
-fprRegNum (DoubleReg i) = i
-fprRegNum (XmmReg i)    = i
-fprRegNum (YmmReg i)    = i
-fprRegNum (ZmmReg i)    = i
-fprRegNum _ = error "fprRegNum expected only FPR regs"
-
--- | Input: dynflags, and the list of live registers
+-- | Return a list of "padding" registers for LLVM function calls.
 --
--- Output: An augmented list of live registers, where padding was
--- added to the list of registers to ensure the calling convention is
--- correctly used by LLVM.
+-- When we generate LLVM function signatures, we can't just make any register
+-- alive on function entry. Instead, we need to insert fake arguments of the
+-- same register class until we are sure that one of them is mapped to the
+-- register we want alive. E.g. to ensure that F5 is alive, we may need to
+-- insert fake arguments mapped to F1, F2, F3 and F4.
 --
--- Each global reg in the returned list is tagged with a bool, which
--- indicates whether the global reg was added as padding, or was an original
--- live register.
---
--- That is, True => padding, False => a real, live global register.
---
--- Also, the returned list is not sorted in any particular order.
---
-padLiveArgs :: Platform -> LiveGlobalRegs -> [(Bool, GlobalReg)]
-padLiveArgs plat live =
-      if platformUnregisterised plat
-        then taggedLive -- not using GHC's register convention for platform.
-        else padding ++ taggedLive
+-- Invariant: Cmm FPR regs with number "n" maps to real registers with number
+-- "n" If the calling convention uses registers in a different order or if the
+-- invariant doesn't hold, this code probably won't be correct.
+padLiveArgs :: Platform -> LiveGlobalRegs -> LiveGlobalRegs
+padLiveArgs platform live =
+      if platformUnregisterised platform
+        then [] -- not using GHC's register convention for platform.
+        else padded
   where
-    taggedLive = map (\x -> (False, x)) live
-
-    fprLive = filter isFPR live
-    padding = concatMap calcPad $ groupBy sharesClass fprLive
-
-    sharesClass :: GlobalReg -> GlobalReg -> Bool
-    sharesClass a b = sameFPRClass a b || overlappingClass
+    ----------------------------------
+    -- handle floating-point registers (FPR)
+
+    fprLive = filter isFPR live  -- real live FPR registers
+
+    -- we group live registers sharing the same classes, i.e. that use the same
+    -- set of real registers to be passed. E.g. FloatReg, DoubleReg and XmmReg
+    -- all use the same real regs on X86-64 (XMM registers).
+    --
+    classes         = groupBy sharesClass fprLive
+    sharesClass a b = regsOverlap platform (norm a) (norm b) -- check if mapped to overlapping registers
+    norm x          = CmmGlobal ((fpr_ctor x) 1)             -- get the first register of the family
+
+    -- For each class, we just have to fill missing registers numbers. We use
+    -- the constructor of the greatest register to build padding registers.
+    --
+    -- E.g. sortedRs = [   F2,   XMM4, D5]
+    --      output   = [D1,   D3]
+    padded      = concatMap padClass classes
+    padClass rs = go sortedRs [1..]
       where
-        overlappingClass = regsOverlap plat (norm a) (norm b)
-        norm = CmmGlobal . normalizeFPRNum
-
-    calcPad :: [GlobalReg] -> [(Bool, GlobalReg)]
-    calcPad rs = getFPRPadding (getFPRCtor $ head rs) rs
-
-getFPRPadding :: (Int -> GlobalReg) -> LiveGlobalRegs -> [(Bool, GlobalReg)]
-getFPRPadding paddingCtor live = padding
-    where
-        fprRegNums = sort $ map fprRegNum live
-        (_, padding) = foldl assignSlots (1, []) $ fprRegNums
-
-        assignSlots (i, acc) regNum
-            | i == regNum = -- don't need padding here
-                  (i+1, acc)
-            | i < regNum = let -- add padding for slots i .. regNum-1
-                  numNeeded = regNum-i
-                  acc' = genPad i numNeeded ++ acc
-                in
-                  (regNum+1, acc')
-            | otherwise = error "padLiveArgs -- i > regNum ??"
-
-        genPad start n =
-            take n $ flip map (iterate (+1) start) (\i ->
-                (True, paddingCtor i))
+         sortedRs = sortBy (comparing fpr_num) rs
+         maxr     = last sortedRs
+         ctor     = fpr_ctor maxr
+
+         go [] _ = []
+         go (c1:c2:_) _   -- detect bogus case (see #17920)
+            | fpr_num c1 == fpr_num c2
+            , Just real <- globalRegMaybe platform c1
+            = sorryDoc "LLVM code generator" $
+               text "Found two different Cmm registers (" <> ppr c1 <> text "," <> ppr c2 <>
+               text ") both alive AND mapped to the same real register: " <> ppr real <>
+               text ". This isn't currently supported by the LLVM backend."
+         go (c:cs) (f:fs)
+            | fpr_num c == f = go cs fs              -- already covered by a real register
+            | otherwise      = ctor f : go (c:cs) fs -- add padding register
+         go _ _ = undefined -- unreachable
+
+    fpr_ctor :: GlobalReg -> Int -> GlobalReg
+    fpr_ctor (FloatReg _)  = FloatReg
+    fpr_ctor (DoubleReg _) = DoubleReg
+    fpr_ctor (XmmReg _)    = XmmReg
+    fpr_ctor (YmmReg _)    = YmmReg
+    fpr_ctor (ZmmReg _)    = ZmmReg
+    fpr_ctor _ = error "fpr_ctor expected only FPR regs"
+
+    fpr_num :: GlobalReg -> Int
+    fpr_num (FloatReg i)  = i
+    fpr_num (DoubleReg i) = i
+    fpr_num (XmmReg i)    = i
+    fpr_num (YmmReg i)    = i
+    fpr_num (ZmmReg i)    = i
+    fpr_num _ = error "fpr_num expected only FPR regs"
 
 
 -- | Llvm standard fun attributes
diff --git a/compiler/GHC/CmmToLlvm/CodeGen.hs b/compiler/GHC/CmmToLlvm/CodeGen.hs
index 53f17f545c..6e424b7e48 100644
--- a/compiler/GHC/CmmToLlvm/CodeGen.hs
+++ b/compiler/GHC/CmmToLlvm/CodeGen.hs
@@ -1,4 +1,4 @@
-{-# LANGUAGE CPP, GADTs #-}
+{-# LANGUAGE CPP, GADTs, MultiWayIf #-}
 {-# OPTIONS_GHC -fno-warn-type-defaults #-}
 {-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}
 -- ----------------------------------------------------------------------------
@@ -38,6 +38,7 @@ import GHC.Utils.Misc
 
 import Control.Monad.Trans.Class
 import Control.Monad.Trans.Writer
+import Control.Monad
 
 import qualified Data.Semigroup as Semigroup
 import Data.List ( nub )
@@ -1848,7 +1849,7 @@ funPrologue live cmmBlocks = do
       isLive r     = r `elem` alwaysLive || r `elem` live
 
   platform <- getPlatform
-  stmtss <- flip mapM assignedRegs $ \reg ->
+  stmtss <- forM assignedRegs $ \reg ->
     case reg of
       CmmLocal (LocalReg un _) -> do
         let (newv, stmts) = allocReg reg
@@ -1875,9 +1876,7 @@ funEpilogue :: LiveGlobalRegs -> LlvmM ([LlvmVar], LlvmStatements)
 funEpilogue live = do
     platform <- getPlatform
 
-    -- the bool indicates whether the register is padding.
-    let alwaysNeeded = map (\r -> (False, r)) alwaysLive
-        livePadded = alwaysNeeded ++ padLiveArgs platform live
+    let paddingRegs = padLiveArgs platform live
 
     -- Set to value or "undef" depending on whether the register is
     -- actually live
@@ -1887,14 +1886,25 @@ funEpilogue live = do
         loadUndef r = do
           let ty = (pLower . getVarType $ lmGlobalRegVar platform r)
           return (Just $ LMLitVar $ LMUndefLit ty, nilOL)
-    platform <- getPlatform
+
+    -- Note that floating-point registers in `activeStgRegs` must be sorted
+    -- according to the calling convention.
+    --  E.g. for X86:
+    --     GOOD: F1,D1,XMM1,F2,D2,XMM2,...
+    --     BAD : F1,F2,F3,D1,D2,D3,XMM1,XMM2,XMM3,...
+    --  As Fn, Dn and XMMn use the same register (XMMn) to be passed, we don't
+    --  want to pass F2 before D1 for example, otherwise we could get F2 -> XMM1
+    --  and D1 -> XMM2.
     let allRegs = activeStgRegs platform
-    loads <- flip mapM allRegs $ \r -> case () of
-      _ | (False, r) `elem` livePadded
-                             -> loadExpr r   -- if r is not padding, load it
-        | not (isFPR r) || (True, r) `elem` livePadded
-                             -> loadUndef r
-        | otherwise          -> return (Nothing, nilOL)
+    loads <- forM allRegs $ \r -> if
+      -- load live registers
+      | r `elem` alwaysLive  -> loadExpr r
+      | r `elem` live        -> loadExpr r
+      -- load all non Floating-Point Registers
+      | not (isFPR r)        -> loadUndef r
+      -- load padding Floating-Point Registers
+      | r `elem` paddingRegs -> loadUndef r
+      | otherwise            -> return (Nothing, nilOL)
 
     let (vars, stmts) = unzip loads
     return (catMaybes vars, concatOL stmts)