repa-plugin-1.0.0.1: Data/Array/Repa/Plugin/Pass/Lower.hs
module Data.Array.Repa.Plugin.Pass.Lower
(passLower)
where
import Data.Array.Repa.Plugin.Primitives
import Data.Array.Repa.Plugin.ToDDC.Detect
import Data.Array.Repa.Plugin.ToDDC
import Data.Array.Repa.Plugin.ToGHC
import Data.Array.Repa.Plugin.GHC.Pretty
import DDC.Core.Exp
import qualified DDC.Core.Flow as Flow
import qualified DDC.Core.Flow.Profile as Flow
import qualified DDC.Core.Flow.Transform.Prep as Flow
import qualified DDC.Core.Flow.Transform.Slurp as Flow
import qualified DDC.Core.Flow.Transform.Schedule as Flow
import qualified DDC.Core.Flow.Transform.Extract as Flow
import qualified DDC.Core.Flow.Transform.Concretize as Flow
import qualified DDC.Core.Flow.Transform.Thread as Flow
import qualified DDC.Core.Flow.Transform.Wind as Flow
import qualified DDC.Core.Module as Core
import qualified DDC.Core.Check as Core
import qualified DDC.Core.Simplifier as Core
import qualified DDC.Core.Fragment as Core
import qualified DDC.Core.Transform.Namify as Core
import qualified DDC.Core.Transform.Flatten as Core
import qualified DDC.Core.Transform.Forward as Forward
import qualified DDC.Core.Transform.Thread as Core
import qualified DDC.Core.Transform.Reannotate as Core
import qualified DDC.Core.Transform.Snip as Snip
import qualified DDC.Core.Transform.Eta as Eta
import qualified HscTypes as G
import qualified CoreMonad as G
import qualified UniqSupply as G
import qualified DDC.Base.Pretty as D
import qualified Data.Map as Map
import System.IO.Unsafe
import Control.Monad.State.Strict
import Data.List
-- | We use this unique when generating fresh names.
--
-- If this is not actually unique relative to the rest of the compiler
-- then we're completely screwed.
--
-- GHC doesn't seem to have an API to generate unique prefixes.
--
letsHopeThisIsUnique :: Char
letsHopeThisIsUnique = 's'
-- | Run the lowering pass on this module.
passLower :: [G.CommandLineOption] -> String -> G.ModGuts -> G.CoreM G.ModGuts
passLower options name guts0
= unsafePerformIO
$ do
-- Here's hoping this is really unique
us <- G.mkSplitUniqSupply letsHopeThisIsUnique
-- Decide whether to dump intermediate files
let shouldDump = elem "dump" options
let dump thing str = when shouldDump
$ writeFile ("dump." ++ name ++ "." ++ thing) str
-- Input ------------------------------------------
-- Dump the GHC core code that we start with.
dump "01-ghc.hs"
$ D.renderIndent (pprModGuts guts0)
-- Primitives -------------------------------------
-- Build a table of expressions to access our primitives.
let (Just (primitives, guts), us2)
= G.initUs us (slurpPrimitives guts0)
-- Convert ----------------------------------------
-- Convert the GHC Core module to Disciple Core.
let (mm_dc, failsConvert) = convertModGuts guts
dump "02-raw.dcf"
$ D.renderIndent (D.ppr mm_dc)
dump "02-raw.fails"
$ D.renderIndent (D.vcat $ intersperse D.empty $ map D.ppr failsConvert)
-- Detect -----------------------------------------
-- Detect flow operators and primitives.
-- We also get a map of DDC names to GHC names
let (mm_detect, names) = detectModule mm_dc
dump "03-detect.dcf"
$ D.renderIndent (D.ppr mm_detect)
dump "03-detect.names"
$ D.renderIndent (D.vcat $ map D.ppr $ Map.toList names)
-- Norm -------------------------------------------
-- Eta expand everything so we have names for parameters.
let etaConfig = Eta.configZero { Eta.configExpand = True }
let mm_eta = Core.result $ Eta.etaModule etaConfig Flow.profile mm_detect
dump "04-norm.1-eta.dcf"
$ D.renderIndent (D.ppr mm_eta)
-- A-normalize module for the Prep transform.
let mkNamT = Core.makeNamifier Flow.freshT
let mkNamX = Core.makeNamifier Flow.freshX
-- Snip and flatten the code to create new let-bindings
-- for flow combinators. This ensures all the flow combinators
-- and workers are bound at the top-level of the function.
let snipConfig = Snip.configZero { Snip.configSnipLetBody = True }
let mm_snip' = Core.flatten $ Snip.snip snipConfig mm_eta
let mm_snip = evalState (Core.namifyUnique mkNamT mkNamX mm_snip') 0
dump "04-norm.dcf"
$ D.renderIndent (D.ppr mm_snip)
-- Prep -------------------------------------------
-- 1. Eta-expand worker functions passed to flow combinators.
-- We also get back a map containing the types of parameters
-- to worker functions.
-- NOTE: We're not using the module result of prep now that
-- we have real eta-expansion.
let (_, workerNameArgs)
= Flow.prepModule mm_snip
-- 2. Move worker functions forward so they are directly
-- applied to flow combinators.
let isFloatable lts
= case lts of
LLet (BName n _) _
| Just{} <- Map.lookup n workerNameArgs
-> Forward.FloatForce
_ -> Forward.FloatAllow
let config = Forward.Config isFloatable False
let result_forward = Forward.forwardModule Flow.profile config mm_snip
let mm_forward = Core.result result_forward
dump "05-prep.1-forward.dcf"
$ D.renderIndent (D.ppr mm_forward)
-- 3. Create fresh names for anonymous binders.
-- The lowering pass needs them all to have real names.
let mm_namify = evalState (Core.namifyUnique mkNamT mkNamX mm_forward) 0
dump "05-prep.2-namify.dcf"
$ D.renderIndent (D.ppr mm_namify)
-- 4. Type check add type annots on all binders.
let mm_prep = checkFlowModule_ mm_namify
dump "05-prep.3-check.dcf"
$ D.renderIndent (D.ppr mm_prep)
-- Lower ------------------------------------------
-- Slurp out flow processes from the preped module.
let processes = Flow.slurpProcesses mm_prep
-- Schedule processes into abstract loops.
let procs = map Flow.scheduleProcess processes
-- Extract concrete code from the abstract loops.
let mm_lowered' = Flow.extractModule mm_prep procs
let mm_lowered = evalState (Core.namifyUnique mkNamT mkNamX mm_lowered') 0
dump "06-lowered.1-processes.txt"
$ D.renderIndent (D.vcat $ intersperse D.empty $ map D.ppr $ processes)
dump "06-lowered.dcf"
$ D.renderIndent (D.ppr mm_lowered)
-- Concretize ------------------------------------
-- Concretize rate variables.
let mm_concrete = Flow.concretizeModule mm_lowered
dump "07-concrete.dcf"
$ D.renderIndent (D.ppr mm_concrete)
-- Wind ------------------------------------------
-- Convert uses of the loop# and guard# combinator to real tail-recursive
-- loops.
let mm_wind = Core.result $ Forward.forwardModule Flow.profile
(Forward.Config (const Forward.FloatAllow) True)
$ Core.result $ Forward.forwardModule Flow.profile
(Forward.Config (const Forward.FloatAllow) True)
$ Core.result $ Forward.forwardModule Flow.profile
(Forward.Config (const Forward.FloatAllow) True)
$ Core.result $ Forward.forwardModule Flow.profile
(Forward.Config (const Forward.FloatAllow) True)
$ Flow.windModule mm_concrete
dump "08-wind.dcf"
$ D.renderIndent (D.ppr mm_wind)
-- Check -----------------------------------------
-- Type check the module,
-- the thread transform wants type annotations at each node.
let mm_checked = checkFlowModule mm_wind
dump "09-checked.dcf"
$ D.renderIndent (D.ppr mm_checked)
-- Thread -----------------------------------------
-- Thread the World# token through stateful functions in preparation
-- for conversion back to GHC core.
let mm_thread' = Core.thread Flow.threadConfig
(Core.profilePrimKinds Flow.profile)
(Core.profilePrimTypes Flow.profile)
mm_checked
let mm_thread = evalState (Core.namifyUnique mkNamT mkNamX mm_thread') 0
dump "10-threaded.dcf"
$ D.renderIndent (D.ppr mm_thread)
-- Splice -----------------------------------------
-- Splice the lowered functions back into the GHC core program.
let guts' = G.initUs_ us2
$ spliceModGuts primitives names mm_thread guts
dump "11-spliced.fc"
$ D.renderIndent (pprModGuts guts')
return (return guts')
-- | Type check a Core Flow module
checkFlowModule_
:: Core.Module () Flow.Name
-> Core.Module () Flow.Name
checkFlowModule_ mm
= Core.reannotate Core.annotTail
$ checkFlowModule mm
-- | Type check a Core Flow module, producing type annotations on every node.
checkFlowModule
:: Core.Module () Flow.Name
-> Core.Module (Core.AnTEC () Flow.Name) Flow.Name
checkFlowModule mm
= let result = Core.checkModule
(Core.configOfProfile Flow.profile)
mm
in case result of
Right mm' -> mm'
Left err
-> error $ D.renderIndent $ D.indent 8 $ D.vcat
[ D.empty
, D.text "repa-plugin:"
, D.indent 2
$ D.vcat [ D.text "Type error in generated code"
, D.ppr err ] ]