Add support for -mavx and -mavx=2 flags.

3 files changed, 55 insertions, 16 deletions

diff --git a/compiler/main/CmdLineParser.hs b/compiler/main/CmdLineParser.hs
index 252a376432..b8086fa33e 100644
--- a/compiler/main/CmdLineParser.hs
+++ b/compiler/main/CmdLineParser.hs
@@ -54,6 +54,7 @@ data OptKind m                             -- Suppose the flag is -f
     | PrefixPred    (String -> Bool) (String -> EwM m ())
     | AnySuffixPred (String -> Bool) (String -> EwM m ())
     | VersionSuffix (Int -> Int -> EwM m ())
+    | OptVersionSuffix (Int -> Int -> EwM m ())
       -- -f or -f=maj.min; pass major and minor version to fn
 
 
@@ -205,6 +206,13 @@ processOneArg opt_kind rest arg args
                           Just maj <- parseInt maj_s -> Right (f maj 0, args)
                         | otherwise -> Left ("malformed version argument in " ++ dash_arg)
 
+        OptVersionSuffix f | [maj_s, min_s] <- split '.' rest_no_eq,
+                             Just maj <- parseInt maj_s,
+                             Just min <- parseInt min_s -> Right (f maj min, args)
+                           | [maj_s] <- split '.' rest_no_eq,
+                             Just maj <- parseInt maj_s -> Right (f maj 0, args)
+                           | otherwise -> Right (f 1 0, args)
+
 
 findArg :: [Flag m] -> String -> Maybe (String, OptKind m)
 findArg spec arg =
@@ -219,19 +227,20 @@ findArg spec arg =
         (one:_) -> Just one
 
 arg_ok :: OptKind t -> [Char] -> String -> Bool
-arg_ok (NoArg           _)  rest _   = null rest
-arg_ok (HasArg          _)  _    _   = True
-arg_ok (SepArg          _)  rest _   = null rest
-arg_ok (Prefix          _)  rest _   = notNull rest
-arg_ok (PrefixPred p    _)  rest _   = notNull rest && p (dropEq rest)
-arg_ok (OptIntSuffix    _)  _    _   = True
-arg_ok (IntSuffix       _)  _    _   = True
-arg_ok (FloatSuffix     _)  _    _   = True
-arg_ok (OptPrefix       _)  _    _   = True
-arg_ok (PassFlag        _)  rest _   = null rest
-arg_ok (AnySuffix       _)  _    _   = True
-arg_ok (AnySuffixPred p _)  _    arg = p arg
-arg_ok (VersionSuffix   _)  _    _   = True
+arg_ok (NoArg            _)  rest _   = null rest
+arg_ok (HasArg           _)  _    _   = True
+arg_ok (SepArg           _)  rest _   = null rest
+arg_ok (Prefix           _)  rest _   = notNull rest
+arg_ok (PrefixPred p     _)  rest _   = notNull rest && p (dropEq rest)
+arg_ok (OptIntSuffix     _)  _    _   = True
+arg_ok (IntSuffix        _)  _    _   = True
+arg_ok (FloatSuffix      _)  _    _   = True
+arg_ok (OptPrefix        _)  _    _   = True
+arg_ok (PassFlag         _)  rest _   = null rest
+arg_ok (AnySuffix        _)  _    _   = True
+arg_ok (AnySuffixPred p  _)  _    arg = p arg
+arg_ok (VersionSuffix    _)  _    _   = True
+arg_ok (OptVersionSuffix _)  _    _   = True
 
 -- | Parse an Int
 --
diff --git a/compiler/main/DriverPipeline.hs b/compiler/main/DriverPipeline.hs
index 67377e68c2..8fef336574 100644
--- a/compiler/main/DriverPipeline.hs
+++ b/compiler/main/DriverPipeline.hs
@@ -1434,7 +1434,8 @@ runPhase (RealPhase LlvmLlc) input_fn dflags
                 ++ [SysTools.Option tbaa]
                 ++ map SysTools.Option fpOpts
                 ++ map SysTools.Option abiOpts
-                ++ map SysTools.Option sseOpts)
+                ++ map SysTools.Option sseOpts
+                ++ map SysTools.Option avxOpts)
 
     return (RealPhase next_phase, output_fn)
   where
@@ -1467,6 +1468,10 @@ runPhase (RealPhase LlvmLlc) input_fn dflags
                 | isSse2Enabled dflags   = ["-mattr=+sse2"]
                 | otherwise              = []
 
+        avxOpts | isAvxEnabled dflags    = ["-mattr=+avx"]
+                | isAvx2Enabled dflags   = ["-mattr=+avx2"]
+                | otherwise              = []
+
 -----------------------------------------------------------------------------
 -- LlvmMangle phase
 
@@ -2033,6 +2038,12 @@ doCpp dflags raw input_fn output_fn = do
           [ "-D__SSE2__=1" | sse2 || sse4_2 ] ++
           [ "-D__SSE4_2__=1" | sse4_2 ]
 
+    let avx = isAvxEnabled dflags
+        avx2 = isAvx2Enabled dflags
+        avx_defs =
+          [ "-D__AVX__=1" | avx || avx2 ] ++
+          [ "-D__AVX2__=1" | avx2 ]
+
     backend_defs <- getBackendDefs dflags
 
     cpp_prog       (   map SysTools.Option verbFlags
@@ -2042,6 +2053,7 @@ doCpp dflags raw input_fn output_fn = do
                     ++ map SysTools.Option backend_defs
                     ++ map SysTools.Option hscpp_opts
                     ++ map SysTools.Option sse_defs
+                    ++ map SysTools.Option avx_defs
         -- Set the language mode to assembler-with-cpp when preprocessing. This
         -- alleviates some of the C99 macro rules relating to whitespace and the hash
         -- operator, which we tend to abuse. Clang in particular is not very happy
diff --git a/compiler/main/DynFlags.hs b/compiler/main/DynFlags.hs
index e1236840a0..7a675edc1f 100644
--- a/compiler/main/DynFlags.hs
+++ b/compiler/main/DynFlags.hs
@@ -127,9 +127,11 @@ module DynFlags (
 
         unsafeGlobalDynFlags, setUnsafeGlobalDynFlags,
 
-        -- * SSE
+        -- * SSE and AVX
         isSse2Enabled,
         isSse4_2Enabled,
+        isAvxEnabled,
+        isAvx2Enabled,
 
         -- * Linker information
         LinkerInfo(..),
@@ -747,6 +749,7 @@ data DynFlags = DynFlags {
 
   -- | Machine dependant flags (-m<blah> stuff)
   sseVersion            :: Maybe (Int, Int),  -- (major, minor)
+  avxVersion            :: Maybe (Int, Int),  -- (major, minor)
 
   -- | Run-time linker information (what options we need, etc.)
   rtldFlags             :: IORef (Maybe LinkerInfo)
@@ -1359,6 +1362,7 @@ defaultDynFlags mySettings =
         interactivePrint = Nothing,
         nextWrapperNum = panic "defaultDynFlags: No nextWrapperNum",
         sseVersion = Nothing,
+        avxVersion = Nothing,
         rtldFlags = panic "defaultDynFlags: no rtldFlags"
       }
 
@@ -2259,6 +2263,7 @@ dynamic_flags = [
   , Flag "monly-3-regs" (NoArg (addWarn "The -monly-3-regs flag does nothing; it will be removed in a future GHC release"))
   , Flag "monly-4-regs" (NoArg (addWarn "The -monly-4-regs flag does nothing; it will be removed in a future GHC release"))
   , Flag "msse"         (versionSuffix (\maj min d -> d{ sseVersion = Just (maj, min) }))
+  , Flag "mavx"         (optVersionSuffix (\maj min d -> d{ avxVersion = Just (maj, min) }))
 
      ------ Warning opts -------------------------------------------------
   , Flag "W"      (NoArg (mapM_ setWarningFlag minusWOpts))
@@ -3035,6 +3040,9 @@ optIntSuffixM fn = OptIntSuffix (\mi -> updM (fn mi))
 versionSuffix :: (Int -> Int -> DynFlags -> DynFlags) -> OptKind (CmdLineP DynFlags)
 versionSuffix fn = VersionSuffix (\maj min -> upd (fn maj min))
 
+optVersionSuffix :: (Int -> Int -> DynFlags -> DynFlags) -> OptKind (CmdLineP DynFlags)
+optVersionSuffix fn = OptVersionSuffix (\maj min -> upd (fn maj min))
+
 setDumpFlag :: DumpFlag -> OptKind (CmdLineP DynFlags)
 setDumpFlag dump_flag = NoArg (setDumpFlag' dump_flag)
 
@@ -3519,7 +3527,7 @@ setUnsafeGlobalDynFlags :: DynFlags -> IO ()
 setUnsafeGlobalDynFlags = writeIORef v_unsafeGlobalDynFlags
 
 -- -----------------------------------------------------------------------------
--- SSE
+-- SSE and AVX
 
 -- TODO: Instead of using a separate predicate (i.e. isSse2Enabled) to
 -- check if SSE is enabled, we might have x86-64 imply the -msse2
@@ -3539,6 +3547,16 @@ isSse2Enabled dflags = case platformArch (targetPlatform dflags) of
 isSse4_2Enabled :: DynFlags -> Bool
 isSse4_2Enabled dflags = sseVersion dflags >= Just (4,2)
 
+isAvxEnabled :: DynFlags -> Bool
+isAvxEnabled dflags = case platformArch (targetPlatform dflags) of
+    ArchX86_64 -> avxVersion dflags >= Just (1,0)
+    _          -> False
+
+isAvx2Enabled :: DynFlags -> Bool
+isAvx2Enabled dflags = case platformArch (targetPlatform dflags) of
+    ArchX86_64 -> avxVersion dflags >= Just (2,0)
+    _          -> False
+
 -- -----------------------------------------------------------------------------
 -- Linker information