Simple conversion vectorised -> unvectorised

2 files changed, 47 insertions, 3 deletions

diff --git a/compiler/vectorise/VectType.hs b/compiler/vectorise/VectType.hs
index 781131e66c..912eacfad7 100644
--- a/compiler/vectorise/VectType.hs
+++ b/compiler/vectorise/VectType.hs
@@ -7,7 +7,8 @@
 
 module VectType ( vectTyCon, vectType, vectTypeEnv,
                   mkRepr, arrShapeTys, arrShapeVars, arrSelector,
-                  PAInstance, buildPADict )
+                  PAInstance, buildPADict,
+                  fromVect )
 where
 
 #include "HsVersions.h"
@@ -982,3 +983,40 @@ tyConsOfType other             = pprPanic "ClosureConv.tyConsOfType" $ ppr other
 tyConsOfTypes :: [Type] -> UniqSet TyCon
 tyConsOfTypes = unionManyUniqSets . map tyConsOfType
 
+
+-- ----------------------------------------------------------------------------
+-- Conversions
+
+fromVect :: Type -> CoreExpr -> VM CoreExpr
+fromVect ty expr | Just ty' <- coreView ty = fromVect ty' expr
+fromVect (FunTy arg_ty res_ty) expr
+  = do
+      arg     <- newLocalVar FSLIT("x") arg_ty
+      varg    <- toVect arg_ty (Var arg)
+      varg_ty <- vectType arg_ty
+      vres_ty <- vectType res_ty
+      apply   <- builtin applyClosureVar
+      body    <- fromVect res_ty
+               $ Var apply `mkTyApps` [arg_ty, res_ty] `mkApps` [expr, Var arg]
+      return $ Lam arg body
+fromVect ty expr
+  = identityConv ty >> return expr
+
+toVect :: Type -> CoreExpr -> VM CoreExpr
+toVect ty expr = identityConv ty >> return expr
+
+identityConv :: Type -> VM ()
+identityConv ty | Just ty' <- coreView ty = identityConv ty'
+identityConv (TyConApp tycon tys)
+  = do
+      mapM_ identityConv tys
+      identityConvTyCon tycon
+identityConv ty = noV
+
+identityConvTyCon :: TyCon -> VM ()
+identityConvTyCon tc
+  | isBoxedTupleTyCon tc = return ()
+  | isUnLiftedTyCon tc   = return ()
+  | otherwise            = maybeV (lookupTyCon tc) >> return ()
+
+
diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index 9f2e2b7b77..63575b90f1 100644
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -89,7 +89,8 @@ vectTopBind b@(NonRec var expr)
       var'  <- vectTopBinder var
       expr' <- vectTopRhs var expr
       hs    <- takeHoisted
-      return . Rec $ (var, expr) : (var', expr') : hs
+      cexpr <- tryConvert var var' expr
+      return . Rec $ (var, cexpr) : (var', expr') : hs
   `orElseV`
     return b
 
@@ -98,7 +99,8 @@ vectTopBind b@(Rec bs)
       vars'  <- mapM vectTopBinder vars
       exprs' <- zipWithM vectTopRhs vars exprs
       hs     <- takeHoisted
-      return . Rec $ bs ++ zip vars' exprs' ++ hs
+      cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
+      return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
   `orElseV`
     return b
   where
@@ -119,6 +121,10 @@ vectTopRhs var expr
               . inBind var
               $ vectPolyExpr (freeVars expr)
 
+tryConvert :: Var -> Var -> CoreExpr -> VM CoreExpr
+tryConvert var vect_var rhs
+  = fromVect (idType var) (Var vect_var) `orElseV` return rhs
+
 -- ----------------------------------------------------------------------------
 -- Bindings