vectoriser: refactoring and cleanups for PRepr
authorBen Lippmeier <benl@ouroborus.net>
Mon, 14 Nov 2011 05:48:32 +0000 (16:48 +1100)
committerBen Lippmeier <benl@ouroborus.net>
Mon, 14 Nov 2011 05:48:32 +0000 (16:48 +1100)
compiler/vectorise/Vectorise/Generic/PADict.hs
compiler/vectorise/Vectorise/Generic/PAMethods.hs
compiler/vectorise/Vectorise/Type/Env.hs
compiler/vectorise/Vectorise/Type/PData.hs
compiler/vectorise/Vectorise/Utils/Base.hs

index b4c6931..0af5fe0 100644 (file)
@@ -21,21 +21,11 @@ import Var
 import Name
 
 
--- debug                = False
--- dtrace s x   = if debug then pprTrace "Vectoris.Type.PADict" s x else x
-
 -- |Build the PA dictionary function for some type and hoist it to top level.
 --
--- The PA dictionary holds fns that convert values to and from their vectorised representations.
+--  The PA dictionary holds fns that convert values to and from their vectorised representations.
 --
-buildPADict
-  :: TyCon  -- ^ tycon of the type being vectorised.
-  -> TyCon  -- ^ tycon of the type used for the vectorised representation.
-  -> TyCon  -- ^ PRepr instance tycon
-  -> SumRepr  -- ^ representation used for the type being vectorised.
-  -> VM Var -- ^ name of the top-level dictionary function.
-
--- Recall the definition:
+-- @Recall the definition:
 --    class class PR (PRepr a) => PA a where
 --      toPRepr      :: a -> PRepr a
 --      fromPRepr    :: PRepr a -> a
@@ -50,8 +40,17 @@ buildPADict
 --    $toRepr :: forall a. PA a -> T a -> PRepr (T a)
 --    $toPRepr = ...
 -- The "..." stuff is filled in by buildPAScAndMethods
+-- @
+--
+buildPADict
+        :: TyCon        -- ^ tycon of the type being vectorised.
+        -> TyCon        -- ^ tycon of the type used for the vectorised representation.
+        -> TyCon        -- ^ PData  instance tycon
+        -> TyCon        -- ^ PDatas instance tycon
+        -> SumRepr      -- ^ representation used for the type being vectorised.
+        -> VM Var       -- ^ name of the top-level dictionary function.
 
-buildPADict vect_tc prepr_tc arr_tc repr
+buildPADict vect_tc prepr_tc pdata_tc pdatas_tc repr
  = polyAbstract tvs $ \args ->    -- The args are the dictionaries we lambda
                                   -- abstract over; and they are put in the
                                   -- envt, so when we need a (PA a) we can 
@@ -88,23 +87,21 @@ buildPADict vect_tc prepr_tc arr_tc repr
       ; return dfun
       }
   where
-    tvs       = tyConTyVars vect_tc
-    arg_tys   = mkTyVarTys tvs
-    inst_ty   = mkTyConApp vect_tc arg_tys
-
+    tvs          = tyConTyVars vect_tc
+    arg_tys      = mkTyVarTys tvs
+    inst_ty      = mkTyConApp vect_tc arg_tys
     vect_tc_name = getName vect_tc
 
     method args dfun_name (name, build)
-      = localV
-      $ do
-          expr     <- build vect_tc prepr_tc arr_tc repr
-          let body = mkLams (tvs ++ args) expr
-          raw_var  <- newExportedVar (method_name dfun_name name) (exprType body)
-          let var  = raw_var
+     = localV
+     $ do  expr     <- build vect_tc prepr_tc pdata_tc pdatas_tc repr
+           let body = mkLams (tvs ++ args) expr
+           raw_var  <- newExportedVar (method_name dfun_name name) (exprType body)
+           let var  = raw_var
                       `setIdUnfolding` mkInlineUnfolding (Just (length args)) body
                       `setInlinePragma` alwaysInlinePragma
-          hoistBinding var body
-          return var
+           hoistBinding var body
+           return var
 
     method_call args id        = mkApps (Var id) (map Type arg_tys ++ map Var args)
     method_name dfun_name name = mkVarOcc $ occNameString dfun_name ++ ('$' : name)
index 832c839..a5e999b 100644 (file)
@@ -64,7 +64,16 @@ mk_fam_inst fam_tc arg_tc
 --  Not all lifted backends use the 'toArrPReprs' and 'fromArrPReprs' methods, 
 --  so we only generate these if the 'PDatas' type family is defined.
 --
-buildPAScAndMethods :: VM [( String, TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr)]
+type PAInstanceBuilder
+        =  TyCon        -- ^ Vectorised TyCon 
+        -> TyCon        -- ^ Representation TyCon
+        -> TyCon        -- ^ 'PData'  TyCon
+        -> TyCon        -- ^ 'PDatas' TyCon
+        -> SumRepr      -- ^ Description of generic representation.
+        -> VM CoreExpr  -- ^ Instance function.
+
+
+buildPAScAndMethods :: VM [(String, PAInstanceBuilder)]
 buildPAScAndMethods
  = do   hasPDatas <- liftM isJust $ builtin pdatasTyCon
         return 
@@ -76,12 +85,11 @@ buildPAScAndMethods
          ++ (if hasPDatas then
               [ ("toArrPReprs",   buildToArrPReprs)
               , ("fromArrPReprs", buildFromArrPReprs)]
-              else [])
-             
+              else [])             
 
 
-buildPRDict :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
-buildPRDict vect_tc prepr_tc _ _
+buildPRDict :: PAInstanceBuilder
+buildPRDict vect_tc prepr_tc _ _ _
   = prDictOfPReprInstTyCon inst_ty prepr_tc arg_tys
   where
     arg_tys = mkTyVarTys (tyConTyVars vect_tc)
@@ -90,8 +98,8 @@ buildPRDict vect_tc prepr_tc _ _
 
 -- buildToPRepr ---------------------------------------------------------------
 -- | Build the 'toRepr' method of the PA class.
-buildToPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
-buildToPRepr vect_tc repr_tc _ repr
+buildToPRepr :: PAInstanceBuilder
+buildToPRepr vect_tc repr_tc _ repr
  = do let arg_ty = mkTyConApp vect_tc ty_args
 
       -- Get the representation type of the argument.
@@ -164,8 +172,8 @@ buildToPRepr vect_tc repr_tc _ repr
 
 -- buildFromPRepr -------------------------------------------------------------
 -- | Build the 'fromPRepr' method of the PA class.
-buildFromPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
-buildFromPRepr vect_tc repr_tc _ repr
+buildFromPRepr :: PAInstanceBuilder
+buildFromPRepr vect_tc repr_tc _ repr
   = do
       arg_ty <- mkPReprType res_ty
       arg <- newLocalVar (fsLit "x") arg_ty
@@ -218,8 +226,8 @@ buildFromPRepr vect_tc repr_tc _ repr
 
 -- buildToArrRepr -------------------------------------------------------------
 -- | Build the 'toArrRepr' method of the PA class.
-buildToArrPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
-buildToArrPRepr vect_tc prepr_tc pdata_tc r
+buildToArrPRepr :: PAInstanceBuilder
+buildToArrPRepr vect_tc prepr_tc pdata_tc r
  = do arg_ty <- mkPDataType el_ty
       res_ty <- mkPDataType =<< mkPReprType el_ty
       arg    <- newLocalVar (fsLit "xs") arg_ty
@@ -301,8 +309,8 @@ buildToArrPRepr vect_tc prepr_tc pdata_tc r
 
 -- buildFromArrPRepr ----------------------------------------------------------
 -- | Build the 'fromArrPRepr' method for the PA class.
-buildFromArrPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
-buildFromArrPRepr vect_tc prepr_tc pdata_tc r
+buildFromArrPRepr :: PAInstanceBuilder
+buildFromArrPRepr vect_tc prepr_tc pdata_tc r
  = do arg_ty <- mkPDataType =<< mkPReprType el_ty
       res_ty <- mkPDataType el_ty
       arg    <- newLocalVar (fsLit "xs") arg_ty
@@ -378,8 +386,54 @@ buildFromArrPRepr vect_tc prepr_tc pdata_tc r
          = do (res', args') <- f res_ty res expr r
               return (res', args' ++ args)
 
--- buildToArrPReprs -----------------------------------------------------------
-buildToArrPReprs        = error "buildToArrPReprs not done yet"
 
+-- buildToArrPReprs -----------------------------------------------------------
+-- | Build the 'toArrPReprs' instance for the PA class.
+--   This converts a PData of elements into the generic representation.
+buildToArrPReprs :: PAInstanceBuilder
+buildToArrPReprs vect_tc prepr_tc _ pdatas_tc r
+ = do
+        -- The element type of the argument.
+        --  eg: 'Tree a b'.
+        let ty_args = mkTyVarTys $ tyConTyVars vect_tc
+        let el_ty   = mkTyConApp vect_tc ty_args
+
+        -- The argument type of the instance.
+        --  eg: 'PDatas (Tree a b)'
+        Just arg_ty      <- mkPDatasType el_ty
+
+        -- The result type. 
+        --  eg: 'PDatas (PRepr (Tree a b))'
+        Just res_ty      <- mkPDatasType =<< mkPReprType el_ty
+        
+        -- Variable to bind the argument to the instance
+        -- eg: (xss :: PDatas (Tree a b))
+        varg         <- newLocalVar (fsLit "xss") arg_ty
+        
+        return  $ Lam varg (Var varg)
+
+        
 -- buildFromArrPReprs ---------------------------------------------------------
-buildFromArrPReprs      = error "buildFromArrPReprs not done yet"
\ No newline at end of file
+buildFromArrPReprs :: PAInstanceBuilder
+buildFromArrPReprs vect_tc prepr_tc _ pdatas_tc r
+ = do   
+        -- The element type of the argument.
+        --  eg: 'Tree a b'.
+        let ty_args = mkTyVarTys $ tyConTyVars vect_tc
+        let el_ty   = mkTyConApp vect_tc ty_args
+
+        -- The argument type of the instance.
+        --  eg: 'PDatas (PRepr (Tree a b))'
+        Just arg_ty      <- mkPDatasType =<< mkPReprType el_ty
+
+        -- The result type. 
+        --  eg: 'PDatas (Tree a b)'
+        Just res_ty      <- mkPDatasType el_ty
+        
+        -- Variable to bind the argument to the instance
+        -- eg: (xss :: PDatas (PRepr (Tree a b)))
+        varg         <- newLocalVar (fsLit "xss") arg_ty
+        
+        return  $ Lam varg (Var varg)
+
+
index 5f76295..efc84fa 100644 (file)
@@ -204,9 +204,10 @@ vectTypeEnv tycons vectTypeDecls vectClassDecls
 
            -- Build 'PRepr' and 'PData' instance type constructors and family instances for all
            -- type constructors with vectorised representations.
-       ; reprs     <- mapM tyConRepr vect_tcs
-       ; repr_tcs  <- zipWith3M buildPReprTyCon orig_tcs vect_tcs reprs
-       ; pdata_tcs <- zipWith3M buildPDataTyCon orig_tcs vect_tcs reprs
+       ; reprs      <- mapM tyConRepr vect_tcs
+       ; repr_tcs   <- zipWith3M buildPReprTyCon  orig_tcs vect_tcs reprs
+       ; pdata_tcs  <- zipWith3M buildPDataTyCon  orig_tcs vect_tcs reprs
+       ; pdatas_tcs <- zipWith3M buildPDatasTyCon orig_tcs vect_tcs reprs
        ; let inst_tcs  = repr_tcs ++ pdata_tcs
              fam_insts = map mkLocalFamInst inst_tcs
        ; updGEnv $ extendFamEnv fam_insts
@@ -217,11 +218,12 @@ vectTypeEnv tycons vectTypeDecls vectClassDecls
        ; (_, binds) <- fixV $ \ ~(dfuns, _) ->
            do { defTyConPAs (zipLazy vect_tcs dfuns)
               ; dfuns <- sequence 
-                      $  zipWith4 buildTyConBindings
+                      $  zipWith5 buildTyConBindings
                                   orig_tcs
                                   vect_tcs
                                   repr_tcs
                                   pdata_tcs
+                                  pdatas_tcs
 
               ; binds <- takeHoisted
               ; return (dfuns, binds)
@@ -233,14 +235,13 @@ vectTypeEnv tycons vectTypeDecls vectClassDecls
        }
 
 
--- Helpers -------------------
-
-buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> VM Var
-buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc
- = do { vectDataConWorkers orig_tc vect_tc pdata_tc
-      ; repr <- tyConRepr vect_tc
-      ; buildPADict vect_tc prepr_tc pdata_tc repr
-      }
+-- Helpers --------------------------------------------------------------------
+buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> TyCon -> VM Var
+buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc pdatas_tc
+ = do   vectDataConWorkers orig_tc vect_tc pdata_tc
+        repr <- tyConRepr vect_tc
+        buildPADict vect_tc prepr_tc pdata_tc pdatas_tc repr
+      
 
 vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
 vectDataConWorkers orig_tc vect_tc arr_tc
index 7f37929..cbc74f5 100644 (file)
@@ -1,7 +1,7 @@
 
 module Vectorise.Type.PData
   ( buildPDataTyCon
-  ) 
+  , buildPDatasTyCon 
 where
 
 import Vectorise.Monad
@@ -20,10 +20,13 @@ import MonadUtils
 import Control.Monad
 
 
+buildPDatasTyCon :: TyCon -> TyCon -> SumRepr -> VM TyCon
+buildPDatasTyCon = buildPDataTyCon -- error "buildPDatasTyCon: not finished"
+
 buildPDataTyCon :: TyCon -> TyCon -> SumRepr -> VM TyCon
-buildPDataTyCon orig_tc vect_tc repr = fixV $ \repr_tc ->
-  do
   name' <- mkLocalisedName mkPDataTyConOcc orig_name
+buildPDataTyCon orig_tc vect_tc repr 
+ = fixV $ \repr_tc ->
do name' <- mkLocalisedName mkPDataTyConOcc orig_name
     rhs   <- buildPDataTyConRhs orig_name vect_tc repr_tc repr
     pdata <- builtin pdataTyCon
 
@@ -35,22 +38,20 @@ buildPDataTyCon orig_tc vect_tc repr = fixV $ \repr_tc ->
                            False       -- not GADT syntax
                            NoParentTyCon
                            (Just $ mk_fam_inst pdata vect_tc)
 where
+ where
     orig_name = tyConName orig_tc
-    tyvars = tyConTyVars vect_tc
-    rec_flag = boolToRecFlag (isRecursiveTyCon vect_tc)
+    tyvars    = tyConTyVars vect_tc
+    rec_flag  = boolToRecFlag (isRecursiveTyCon vect_tc)
 
 
 buildPDataTyConRhs :: Name -> TyCon -> TyCon -> SumRepr -> VM AlgTyConRhs
 buildPDataTyConRhs orig_name vect_tc repr_tc repr
-  = do
-      data_con <- buildPDataDataCon orig_name vect_tc repr_tc repr
+ = do data_con <- buildPDataDataCon orig_name vect_tc repr_tc repr
       return $ DataTyCon { data_cons = [data_con], is_enum = False }
 
 buildPDataDataCon :: Name -> TyCon -> TyCon -> SumRepr -> VM DataCon
 buildPDataDataCon orig_name vect_tc repr_tc repr
-  = do
-      dc_name  <- mkLocalisedName mkPDataDataConOcc orig_name
+ = do dc_name  <- mkLocalisedName mkPDataDataConOcc orig_name
       comp_tys <- sum_tys repr
 
       liftDs $ buildDataCon dc_name
@@ -64,10 +65,10 @@ buildPDataDataCon orig_name vect_tc repr_tc repr
                             comp_tys
                             (mkFamilyTyConApp repr_tc (mkTyVarTys tvs))
                             repr_tc
 where
+ where
     tvs   = tyConTyVars vect_tc
 
-    sum_tys EmptySum = return []
+    sum_tys EmptySum     = return []
     sum_tys (UnarySum r) = con_tys r
     sum_tys (Sum { repr_sel_ty = sel_ty
                  , repr_cons   = cons })
@@ -75,7 +76,7 @@ buildPDataDataCon orig_name vect_tc repr_tc repr
 
     con_tys (ConRepr _ r) = prod_tys r
 
-    prod_tys EmptyProd = return []
+    prod_tys EmptyProd     = return []
     prod_tys (UnaryProd r) = liftM singleton (comp_ty r)
     prod_tys (Prod { repr_comps = comps }) = mapM comp_ty comps
 
index f2b3f80..6472253 100644 (file)
@@ -11,7 +11,7 @@ module Vectorise.Utils.Base (
   mkPReprType,
   mkPArrayType, splitPrimTyCon,
   mkPArray,
-  mkPDataType,
+  mkPDataType,  mkPDatasType,
   mkBuiltinCo,
   mkVScrut,
 
@@ -37,7 +37,7 @@ import Outputable
 import FastString
 
 import Control.Monad (liftM)
-
+import Data.Maybe
 
 -- Simple Types ---------------------------------------------------------------
 voidType :: VM Type
@@ -67,36 +67,38 @@ dataConTagZ :: DataCon -> Int
 dataConTagZ con = dataConTag con - fIRST_TAG
 
 
+-- Type Construction ----------------------------------------------------------
+-- | Make an application of a builtin type constructor to some arguments.
 mkBuiltinTyConApp :: (Builtins -> TyCon) -> [Type] -> VM Type
 mkBuiltinTyConApp get_tc tys
-  = do
-      tc <- builtin get_tc
+ = do tc     <- builtin get_tc
       return $ mkTyConApp tc tys
 
 
 mkBuiltinTyConApps :: (Builtins -> TyCon) -> [Type] -> Type -> VM Type
 mkBuiltinTyConApps get_tc tys ty
-  = do
-      tc <- builtin get_tc
+ = do tc     <- builtin get_tc
       return $ foldr (mk tc) ty tys
   where
     mk tc ty1 ty2 = mkTyConApp tc [ty1,ty2]
 
 
+-- | Make an application of the 'Wrap' type constructor.
 mkWrapType :: Type -> VM Type
-mkWrapType ty = mkBuiltinTyConApp wrapTyCon [ty]
+mkWrapType ty  = mkBuiltinTyConApp wrapTyCon [ty]
 
 
+-- | Make an application of the closure type constructor.
 mkClosureTypes :: [Type] -> Type -> VM Type
 mkClosureTypes = mkBuiltinTyConApps closureTyCon
 
 
+-- | Make an application of the 'PRepr' type constructor.
 mkPReprType :: Type -> VM Type
 mkPReprType ty = mkBuiltinTyConApp preprTyCon [ty]
 
 
--- |Wrap a type into 'PArray', treating unboxed types specially.
---
+-- | Wrap a type into 'PArray', treating unboxed types specially.
 mkPArrayType :: Type -> VM Type
 mkPArrayType ty
   | Just tycon <- splitPrimTyCon ty
@@ -105,8 +107,23 @@ mkPArrayType ty
        }
 mkPArrayType ty = mkBuiltinTyConApp parrayTyCon [ty]
 
+
+-- | Make an appliction of the 'PData' tycon to some argument.
+mkPDataType :: Type -> VM Type
+mkPDataType ty
+        = mkBuiltinTyConApp pdataTyCon [ty]
+
+
+-- | Make an application of the 'PDatas' tycon to some argument.
+mkPDatasType :: Type -> VM (Maybe Type)
+mkPDatasType ty
+ = do   mtc      <- builtin pdatasTyCon
+        case mtc of
+         Nothing        -> return Nothing
+         Just tc'       -> return $ Just $ mkTyConApp tc' [ty]
+
+
 -- |Checks if a type constructor is defined in 'GHC.Prim' (e.g., 'Int#'); if so, returns it.
---
 splitPrimTyCon :: Type -> Maybe TyCon
 splitPrimTyCon ty
   | Just (tycon, []) <- splitTyConApp_maybe ty
@@ -115,22 +132,30 @@ splitPrimTyCon ty
   | otherwise = Nothing
 
 
-------
-mkPArray :: Type -> CoreExpr -> CoreExpr -> VM CoreExpr
-mkPArray ty len dat = do
-                        tc <- builtin parrayTyCon
-                        let [dc] = tyConDataCons tc
-                        return $ mkConApp dc [Type ty, len, dat]
 
-mkPDataType :: Type -> VM Type
-mkPDataType ty = mkBuiltinTyConApp pdataTyCon [ty]
+-- CoreExpr Construction ------------------------------------------------------
+-- | Make an application of the 'PArray' data constructor.
+mkPArray 
+        :: Type         -- ^ Element type
+        -> CoreExpr     -- ^ 'Int'   for the array length.
+        -> CoreExpr     -- ^ 'PData' for the array data.
+        -> VM CoreExpr
 
+mkPArray ty len dat 
+ = do   tc <- builtin parrayTyCon
+        let [dc] = tyConDataCons tc
+        return $ mkConApp dc [Type ty, len, dat]
+
+
+-- Coercion Construction -----------------------------------------------------
+-- | Make a coersion to some builtin type.
 mkBuiltinCo :: (Builtins -> TyCon) -> VM Coercion
 mkBuiltinCo get_tc
-  = do
-      tc <- builtin get_tc
+ = do tc     <- builtin get_tc
       return $ mkTyConAppCo tc []
 
+
+-------------------------------------------------------------------------------
 mkVScrut :: VExpr -> VM (CoreExpr, CoreExpr, TyCon, [Type])
 mkVScrut (ve, le)
   = do