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 ] ]