Add new mbmi and mbmi2 compiler flags

This adds support for the bit deposit and extraction operations provided
by the BMI and BMI2 instruction set extensions on modern amd64 machines.

Implement x86 code generator for pdep and pext.  Properly initialise
bmiVersion field.

pdep and pext test cases

Fix pattern match for pdep and pext instructions

Fix build of pdep and pext code for 32-bit architectures

Test Plan: Validate

Reviewers: austin, simonmar, bgamari, angerman

Reviewed By: bgamari

Subscribers: trommler, carter, angerman, thomie, rwbarton, newhoggy

GHC Trac Issues: #14206

Differential Revision: https://phabricator.haskell.org/D4236

1 files changed, 77 insertions, 20 deletions

diff --git a/compiler/llvmGen/LlvmCodeGen/CodeGen.hs b/compiler/llvmGen/LlvmCodeGen/CodeGen.hs
index a88642b531..e812dd445f 100644
--- a/compiler/llvmGen/LlvmCodeGen/CodeGen.hs
+++ b/compiler/llvmGen/LlvmCodeGen/CodeGen.hs
@@ -46,6 +46,8 @@ import Data.Maybe ( catMaybes )
 type Atomic = Bool
 type LlvmStatements = OrdList LlvmStatement
 
+data Signage = Signed | Unsigned deriving (Eq, Show)
+
 -- -----------------------------------------------------------------------------
 -- | Top-level of the LLVM proc Code generator
 --
@@ -207,7 +209,7 @@ genCall t@(PrimTarget (MO_Prefetch_Data localityInt)) [] args
     let args_hints' = zip args arg_hints
     argVars <- arg_varsW args_hints' ([], nilOL, [])
     fptr    <- liftExprData $ getFunPtr funTy t
-    argVars' <- castVarsW $ zip argVars argTy
+    argVars' <- castVarsW Signed $ zip argVars argTy
 
     doTrashStmts
     let argSuffix = [mkIntLit i32 0, mkIntLit i32 localityInt, mkIntLit i32 1]
@@ -218,6 +220,11 @@ genCall t@(PrimTarget (MO_Prefetch_Data localityInt)) [] args
 -- and return types
 genCall t@(PrimTarget (MO_PopCnt w)) dsts args =
     genCallSimpleCast w t dsts args
+
+genCall t@(PrimTarget (MO_Pdep w)) dsts args =
+    genCallSimpleCast2 w t dsts args
+genCall t@(PrimTarget (MO_Pext w)) dsts args =
+    genCallSimpleCast2 w t dsts args
 genCall t@(PrimTarget (MO_Clz w)) dsts args =
     genCallSimpleCast w t dsts args
 genCall t@(PrimTarget (MO_Ctz w)) dsts args =
@@ -285,7 +292,7 @@ genCall t@(PrimTarget op) [] args
     let args_hints = zip args arg_hints
     argVars       <- arg_varsW args_hints ([], nilOL, [])
     fptr          <- getFunPtrW funTy t
-    argVars' <- castVarsW $ zip argVars argTy
+    argVars' <- castVarsW Signed $ zip argVars argTy
 
     doTrashStmts
     let alignVal = mkIntLit i32 align
@@ -518,7 +525,7 @@ genCallExtract target@(PrimTarget op) w (argA, argB) (llvmTypeA, llvmTypeB) = do
     -- Process the arguments.
     let args_hints = zip [argA, argB] (snd $ foreignTargetHints target)
     (argsV1, args1, top1) <- arg_vars args_hints ([], nilOL, [])
-    (argsV2, args2) <- castVars $ zip argsV1 argTy
+    (argsV2, args2) <- castVars Signed $ zip argsV1 argTy
 
     -- Get the function and make the call.
     fname <- cmmPrimOpFunctions op
@@ -558,9 +565,9 @@ genCallSimpleCast w t@(PrimTarget op) [dst] args = do
     let (_, arg_hints) = foreignTargetHints t
     let args_hints = zip args arg_hints
     (argsV, stmts2, top2)       <- arg_vars args_hints ([], nilOL, [])
-    (argsV', stmts4)            <- castVars $ zip argsV [width]
+    (argsV', stmts4)            <- castVars Signed $ zip argsV [width]
     (retV, s1)                  <- doExpr width $ Call StdCall fptr argsV' []
-    ([retV'], stmts5)           <- castVars [(retV,dstTy)]
+    ([retV'], stmts5)           <- castVars (cmmPrimOpRetValSignage op) [(retV,dstTy)]
     let s2                       = Store retV' dstV
 
     let stmts = stmts2 `appOL` stmts4 `snocOL`
@@ -569,6 +576,37 @@ genCallSimpleCast w t@(PrimTarget op) [dst] args = do
 genCallSimpleCast _ _ dsts _ =
     panic ("genCallSimpleCast: " ++ show (length dsts) ++ " dsts")
 
+-- Handle simple function call that only need simple type casting, of the form:
+--   truncate arg >>= \a -> call(a) >>= zext
+--
+-- since GHC only really has i32 and i64 types and things like Word8 are backed
+-- by an i32 and just present a logical i8 range. So we must handle conversions
+-- from i32 to i8 explicitly as LLVM is strict about types.
+genCallSimpleCast2 :: Width -> ForeignTarget -> [CmmFormal] -> [CmmActual]
+              -> LlvmM StmtData
+genCallSimpleCast2 w t@(PrimTarget op) [dst] args = do
+    let width = widthToLlvmInt w
+        dstTy = cmmToLlvmType $ localRegType dst
+
+    fname                       <- cmmPrimOpFunctions op
+    (fptr, _, top3)             <- getInstrinct fname width (const width <$> args)
+
+    dstV                        <- getCmmReg (CmmLocal dst)
+
+    let (_, arg_hints) = foreignTargetHints t
+    let args_hints = zip args arg_hints
+    (argsV, stmts2, top2)       <- arg_vars args_hints ([], nilOL, [])
+    (argsV', stmts4)            <- castVars Signed $ zip argsV (const width <$> argsV)
+    (retV, s1)                  <- doExpr width $ Call StdCall fptr argsV' []
+    ([retV'], stmts5)           <- castVars (cmmPrimOpRetValSignage op) [(retV,dstTy)]
+    let s2                       = Store retV' dstV
+
+    let stmts = stmts2 `appOL` stmts4 `snocOL`
+                s1 `appOL` stmts5 `snocOL` s2
+    return (stmts, top2 ++ top3)
+genCallSimpleCast2 _ _ dsts _ =
+    panic ("genCallSimpleCast2: " ++ show (length dsts) ++ " dsts")
+
 -- | Create a function pointer from a target.
 getFunPtrW :: (LMString -> LlvmType) -> ForeignTarget
            -> WriterT LlvmAccum LlvmM LlvmVar
@@ -638,31 +676,32 @@ arg_vars ((e, _):rest) (vars, stmts, tops)
 
 
 -- | Cast a collection of LLVM variables to specific types.
-castVarsW :: [(LlvmVar, LlvmType)]
+castVarsW :: Signage
+          -> [(LlvmVar, LlvmType)]
           -> WriterT LlvmAccum LlvmM [LlvmVar]
-castVarsW vars = do
-    (vars, stmts) <- lift $ castVars vars
+castVarsW signage vars = do
+    (vars, stmts) <- lift $ castVars signage vars
     tell $ LlvmAccum stmts mempty
     return vars
 
 -- | Cast a collection of LLVM variables to specific types.
-castVars :: [(LlvmVar, LlvmType)]
+castVars :: Signage -> [(LlvmVar, LlvmType)]
          -> LlvmM ([LlvmVar], LlvmStatements)
-castVars vars = do
-                done <- mapM (uncurry castVar) vars
+castVars signage vars = do
+                done <- mapM (uncurry (castVar signage)) vars
                 let (vars', stmts) = unzip done
                 return (vars', toOL stmts)
 
 -- | Cast an LLVM variable to a specific type, panicing if it can't be done.
-castVar :: LlvmVar -> LlvmType -> LlvmM (LlvmVar, LlvmStatement)
-castVar v t | getVarType v == t
+castVar :: Signage -> LlvmVar -> LlvmType -> LlvmM (LlvmVar, LlvmStatement)
+castVar signage v t | getVarType v == t
             = return (v, Nop)
 
             | otherwise
             = do dflags <- getDynFlags
                  let op = case (getVarType v, t) of
                       (LMInt n, LMInt m)
-                          -> if n < m then LM_Sext else LM_Trunc
+                          -> if n < m then extend else LM_Trunc
                       (vt, _) | isFloat vt && isFloat t
                           -> if llvmWidthInBits dflags vt < llvmWidthInBits dflags t
                                 then LM_Fpext else LM_Fptrunc
@@ -676,7 +715,16 @@ castVar v t | getVarType v == t
                       (vt, _) -> panic $ "castVars: Can't cast this type ("
                                   ++ showSDoc dflags (ppr vt) ++ ") to (" ++ showSDoc dflags (ppr t) ++ ")"
                  doExpr t $ Cast op v t
+    where extend = case signage of
+            Signed      -> LM_Sext
+            Unsigned    -> LM_Zext
+
 
+cmmPrimOpRetValSignage :: CallishMachOp -> Signage
+cmmPrimOpRetValSignage mop = case mop of
+    MO_Pdep _   -> Unsigned
+    MO_Pext _   -> Unsigned
+    _           -> Signed
 
 -- | Decide what C function to use to implement a CallishMachOp
 cmmPrimOpFunctions :: CallishMachOp -> LlvmM LMString
@@ -735,6 +783,15 @@ cmmPrimOpFunctions mop = do
     (MO_Clz w)    -> fsLit $ "llvm.ctlz."   ++ showSDoc dflags (ppr $ widthToLlvmInt w)
     (MO_Ctz w)    -> fsLit $ "llvm.cttz."   ++ showSDoc dflags (ppr $ widthToLlvmInt w)
 
+    (MO_Pdep w)   ->  let w' = showSDoc dflags (ppr $ widthInBits w)
+                      in  if isBmi2Enabled dflags
+                            then fsLit $ "llvm.x86.bmi.pdep."   ++ w'
+                            else fsLit $ "hs_pdep"              ++ w'
+    (MO_Pext w)   ->  let w' = showSDoc dflags (ppr $ widthInBits w)
+                      in  if isBmi2Enabled dflags
+                            then fsLit $ "llvm.x86.bmi.pext."   ++ w'
+                            else fsLit $ "hs_pext"              ++ w'
+
     (MO_Prefetch_Data _ )-> fsLit "llvm.prefetch"
 
     MO_AddIntC w    -> fsLit $ "llvm.sadd.with.overflow."
@@ -1212,7 +1269,7 @@ genMachOp _ op [x] = case op of
 
         negateVec ty v2 negOp = do
             (vx, stmts1, top) <- exprToVar x
-            ([vx'], stmts2) <- castVars [(vx, ty)]
+            ([vx'], stmts2) <- castVars Signed [(vx, ty)]
             (v1, s1) <- doExpr ty $ LlvmOp negOp v2 vx'
             return (v1, stmts1 `appOL` stmts2 `snocOL` s1, top)
 
@@ -1275,7 +1332,7 @@ genMachOp_slow :: EOption -> MachOp -> [CmmExpr] -> LlvmM ExprData
 genMachOp_slow _ (MO_V_Extract l w) [val, idx] = runExprData $ do
     vval <- exprToVarW val
     vidx <- exprToVarW idx
-    [vval'] <- castVarsW [(vval, LMVector l ty)]
+    [vval'] <- castVarsW Signed [(vval, LMVector l ty)]
     doExprW ty $ Extract vval' vidx
   where
     ty = widthToLlvmInt w
@@ -1283,7 +1340,7 @@ genMachOp_slow _ (MO_V_Extract l w) [val, idx] = runExprData $ do
 genMachOp_slow _ (MO_VF_Extract l w) [val, idx] = runExprData $ do
     vval <- exprToVarW val
     vidx <- exprToVarW idx
-    [vval'] <- castVarsW [(vval, LMVector l ty)]
+    [vval'] <- castVarsW Signed [(vval, LMVector l ty)]
     doExprW ty $ Extract vval' vidx
   where
     ty = widthToLlvmFloat w
@@ -1293,7 +1350,7 @@ genMachOp_slow _ (MO_V_Insert l w) [val, elt, idx] = runExprData $ do
     vval <- exprToVarW val
     velt <- exprToVarW elt
     vidx <- exprToVarW idx
-    [vval'] <- castVarsW [(vval, ty)]
+    [vval'] <- castVarsW Signed [(vval, ty)]
     doExprW ty $ Insert vval' velt vidx
   where
     ty = LMVector l (widthToLlvmInt w)
@@ -1302,7 +1359,7 @@ genMachOp_slow _ (MO_VF_Insert l w) [val, elt, idx] = runExprData $ do
     vval <- exprToVarW val
     velt <- exprToVarW elt
     vidx <- exprToVarW idx
-    [vval'] <- castVarsW [(vval, ty)]
+    [vval'] <- castVarsW Signed [(vval, ty)]
     doExprW ty $ Insert vval' velt vidx
   where
     ty = LMVector l (widthToLlvmFloat w)
@@ -1414,7 +1471,7 @@ genMachOp_slow opt op [x, y] = case op of
         binCastLlvmOp ty binOp = runExprData $ do
             vx <- exprToVarW x
             vy <- exprToVarW y
-            [vx', vy'] <- castVarsW [(vx, ty), (vy, ty)]
+            [vx', vy'] <- castVarsW Signed [(vx, ty), (vy, ty)]
             doExprW ty $ binOp vx' vy'
 
         -- | Need to use EOption here as Cmm expects word size results from