Add Trace

src/Data/Array/Accelerate/AST/IdxSet.hs

@@ -52,6 +52,12 @@ instance Eq (IdxSet env) where
   IdxSet (PNone a) == IdxSet (PNone b) = IdxSet a == IdxSet b
   _ == _ = False
 
+instance Semigroup (IdxSet env) where
+  (<>) = union
+
+instance Monoid (IdxSet env) where
+  mempty = empty
+
 overlaps :: IdxSet env -> IdxSet env -> Bool
 overlaps (IdxSet (PPush _ _)) (IdxSet (PPush _ _)) = True
 overlaps (IdxSet (PPush a _)) (IdxSet (PNone b  )) = overlaps (IdxSet a) (IdxSet b)

src/Data/Array/Accelerate/AST/Operation.hs

@@ -39,22 +39,23 @@
 
   Var', Exp', Fun', In, Out, Mut,
 
-  ArrayDescriptor(..), weakenArrayDescriptor,
+  ArrayDescriptor(..), ArrayDescriptors, weakenArrayDescriptor, weakenArrayDescriptors,
 
   OpenExp, OpenFun, Exp, Fun, ArrayInstr(..),
   expGroundVars, funGroundVars, arrayInstrsInExp, arrayInstrsInFun,
 
-  encodeGroundR, encodeGroundsR, encodeGroundVar, encodeGroundVars,
+  encodeGroundR, encodeGroundsR, encodeGroundVar, encodeGroundVars, encodeArrayDescriptors,
   rnfGroundR, rnfGroundsR, rnfGroundVar, rnfGroundVars, rnfUniqueness,
   liftGroundR, liftGroundsR, liftGroundVar, liftGroundVars,
 
   bufferImpossible, groundFunctionImpossible,
 
   paramIn, paramsIn, paramIn', paramsIn',
 
-  ReindexPartial, reindexArg, reindexArgs, reindexExp, reindexPreArgs, reindexVar, reindexVars,
+  ReindexPartial, reindexArg, reindexArgs, reindexExp, reindexPreArgs, reindexVar, reindexVars, reindexArrayDescriptors,
   reindexIdxSet,
   weakenReindex,
+  arrayDescriptorsIdxSet,
   argVars, argsVars, argsInputs, argsOutputs, AccessGroundR(..),
 
   mapAccExecutable, mapAfunExecutable,
@@ -187,7 +188,7 @@
           -> PreOpenAcc  op env a
 
   Atrace :: Text
-         -> TupR (ArrayDescriptor env) t
+         -> ArrayDescriptors env t
          -> PreOpenAcc op env Word8
 
   -- | Asserts that the given expression evaluates to true.
@@ -295,9 +296,26 @@
                   -> GroundVars env (Buffers e)
                   -> ArrayDescriptor env (Array sh e)
 
+type ArrayDescriptors env = TupR (ArrayDescriptor env)
+
+weakenArrayDescriptors :: env :> env' -> ArrayDescriptors env a -> ArrayDescriptors env' a
+weakenArrayDescriptors k = mapTupR (weakenArrayDescriptor k)
+
 weakenArrayDescriptor :: env :> env' -> ArrayDescriptor env a -> ArrayDescriptor env' a
 weakenArrayDescriptor k (ArrayDescriptor shr sh buffers) = ArrayDescriptor shr (weakenVars k sh) (weakenVars k buffers)
 
+arrayDescriptorIdxSet :: ArrayDescriptor env t -> IdxSet env
+arrayDescriptorIdxSet (ArrayDescriptor _ sh buffers) = IdxSet.fromVars sh `IdxSet.union` IdxSet.fromVars buffers
+
+arrayDescriptorsIdxSet :: ArrayDescriptors env t -> IdxSet env
+arrayDescriptorsIdxSet = foldMapTupR arrayDescriptorIdxSet
+
+encodeArrayDescriptors :: ArrayDescriptors env t -> Builder
+encodeArrayDescriptors = encodeTupR encodeArrayDescriptor
+
+encodeArrayDescriptor :: ArrayDescriptor env t -> Builder
+encodeArrayDescriptor (ArrayDescriptor shape sh buffer) = encodeShapeR shape <> encodeGroundVars sh <> encodeGroundVars buffer
+
 -- | The arguments to be passed to an operation of type `t`.
 -- This type is represented as a cons list, separated by (->) and ending
 -- in (). This function type represents the type of the operations.
@@ -418,7 +436,7 @@
 instance HasGroundsR (PreOpenAcc op env) where
   groundsR (Exec _ _)        = TupRunit
   groundsR (Return vars)     = groundsR vars
-  groundsR (Manifest var)    = groundsR var     
+  groundsR (Manifest var)    = groundsR var
   groundsR (Compute e)       = groundsR e
   groundsR (Alet _ _ _ a)    = groundsR a
   groundsR (Alloc _ tp _)    = TupRsingle $ GroundRbuffer tp
@@ -428,6 +446,7 @@
   groundsR (Awhile _ _ _ a)  = groundsR a
   groundsR (Aassert _ _)     = TupRsingle $ GroundRscalar scalarTypeWord8
   groundsR (Aassume _)       = TupRsingle $ GroundRscalar scalarTypeWord8
+  groundsR (Atrace _ _)      = TupRsingle $ GroundRscalar scalarTypeWord8
   groundsR (Fence _ a)       = groundsR a
 
 instance HasGroundsR (GroundVar env) where
@@ -585,13 +604,19 @@
 reindexAcc r (Awhile tr poa poa' tr') = Awhile tr <$> reindexAfun r poa <*> reindexAfun r poa' <*> reindexVars r tr'
 reindexAcc r (Aassert msg cond) = Aassert msg <$> reindexExp r cond
 reindexAcc r (Aassume cond) = Aassume <$> reindexExp r cond
+reindexAcc r (Atrace msg t) = Atrace msg <$> reindexArrayDescriptors r t
 reindexAcc r (Fence set e) = Fence <$> reindexIdxSet r set <*> reindexAcc r e
 
+reindexArrayDescriptors :: (Applicative f) => ReindexPartial f env env' -> ArrayDescriptors env a -> f (ArrayDescriptors env' a)
+reindexArrayDescriptors k = traverseTupR (reindexArrayDescriptor k)
+
+reindexArrayDescriptor :: Applicative f => ReindexPartial f env env' -> ArrayDescriptor env a -> f (ArrayDescriptor env' a)
+reindexArrayDescriptor k (ArrayDescriptor shr sh buffers) = ArrayDescriptor shr <$> reindexVars k sh <*> reindexVars k buffers
+
 reindexAfun :: Applicative f => ReindexPartial f env env' -> PreOpenAfun op env t -> f (PreOpenAfun op env' t)
 reindexAfun r (Abody poa) = Abody <$> reindexAcc r poa
 reindexAfun r (Alam lhs poa) = reindexLHS r lhs $ \lhs' r' -> Alam lhs' <$> reindexAfun r' poa
 
-
 reindexLHS :: Applicative f => ReindexPartial f env env' -> LeftHandSide s t env env1 -> (forall env1'. LeftHandSide s t env' env1' -> ReindexPartial f env1 env1' -> r) -> r
 reindexLHS r (LeftHandSideSingle st) k = k (LeftHandSideSingle st) $ \case
   ZeroIdx -> pure ZeroIdx
@@ -670,6 +695,7 @@
   Awhile uniqueness c g a       -> Awhile uniqueness (mapAfunExecutable f c) (mapAfunExecutable f g) a
   Aassert msg cond              -> Aassert msg cond
   Aassume cond                  -> Aassume cond
+  Atrace msg t                  -> Atrace msg t
   Fence set e                   -> Fence set (mapAccExecutable f e)
 
 mapAfunExecutable :: (forall args benv'. op args -> Args benv' args -> op' args) -> PreOpenAfun op benv t -> PreOpenAfun op' benv t
@@ -703,8 +729,12 @@
   rnf (Awhile us cond step initial) = rnfTupR rnfUniqueness us `seq` rnf cond `seq` rnf step `seq` rnfGroundVars initial
   rnf (Aassert _ cond)              = rnfOpenExp cond
   rnf (Aassume cond)                = rnfOpenExp cond
+  rnf (Atrace _ t)                  = rnfTupR rnf t
   rnf (Fence set a)                 = IdxSet.rnfIdxSet set `seq` rnf a
 
+instance NFData (ArrayDescriptor env a) where
+  rnf (ArrayDescriptor shr sh buffers) = rnfShapeR shr `seq` rnfGroundVars sh `seq` rnfGroundVars buffers
+
 instance NFData' op => NFData (OperationAfun op env a) where
   rnf (Abody a) = rnf a
   rnf (Alam lhs f) = rnfLeftHandSide rnfGroundR lhs `seq` rnf f

src/Data/Array/Accelerate/AST/Schedule/Uniform.hs

@@ -309,6 +309,8 @@ data Effect kernel env where
 
   Aassert       :: Text -> Exp env PrimBool -> Effect kernel env
 
+  Atrace        :: Text -> ArrayDescriptors env t -> Effect kernel env
+
 -- A base value in the schedule is a scalar, buffer, a signal (resolver)
 -- or a (possibly mutable) reference
 --
@@ -399,6 +401,7 @@ effectFreeVars (SignalAwait signals)     = IdxSet.fromList $ map Exists $ signal
 effectFreeVars (SignalResolve resolvers) = IdxSet.fromList $ map Exists resolvers
 effectFreeVars (RefWrite ref value)      = IdxSet.insertVar ref $ IdxSet.singletonVar value
 effectFreeVars (Aassert _ cond)          = bindingFreeVars $ Compute cond
+effectFreeVars (Atrace _ t)              = arrayDescriptorsIdxSet t
 
 sargVar :: SArg env t -> Exists (Idx env)
 sargVar (SArgScalar   v) = Exists $ varIdx v

src/Data/Array/Accelerate/Analysis/Hash/Schedule/Uniform.hs

@@ -136,6 +136,10 @@ encodeEffect = \case
     intHost $(hashQ "Aassert")
     <> intHost (Hashable.hash msg)
     <> encodeOpenExp cond
+  Atrace msg t ->
+    intHost $(hashQ "Atrace")
+    <> intHost (Hashable.hash msg)
+    <> encodeArrayDescriptors t
 
 encodeIO :: InputOutputR input output -> Builder
 encodeIO = \case

src/Data/Array/Accelerate/Pretty/Operation.hs

@@ -103,6 +103,8 @@ prettyOpenAcc env = \case
         <> hardline <> hang 4 ("  ( " <> prettyOpenAfun env step)
         <> hardline <> "  )"
         <> hardline <> indent 2 (prettyVars (val env) 10 initial)
+  Atrace msg t ->
+    hang 2 (group $ vsep [annotate Statement "atrace", prettyText msg, prettyArrayDescriptor (val env) t])
   Aassert msg g ->
     hang 2 (group $ vsep [annotate Statement "assert", prettyText msg, prettyExp (val env) g])
   Aassume g ->
@@ -112,6 +114,9 @@ prettyOpenAcc env = \case
     <> hardline
     <> prettyOpenAcc env next
 
+prettyArrayDescriptor :: Val env -> ArrayDescriptors env t -> Adoc
+prettyArrayDescriptor env = prettyTupR (\p (ArrayDescriptor _ sh buffer) -> prettyShapeVars env sh <+> " " <+> prettyVars env p buffer) 0
+
 prettyArgs :: Val benv -> Args benv f -> Adoc
 prettyArgs env args = tupled $ map (\(Exists a) -> prettyArg env a) $ argsToList args
 

src/Data/Array/Accelerate/Pretty/Schedule/Uniform.hs

@@ -191,6 +191,7 @@ prettyEffect env = \case
   SignalResolve signals -> hang 2 $ group $ vsep [annotate Statement "resolve", list $ map (prettyIdx env) signals]
   RefWrite ref value    -> hang 2 $ group $ vsep ["*" <> prettyVar env ref <+> "=", prettyVar env value]
   Aassert msg cond      -> hang 2 $ group $ vsep [annotate Statement "assert", prettyText msg, prettyExp (val env) cond]
+  Atrace msg t          -> hang 2 $ group $ vsep [annotate Statement "atrace", prettyText msg, prettyArrayDescriptors env t]
 
 prettyKernelFun :: forall kernel env f. PrettyKernel kernel => Val' env -> KernelFun kernel f -> SArgs env f -> Adoc
 prettyKernelFun env fun args = case prettyKernel of
@@ -233,3 +234,8 @@ prettyShapeVars :: Val' env -> Vars s env sh -> Adoc
 prettyShapeVars _   TupRunit = "Z"
 prettyShapeVars env vars = encloseSep "Z :. " "" " :. " $ map (\(Exists v) -> prettyVar env v) $ flattenTupR vars
 
+prettyArrayDescriptors :: Val' env -> ArrayDescriptors env t -> Adoc
+prettyArrayDescriptors env = foldMapTupR (prettyArrayDescriptor env)
+
+prettyArrayDescriptor :: Val' env -> ArrayDescriptor env t -> Adoc
+prettyArrayDescriptor env (ArrayDescriptor _ sh buffer) = "Array: (" <+> prettyShapeVars env sh <+> prettyVars env 0 buffer <+> ")"

src/Data/Array/Accelerate/Trafo/Operation/Bounds.hs

@@ -186,6 +186,8 @@ boundsOptimizeAcc env@(BoundsEnv _ _ zero bindings) acc = case acc of
     , env'' <- env{ boundsGraph = boundsGraphClearNodes (boundsGraph env) $ accIdxSet modified }
     -> (IdxSet.empty, env'', bottomsGround zero $ groundsR acc, Fence set $ body')
 
+  Atrace msg t -> (IdxSet.empty, env, TupRsingle $ bottom zero scalarTypeWord8, Atrace msg t)
+
   Aassert msg expr
     | (_, expr') <- boundsOptimizeExp env expr
     , expr'' <- simplifyExp expr'

src/Data/Array/Accelerate/Trafo/Operation/LiveVars.hs

@@ -25,7 +25,8 @@
   stronglyLiveVariables, stronglyLiveVariablesFun,
 
   SLVOperation(..), ShrinkOperation(..), ShrunkOperation(..), SubArgs(..), SubArg(..),
-  reEnvArrayInstr,
+  reEnvArrayInstr, 
+  reEnvArrayDescriptors,
   ShrinkArg(..), shrinkArgs, composeSubArgs,
 
   defaultSlvGenerate, defaultSlvMap, defaultSlvBackpermute,
@@ -211,6 +212,15 @@
         liveness3
         $ \re _ ->
           Left $ Awhile us' (condition' re) (step' re) $ expectJust $ reEnvVars re initial
+  Atrace msg t
+    | free <- arrayDescriptorsIdxSet t
+    , liveness1 <- setIdxSetLive free liveness ->
+      LVAnalysis
+        liveness1
+        $ \re s -> case s of
+          SubTupRkeep -> Right $ Atrace msg $ reEnvArrayDescriptors re t
+          SubTupRskip -> Right $ Return TupRunit
+
   Fence set next
     | liveness1 <- setIdxSetLive set liveness
     , LVAnalysis liveness2 next' <- stronglyLiveVariables' liveness1 returns us next ->
@@ -307,7 +317,7 @@
 defaultSlvGenerate
   :: (forall sh' t'. op (Fun' (sh' -> t') -> Out sh' t' -> ()))
   -> Maybe (ShrinkOperation op (Fun' (sh -> t) -> Out sh t -> ()))
 defaultSlvGenerate mkGenerate = Just $ ShrinkOperation $ \subArgs args@(ArgFun f :>: array :>: ArgsNil) _ -> case subArgs of
   SubArgKeep `SubArgsLive` SubArgKeep `SubArgsLive` SubArgsNil
     -> ShrunkOperation mkGenerate args
   SubArgKeep `SubArgsLive` SubArgOut subTp `SubArgsLive` SubArgsNil
@@ -317,7 +327,7 @@
 defaultSlvMap
   :: (forall sh' s' t'. op (Fun' (s' -> t') -> In sh' s' -> Out sh' t' -> ()))
   -> Maybe (ShrinkOperation op (Fun' (s -> t)    -> In sh s -> Out sh  t -> ()))
 defaultSlvMap mkMap = Just $ ShrinkOperation $ \subArgs args@(ArgFun f :>: input :>: output :>: ArgsNil) _ -> case subArgs of
   SubArgKeep `SubArgsLive` SubArgKeep `SubArgsLive` SubArgKeep `SubArgsLive` SubArgsNil
     -> ShrunkOperation mkMap args
   SubArgKeep `SubArgsLive` SubArgKeep `SubArgsLive` SubArgOut subTp `SubArgsLive` SubArgsNil
@@ -327,7 +337,7 @@
 defaultSlvBackpermute
   :: (forall sh1' sh2' t'. op (Fun' (sh2' -> sh1') -> In sh1' t' -> Out sh2' t' -> ()))
   -> Maybe (ShrinkOperation op (Fun' (sh2 -> sh1) -> In sh1 t -> Out sh2 t -> ()))
 defaultSlvBackpermute mkBackpermute = Just $ ShrinkOperation $ \subArgs args@(f :>: ArgArray In (ArrayR shr r) sh buf :>: output :>: ArgsNil) _ -> case subArgs of
     SubArgKeep `SubArgsLive` SubArgKeep `SubArgsLive` SubArgKeep `SubArgsLive` SubArgsNil
       -> ShrunkOperation mkBackpermute args
     SubArgKeep `SubArgsLive` SubArgKeep `SubArgsLive` SubArgOut s `SubArgsLive` SubArgsNil
@@ -348,6 +358,9 @@
 reEnvArg re (ArgFun f)    = ArgFun $ mapArrayInstrFun (reEnvArrayInstr re) f
 reEnvArg re (ArgArray m repr sh buffers) = ArgArray m repr (expectJust $ reEnvVars re sh) (expectJust $ reEnvVars re buffers)
 
+reEnvArrayDescriptors :: ReEnv env subenv -> ArrayDescriptors env t -> TupR (ArrayDescriptor subenv) t
+reEnvArrayDescriptors re = expectJust . traverseTupR (\(ArrayDescriptor shape sh buffer) -> ArrayDescriptor shape <$> reEnvVars re sh <*> reEnvVars re buffer)
+
 -- Captures existential f'
 data ReEnvSubArgs subenv f where
   ReEnvSubArgs :: SubArgs f f'

src/Data/Array/Accelerate/Trafo/Operation/Simplify.hs

@@ -343,6 +343,16 @@ simplify' uniquenesses = \case
               $ awhileSimplifyInvariant us (cond' env') (step' env') $ simplifyReturnVars env us initial
       )
 
+  Atrace msg t ->
+    let
+      set = arrayDescriptorsIdxSet t
+    in
+      ( set
+      , \env ->
+        fence (syncSubstitutes env set)
+        $ Atrace msg $ simplifyArrayDescriptor env t
+    )
+
   Aassert msg cond ->
     ( IdxSet.empty
     , \env ->
@@ -379,6 +389,9 @@ simplify' uniquenesses = \case
               (next' env')
       )
 
+simplifyArrayDescriptor :: InfoEnv env -> ArrayDescriptors env t -> ArrayDescriptors env t
+simplifyArrayDescriptor env = mapTupR (\(ArrayDescriptor shape sh buffers) -> ArrayDescriptor shape (mapTupR (weaken $ substitute env) sh) (mapTupR (weaken $ substitute env) buffers))
+
 -- Given an environment, the set of updated variables and a list of copies of
 -- an operation, checks whether the operation copies all its outputs from
 -- undefined buffers.

src/Data/Array/Accelerate/Trafo/Operation/Substitution.hs

@@ -103,6 +103,9 @@ reindexArrayInstr' k (Parameter v) = Parameter <$> reindexVar' k v
 reindexExp' :: (Applicative f, RebuildableExp e) => SunkReindexPartial f benv benv' -> e (ArrayInstr benv) env t -> f (e (ArrayInstr benv') env t)
 reindexExp' k = rebuildArrayInstrPartial (rebuildArrayInstrMap $ reindexArrayInstr' k)
 
+reindexArrayDescriptor' :: (Applicative f) => SunkReindexPartial f env env' -> ArrayDescriptors env t -> f (TupR (ArrayDescriptor env') t)
+reindexArrayDescriptor' k = traverseTupR (\(ArrayDescriptor shape sh buffer) -> ArrayDescriptor shape <$> reindexVars' k sh <*> reindexVars' k buffer)
+
 reindexIdxSet'
   :: forall f env env' . Applicative f
   => SunkReindexPartial f env env'
@@ -127,6 +130,7 @@ reindexA' k = \case
     Unit var -> Unit <$> reindexVar' k var
     Acond c t f -> Acond <$> reindexVar' k c <*> travA t <*> travA f
     Awhile uniqueness c f i -> Awhile uniqueness <$> reindexAfun' k c <*> reindexAfun' k f <*> reindexVars' k i
+    Atrace msg t -> Atrace msg <$> reindexArrayDescriptor' k t
     Aassert msg g -> Aassert msg <$> reindexExp' k g
     Aassume g -> Aassume <$> reindexExp' k g
     Fence set next -> Fence <$> reindexIdxSet' k set <*> travA next
@@ -183,6 +187,7 @@ makeManifest acc = case acc of
   -- condition is evaluated at compile time.
   Acond c t f -> Acond c (makeManifest t) (makeManifest f)
   Awhile{} -> acc -- Can't fuse anyway
+  Atrace{} -> acc -- Same as compute
   Return vars -> go vars
   Fence set next -> Fence set $ makeManifest next
   where

src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Clustering.hs

@@ -216,6 +216,7 @@ openReconstruct' singletons labelenv graph clusterslist mlab subclustersmap symb
         SBlk {} -> error "wrong type: block"
         SRet env' vars     -> Exists $ Return          (fromJust $ reindexVars (mkReindexPartial' env' env) vars)
         SCmp env' expr     -> Exists $ Compute         (fromJust $ reindexExp  (mkReindexPartial' env' env) expr)
+        SAtr msg env' t    -> Exists $ Atrace msg      (fromJust $ reindexArrayDescriptors (mkReindexPartial' env' env) t)
         SAsr msg env' expr -> Exists $ Aassert msg     (fromJust $ reindexExp  (mkReindexPartial' env' env) expr)
         SAsu env' expr     -> Exists $ Aassume         (fromJust $ reindexExp  (mkReindexPartial' env' env) expr)
         SAlc env' shr e sh -> Exists $ Alloc shr e     (fromJust $ reindexVars (mkReindexPartial' env' env) sh)

src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Graph.hs

@@ -571,6 +571,7 @@ data Symbol (op :: Type -> Type) where
   SCmp  :: Env env -> Exp env a                                                        -> Symbol op
   SAlc  :: Env env -> ShapeR sh -> ScalarType e -> ExpVars env sh                      -> Symbol op
   SUnt  :: Env env -> ExpVar env e                                                     -> Symbol op
+  SAtr  :: Text -> Env env -> ArrayDescriptors env t                                   -> Symbol op
   SAsr  :: Text -> Env env -> Exp env PrimBool                                         -> Symbol op
   SAsu  :: Env env -> Exp env PrimBool                                                 -> Symbol op
   SFen  :: Env env -> IdxSet env                                                       -> Symbol op
@@ -590,6 +591,7 @@ instance Show (Symbol op) where
   show (SCmp {}) = "Cmp"
   show (SAlc {}) = "Alc"
   show (SUnt {}) = "Unt"
+  show (SAtr {}) = "Atr"
   show (SAsr {}) = "Asr"
   show (SAsu {}) = "Asu"
   show (SFen {}) = "Fen"
@@ -1058,6 +1060,13 @@ mkFusionGraph (Awhile u cond body init) = do
     symbol whileN ?= SWhl env condN bodyN init u
   return res                                      -- to return a fresh value of the same type as the initial value.
 
+mkFusionGraph (Atrace msg t) = do
+  c    <- freshComp
+  env  <- use environment
+  c `requiresBuffers` getIdxSetDeps (arrayDescriptorsIdxSet t) env
+  symbol c ?= SAtr msg env t
+  TupRsingle <$> freshVal c (GroundRscalar scalarTypeWord8)
+
 mkFusionGraph (Aassert msg cond) = do
   c    <- freshComp
   env  <- use environment