Plugin.hs

{-# LANGUAGE CPP             #-}
{-# LANGUAGE RecordWildCards #-}
{-# OPTIONS_GHC -fno-warn-unused-imports   #-}
{-# OPTIONS_GHC -fno-warn-unused-top-binds #-}
-- |
-- Module      : Data.Array.Accelerate.LLVM.Native.Plugin
-- Copyright   : [2017..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--

module Data.Array.Accelerate.LLVM.Native.Plugin (

  plugin,

) where

import Data.Array.Accelerate.Error
import Data.Array.Accelerate.LLVM.Native.Plugin.Annotation
import Data.Array.Accelerate.LLVM.Native.Plugin.BuildInfo

import Control.Monad
import Data.IORef
import Data.List
import qualified Data.Map                                           as Map

#if __GLASGOW_HASKELL__ >= 902
import GHC.Driver.Backend
#if __GLASGOW_HASKELL__ < 910
import GHC.Linker
#endif
import GHC.Linker.Loader                                            ( loadCmdLineLibs )
import GHC.Plugins
import GHC.Runtime.Interpreter
#elif __GLASGOW_HASKELL__ >= 900
import GHC.Plugins
import GHC.Runtime.Linker
#else
import GhcPlugins
import Linker
import SysTools
#endif


-- | This GHC plugin is required to support ahead-of-time compilation for the
-- accelerate-llvm-native backend. In particular, it tells GHC about the
-- additional object files generated by
-- 'Data.Array.Accelerate.LLVM.Native.runQ'* which must be linked into the final
-- executable.
--
-- This plugin is automatically installed when using runQ. In older versions of
-- GHC, it was necessary to manually add the plugin using:
--
-- > ghc-options: -fplugin=Data.Array.Accelerate.LLVM.Native.Plugin
--
-- That is no longer needed.

--
plugin :: Plugin
plugin = defaultPlugin
  { installCoreToDos = install
#if __GLASGOW_HASKELL__ >= 806
  , pluginRecompile  = purePlugin
#endif
  }

install :: HasCallStack => [CommandLineOption] -> [CoreToDo] -> CoreM [CoreToDo]
install _ rest = do
  let this (CoreDoPluginPass "accelerate-llvm-native" _) = True
      this _                                             = False
  --
  return $ CoreDoPluginPass "accelerate-llvm-native" pass : filter (not . this) rest

pass :: HasCallStack => ModGuts -> CoreM ModGuts
pass guts = do
  -- Gather annotations for the extra object files which must be supplied to the
  -- linker in order to complete the current module.
  --
  this  <- getModule
  paths <- nub . concat <$> mapM (objectPaths guts) (mg_binds guts)

  unless (null paths)
    $ debugTraceMsg
    $ hang (text "Data.Array.Accelerate.LLVM.Native.Plugin: linking module" <+> quotes (pprModule this) <+> text "with:") 2 (vcat (map text paths))

  -- The linking method depends on the current build target
  -- TODO: Need to update for ghc-8.6: the Backend data type is now abstract
  --
  -- Determine the current build environment
  --
  hscEnv   <- getHscEnv
  dynFlags <- getDynFlags

#if __GLASGOW_HASKELL__ >= 902
  let target = backend dynFlags
#else
  let target = hscTarget dynFlags
#endif

  when (backendGeneratesCode target) $
    if backendWritesFiles target
      then do
        -- The compiler will write files (interface files and object code). This
        -- is true of "real" backends, i.e. not the interpreter.
#if __GLASGOW_HASKELL__ < 806
        -- Because of separate compilation, we will only encounter the annotation
        -- pragmas on files which have changed between invocations. This applies to
        -- both @ghc --make@ as well as the separate compile/link phases of building
        -- with @cabal@ (and @stack@). Note that whenever _any_ file is updated we
        -- must make sure that the linker options contains the complete list of
        -- objects required to build the entire project.
        --

        -- Read the object file index and update (we may have added or removed
        -- objects for the given module)
        --
        let buildInfo = mkBuildInfoFileName (objectMapPath dynFlags)
        abi <- readBuildInfo buildInfo
        --
        let abi'      = if null paths
                          then Map.delete this       abi
                          else Map.insert this paths abi
            allPaths  = nub (concat (Map.elems abi'))
            allObjs   = map optionOfPath allPaths
        --
        writeBuildInfo buildInfo abi'

        -- Make sure the linker flags are up-to-date.
        --
        unless (isNoLink (ghcLink dynFlags)) $ do
          linker_info <- getLinkerInfo dynFlags
          writeIORef (rtldInfo dynFlags)
            $ Just
            $ case linker_info of
                GnuLD     opts -> GnuLD     (nub (opts ++ allObjs))
                GnuGold   opts -> GnuGold   (nub (opts ++ allObjs))
                DarwinLD  opts -> DarwinLD  (nub (opts ++ allObjs))
                SolarisLD opts -> SolarisLD (nub (opts ++ allObjs))
                AixLD     opts -> AixLD     (nub (opts ++ allObjs))
                LlvmLLD   opts -> LlvmLLD   (nub (opts ++ allObjs))
                UnknownLD      -> UnknownLD  -- no linking performed?
#endif
        return ()

      else
        -- We are in interactive mode (ghci)
        --
        unless (null paths) . liftIO $ do
          let opts  = ldInputs dynFlags
              objs  = map optionOfPath paths
          --
#if __GLASGOW_HASKELL__ >= 902
          loadCmdLineLibs (hscInterp hscEnv)
                 $ hscEnv { hsc_dflags = dynFlags { ldInputs = opts ++ objs }}
#else
          linkCmdLineLibs
                 $ hscEnv { hsc_dflags = dynFlags { ldInputs = opts ++ objs }}
#endif
  return guts

#if __GLASGOW_HASKELL__ < 906
backendGeneratesCode :: Backend -> Bool
backendGeneratesCode NoBackend = False
backendGeneratesCode _         = True

backendWritesFiles :: Backend -> Bool
backendWritesFiles Interpreter = False
backendWritesFiles _           = True
#endif

objectPaths :: ModGuts -> CoreBind -> CoreM [FilePath]
objectPaths guts (NonRec b _) = objectAnns guts b
objectPaths guts (Rec bs)     = concat <$> mapM (objectAnns guts . fst) bs

objectAnns :: ModGuts -> CoreBndr -> CoreM [FilePath]
objectAnns guts bndr = do
  anns  <- getAnnotations deserializeWithData guts
#if __GLASGOW_HASKELL__ >= 900
  return [ path | Object path <- lookupWithDefaultUFM (snd anns) [] (varName bndr) ]
#else
  return [ path | Object path <- lookupWithDefaultUFM anns       [] (varUnique bndr) ]
#endif

objectMapPath :: DynFlags -> FilePath
objectMapPath DynFlags{..}
  | Just p <- objectDir = p
  | Just p <- dumpDir   = p
  | otherwise           = "."

optionOfPath :: FilePath -> Option
optionOfPath = FileOption []