import CoreUnfold ( mkInlineUnfolding )
import CoreFVs
import CoreMonad ( CoreM, getHscEnv )
-import FamInstEnv ( extendFamInstEnvList )
import Var
import Id
import OccName
-- TODO: What new binds do we get back here?
(types', fam_insts, tc_binds) <- vectTypeEnv (mg_types guts)
- -- TODO: What is this?
- let fam_inst_env' = extendFamInstEnvList (mg_fam_inst_env guts) fam_insts
- updGEnv (setFamInstEnv fam_inst_env')
+ (_, fam_inst_env) <- readGEnv global_fam_inst_env
-- dicts <- mapM buildPADict pa_insts
-- workers <- mapM vectDataConWorkers pa_insts
return $ guts { mg_types = types'
, mg_binds = Rec tc_binds : binds'
- , mg_fam_inst_env = fam_inst_env'
+ , 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
+ (inline, _, expr') <- vectTopRhs [] var expr
var' <- vectTopBinder var inline expr'
-- Vectorising the body may create other top-level bindings.
<- 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')
-
+ (inlines', areScalars', exprs')
+ <- mapAndUnzip3M (uncurry $ vectTopRhs vars) bs
+ if (and areScalars') || (length bs <= 1)
+ then do
+ return (vars', inlines', exprs')
+ else do
+ _ <- mapM deleteGlobalScalar vars
+ (inlines'', _, exprs'') <- mapAndUnzip3M (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
return b
where
(vars, exprs) = unzip bs
-
+ mapAndUnzip3M f xs = do
+ ys <- mapM f xs
+ return $ unzip3 ys
-- | Make the vectorised version of this top level binder, and add the mapping
-- between it and the original to the state. For some binder @foo@ the vectorised
-- | Vectorise the RHS of a top-level binding, in an empty local environment.
vectTopRhs
- :: Var -- ^ Name of the binding.
+ :: [Var] -- ^ Names of all functions in the rec block
+ -> Var -- ^ Name of the binding.
-> CoreExpr -- ^ Body of the binding.
- -> VM (Inline, CoreExpr)
+ -> VM (Inline, Bool, CoreExpr)
-vectTopRhs var expr
+vectTopRhs recFs var expr
= dtrace (vcat [text "vectTopRhs", ppr expr])
$ closedV
- $ do (inline, vexpr) <- inBind var
- $ vectPolyExpr (isLoopBreaker $ idOccInfo var)
- (freeVars expr)
- return (inline, vectorised vexpr)
+ $ do (inline, isScalar, vexpr) <-
+ inBind var $ vectPolyExpr (isLoopBreaker $ idOccInfo var) recFs (freeVars expr)
+ if isScalar
+ then addGlobalScalar var
+ else deleteGlobalScalar var
+ return (inline, isScalar, vectorised vexpr)
-- | Project out the vectorised version of a binding from some closure,