module Vectorise( vectorise )
where
import Vectorise.Type.Env
import Vectorise.Type.Type
import Vectorise.Convert
import Vectorise.Utils.Hoisting
import Vectorise.Exp
import Vectorise.Vect
import Vectorise.Env
import Vectorise.Monad
import HscTypes hiding ( MonadThings(..) )
import Module ( PackageId )
import CoreSyn
import CoreUnfold ( mkInlineUnfolding )
import CoreFVs
import CoreMonad ( CoreM, getHscEnv )
import FamInstEnv ( extendFamInstEnvList )
import Var
import Id
import OccName
import BasicTypes ( isLoopBreaker )
import Outputable
import Util ( zipLazy )
import Control.Monad
debug = False
dtrace s x = if debug then pprTrace "Vectorise" s x else x
vectorise :: PackageId -> ModGuts -> CoreM ModGuts
vectorise backend guts
= do hsc_env <- getHscEnv
liftIO $ vectoriseIO backend hsc_env guts
vectoriseIO :: PackageId -> HscEnv -> ModGuts -> IO ModGuts
vectoriseIO backend hsc_env guts
= do
eps <- hscEPS hsc_env
let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
Just (info', guts') <- initV backend hsc_env guts info (vectModule guts)
return (guts' { mg_vect_info = info' })
vectModule :: ModGuts -> VM ModGuts
vectModule guts
= do
(types', fam_insts, tc_binds) <- vectTypeEnv (mg_types guts)
let fam_inst_env' = extendFamInstEnvList (mg_fam_inst_env guts) fam_insts
updGEnv (setFamInstEnv fam_inst_env')
binds' <- mapM vectTopBind (mg_binds guts)
return $ guts { mg_types = types'
, mg_binds = Rec tc_binds : binds'
, mg_fam_inst_env = fam_inst_env'
, mg_fam_insts = mg_fam_insts guts ++ fam_insts
}
vectTopBind :: CoreBind -> VM CoreBind
vectTopBind b@(NonRec var expr)
= do
(inline, expr') <- vectTopRhs var expr
var' <- vectTopBinder var inline expr'
hs <- takeHoisted
cexpr <- tryConvert var var' expr
return . Rec $ (var, cexpr) : (var', expr') : hs
`orElseV`
return b
vectTopBind b@(Rec bs)
= do
(vars', _, exprs')
<- fixV $ \ ~(_, inlines, rhss) ->
do vars' <- sequence [vectTopBinder var inline rhs
| (var, ~(inline, rhs)) <- zipLazy vars (zip inlines rhss)]
(inlines', exprs')
<- mapAndUnzipM (uncurry vectTopRhs) bs
return (vars', inlines', exprs')
hs <- takeHoisted
cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
`orElseV`
return b
where
(vars, exprs) = unzip bs
vectTopBinder
:: Var
-> Inline
-> CoreExpr
-> VM Var
vectTopBinder var inline expr
= do
vty <- vectType (idType var)
var' <- liftM (`setIdUnfoldingLazily` unfolding)
$ cloneId mkVectOcc var vty
defGlobalVar var var'
return var'
where
unfolding = case inline of
Inline arity -> mkInlineUnfolding (Just arity) expr
DontInline -> noUnfolding
vectTopRhs
:: Var
-> CoreExpr
-> VM (Inline, CoreExpr)
vectTopRhs var expr
= dtrace (vcat [text "vectTopRhs", ppr expr])
$ closedV
$ do (inline, vexpr) <- inBind var
$ vectPolyExpr (isLoopBreaker $ idOccInfo var)
(freeVars expr)
return (inline, vectorised vexpr)
tryConvert
:: Var
-> Var
-> CoreExpr
-> VM CoreExpr
tryConvert var vect_var rhs
= fromVect (idType var) (Var vect_var) `orElseV` return rhs