Skip to content

Commit

Permalink
slv on partitionedacc
Browse files Browse the repository at this point in the history
  • Loading branch information
@dpvanbalen
dpvanbalen committed Apr 15, 2024
1 parent 5b8e358 commit f82ed09
Show file tree
Hide file tree
Showing 7 changed files with 231 additions and 167 deletions.
34 changes: 17 additions & 17 deletions accelerate.cabal
View file
Original file line number Diff line number Diff line change
Expand Up @@ -144,15 +144,15 @@ extra-source-files:
cbits/xkcp/*.inc
-- TRACY
-- These are referenced directly using the FFI
cbits/tracy/public/*.cpp
cbits/tracy/public/tracy/*.h
cbits/tracy/public/tracy/*.hpp
cbits/tracy/public/common/*.h
cbits/tracy/public/common/*.hpp
cbits/tracy/public/common/*.cpp
cbits/tracy/public/client/*.h
cbits/tracy/public/client/*.hpp
cbits/tracy/public/client/*.cpp
-- cbits/tracy/public/*.cpp
-- cbits/tracy/public/tracy/*.h
-- cbits/tracy/public/tracy/*.hpp
-- cbits/tracy/public/common/*.h
-- cbits/tracy/public/common/*.hpp
-- cbits/tracy/public/common/*.cpp
-- cbits/tracy/public/client/*.h
-- cbits/tracy/public/client/*.hpp
-- cbits/tracy/public/client/*.cpp
-- These are used to build Tracy's client tools in Setup.hs
cbits/tracy/capture/build/unix/Makefile
cbits/tracy/capture/build/unix/*.mk
Expand All @@ -176,12 +176,12 @@ extra-source-files:
cbits/tracy/profiler/src/*.cpp
cbits/tracy/profiler/src/*.h
cbits/tracy/profiler/src/*.hpp
cbits/tracy/profiler/src/font/*.hpp
cbits/tracy/profiler/src/imgui/*.cpp
cbits/tracy/profiler/src/imgui/*.h
cbits/tracy/public/libbacktrace/*.cpp
cbits/tracy/public/libbacktrace/*.h
cbits/tracy/public/libbacktrace/*.hpp
-- cbits/tracy/profiler/src/font/*.hpp
-- cbits/tracy/profiler/src/imgui/*.cpp
-- cbits/tracy/profiler/src/imgui/*.h
-- cbits/tracy/public/libbacktrace/*.cpp
-- cbits/tracy/public/libbacktrace/*.h
-- cbits/tracy/public/libbacktrace/*.hpp
cbits/tracy/server/*.cpp
cbits/tracy/server/*.h
cbits/tracy/server/*.hpp
Expand All @@ -190,7 +190,7 @@ extra-source-files:
cbits/tracy/zstd/common/*.h
cbits/tracy/zstd/compress/*.c
cbits/tracy/zstd/compress/*.h
cbits/tracy/zstd/decompress/*.S
-- cbits/tracy/zstd/decompress/*.S
cbits/tracy/zstd/decompress/*.c
cbits/tracy/zstd/decompress/*.h
cbits/tracy/zstd/dictBuilder/*.c
Expand Down Expand Up @@ -451,6 +451,7 @@ library
Data.Array.Accelerate.Trafo.LiveVars
Data.Array.Accelerate.Trafo.NewNewFusion
Data.Array.Accelerate.Trafo.Operation.Substitution
Data.Array.Accelerate.Trafo.Operation.LiveVars
Data.Array.Accelerate.Trafo.Partitioning.ILP
Data.Array.Accelerate.Trafo.Partitioning.ILP.Clustering
Data.Array.Accelerate.Trafo.Partitioning.ILP.Graph
Expand Down Expand Up @@ -528,7 +529,6 @@ library
Data.Array.Accelerate.Trafo.Exp.Algebra
Data.Array.Accelerate.Trafo.Environment
Data.Array.Accelerate.Trafo.Operation.Simplify
Data.Array.Accelerate.Trafo.Operation.LiveVars
Data.Array.Accelerate.Trafo.Shrink
Data.Atomic

Expand Down
280 changes: 156 additions & 124 deletions src/Data/Array/Accelerate/AST/Partitioned.hs
View file

Large diffs are not rendered by default.

62 changes: 43 additions & 19 deletions src/Data/Array/Accelerate/Eval.hs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,12 +70,14 @@ class ( MakesILP op
=> StaticClusterAnalysis (op :: Type -> Type) where
data BackendClusterArg2 op env arg
onOp :: op args -> BackendArgs op env (OutArgsOf args) -> Args env args -> FEnv op env -> BackendArgs op env args
bcaid :: BackendClusterArg op arg -> BackendClusterArg op arg'
def :: Arg env arg -> FEnv op env -> BackendClusterArg op arg -> BackendClusterArg2 op env arg
valueToIn :: BackendClusterArg2 op env (Value sh e) -> BackendClusterArg2 op env (In sh e)
valueToOut :: BackendClusterArg2 op env (Value sh e) -> BackendClusterArg2 op env (Out sh e)
inToValue :: BackendClusterArg2 op env (In sh e) -> BackendClusterArg2 op env (Value sh e)
outToValue :: BackendClusterArg2 op env (Out sh e) -> BackendClusterArg2 op env (Value sh e)
outToSh :: BackendClusterArg2 op env (Out sh e) -> BackendClusterArg2 op env (Sh sh e)
outToVar :: BackendClusterArg2 op env (Out sh e) -> BackendClusterArg2 op env (Var' sh )
shToOut :: BackendClusterArg2 op env (Sh sh e) -> BackendClusterArg2 op env (Out sh e)
shToValue :: BackendClusterArg2 op env (Sh sh e) -> BackendClusterArg2 op env (Value sh e)
varToValue :: BackendClusterArg2 op env (Var' sh) -> BackendClusterArg2 op env (Value sh e)
Expand All @@ -91,18 +93,37 @@ class ( MakesILP op
unpairinfo :: BackendClusterArg2 op env (m sh (a,b)) -> (BackendClusterArg2 op env (m sh a), BackendClusterArg2 op env (m sh b))
unpairinfo x = (shrinkOrGrow x, shrinkOrGrow x)


foo :: StaticClusterAnalysis op => SubArgs big small -> Args env small -> BackendArgs op env (OutArgsOf small) -> FEnv op env -> BackendCluster op small -> BackendArgs op env (OutArgsOf big)
foo SubArgsNil ArgsNil ArgsNil _ _ = ArgsNil
foo (SubArgKeep `SubArgsLive` subargs) (a:>:as) bs env (_ :>: cs) = case a of
ArgArray Out _ _ _ -> case bs of (b:>:bs') -> b :>: foo subargs as bs' env cs
ArgArray In _ _ _ -> foo subargs as bs env cs
ArgArray Mut _ _ _ -> foo subargs as bs env cs
ArgVar _ -> foo subargs as bs env cs
ArgExp _ -> foo subargs as bs env cs
ArgFun _ -> foo subargs as bs env cs
foo (SubArgOut _ `SubArgsLive` subargs) (_ :>: as) (b:>:bs) env (_ :>: cs) = shrinkOrGrow b :>: foo subargs as bs env cs
foo (SubArgsDead subargs) (a :>: as) bs env (c :>: cs) =
shToOut (varToSh (def a env c)) :>: foo subargs as bs env cs

makeBackendArg :: forall op env args. StaticClusterAnalysis op => Args env args -> FEnv op env -> Cluster op args -> BackendCluster op args -> BackendArgs op env args
makeBackendArg args env c b = go args c (defaultOuts args b)
makeBackendArg args env c b = go args c (defaultOuts args b) b
where
go :: forall args. Args env args -> Cluster op args -> BackendArgs op env (OutArgsOf args) -> BackendArgs op env args
go args (Fused f l r) outputs = let
backR = go (right f args) r (rightB args f outputs)
backL = go (left f args) l (backleft f backR outputs)
go :: forall args. Args env args -> Cluster op args -> BackendArgs op env (OutArgsOf args) -> BackendCluster op args -> BackendArgs op env args
go args (Fused f l r) outputs bs = let
backR = go (right f args) r (rightB args f outputs) (right' bcaid bcaid f bs)
backL = go (left f args) l (backleft f backR outputs) (left' bcaid f bs)
in fuseBack f backL backR
go args (Op (SOp (SOAOp op soa) (SA sort unsort)) sa) outputs =
sort . soaExpand uncombineB soa $ onOp @op op (forgetIn (soaShrink combine soa $ unsort args) $ soaShrink combineB soa $ unsort $ inventIn args outputs) (soaShrink combine soa $ unsort args) env
go args (Op (SLV (SOp (SOAOp op soa) (SA sort unsort)) subargs) _l) outputs bs =
slv outToVar subargs
. sort
. soaExpand uncombineB soa
$ onOp @op
op
(forgetIn (soaShrink combine soa . unsort $ slv' varout subargs args)
. soaShrink combineB soa . unsort $ inventIn (slv' varout subargs args) (foo subargs args outputs env bs))
(soaShrink combine soa . unsort $ slv' varout subargs args)
env

combineB :: BackendClusterArg2 op env (f l) -> BackendClusterArg2 op env (f r) -> BackendClusterArg2 op env (f (l,r))
combineB = unsafeCoerce $ pairinfo @op
Expand Down Expand Up @@ -222,10 +243,12 @@ evalCluster c b args env ix = do

evalOps :: forall op args env. (EvalOp op) => Index op -> Cluster op args -> BackendArgEnv op env (InArgs args) -> Args env args -> FEnv op env -> EvalMonad op (EmbedEnv op env (OutArgs args))
evalOps ix c ba args env = case c of
Op (SOp (SOAOp op soas) (SA f g)) l -> outargs f (g args)
. soaOut splitFromArg' (soaShrink combine soas $ g args) soas
<$> evalOp ix l op env (soaIn pairInArg (g args) soas
$ inargs g ba)
Op (SLV (SOp (SOAOp op soas) (SA f g)) subargs) l
-> slvOut args subargs
. outargs f (g $ slv' varout subargs args)
. soaOut splitFromArg' (soaShrink combine soas $ g $ slv' varout subargs args) soas
<$> evalOp ix l op env (soaIn pairInArg (g $ slv' varout subargs args) soas
$ inargs g $ slvIn (flip bvartosh env) subargs ba)
Fused f l r -> do
lin <- leftIn f ba env
lout <- evalOps ix l lin (left f args) env
Expand Down Expand Up @@ -384,7 +407,7 @@ instance EvalOp op => TupRmonoid (Compose (BackendArgEnvElem op env) (Sh sh)) wh
-- use this to check whether a singleton cluster is a generate, map, etc
peekSingletonCluster :: (forall args'. op args' -> r) -> Cluster op args -> Maybe r
peekSingletonCluster k = \case
Op (SOp (SOAOp op _) _) _ -> Just $ k op
Op (SLV (SOp (SOAOp op _) _) _) _ -> Just $ k op
_ -> Nothing -- not a singleton cluster


Expand All @@ -396,27 +419,28 @@ applySingletonCluster :: forall op env args args' r
-> Args env args
-> r
applySingletonCluster k c args = case c of
Op (SOp (SOAOp op soas) (SA _ unsort)) _ ->
Op (SLV (SOp (SOAOp op soas) (SA _ unsort)) subargs) _ ->
unsafeCoerce @(op args' -> Args env args' -> r) @(op _ -> Args env _ -> r)
k
op
$ soaShrink combine soas $ unsort args
$ soaShrink combine soas $ unsort $ slv' varout subargs args
_ -> error "not singleton"


-- only use this function if you know it is a singleton cluster of the right operation
applySingletonCluster' :: forall op env args args' f
. (op args' -> Args env args' -> PreArgs f args')
-> (forall l r g. f (g (l,r)) -> (f (g l), f (g r)))
-> (forall sh e. f (Out sh e) -> f (Var' sh))
-> Cluster op args
-> Args env args
-> PreArgs f args
applySingletonCluster' k f c args = case c of
Op (SOp (SOAOp op soas) (SA sort unsort)) _ ->
sort $ soaExpand f soas $
applySingletonCluster' k f outvar' c args = case c of
Op (SLV (SOp (SOAOp op soas) (SA sort unsort)) subargs) _ ->
slv outvar' subargs $ sort $ soaExpand f soas $
unsafeCoerce @(op args' -> Args env args' -> PreArgs f args') @(op _ -> Args env _ -> PreArgs f _)
k
op
$ soaShrink combine soas $ unsort args
$ soaShrink combine soas $ unsort $ slv' varout subargs args
_ -> error "not singleton"

9 changes: 6 additions & 3 deletions src/Data/Array/Accelerate/Pretty/Partitioned.hs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,8 @@ import Prelude hiding (exp)
import Data.Array.Accelerate.Representation.Type (TupR (..))
import Data.Array.Accelerate.AST.Idx (Idx (..))
import Data.Bifunctor (second)
import Data.Array.Accelerate.Representation.Array (ArrayR(..))
import Data.Array.Accelerate.AST.Var (varsType)

instance PrettyOp op => PrettyOp (Clustered op) where
prettyOp :: PrettyOp op => Clustered op t -> Adoc
Expand All @@ -45,9 +47,10 @@ instance PrettyOp op => PrettyOp (Clustered op) where

instance PrettyOp op => PrettyOp (Cluster op) where
prettyOp (Fused _ l r) = "Fused (" <> prettyOp l <> ", " <> prettyOp r
prettyOp (Op (SOp (SOAOp op _) _) _) = prettyOp op
prettyOpWithArgs env (Fused f l r) args = "Fused (" <> prettyOpWithArgs env l (left f args) <> ", " <> prettyOpWithArgs env r (right f args)
prettyOpWithArgs env (Op (SOp (SOAOp op soa) (SA _ unsort)) _) args = prettyOpWithArgs env op (soaShrink combine soa . unsort $ args)
prettyOp (Op (SLV (SOp (SOAOp op _) _) _) _) = prettyOp op
prettyOpWithArgs env (Fused f l r) args = "Fused" -- (" <> prettyOpWithArgs env l (left f args) <> ", " <> prettyOpWithArgs env r (right f args)
prettyOpWithArgs env (Op (SLV (SOp (SOAOp op soa) (SA _ unsort)) subargs) _) args =
prettyOpWithArgs env op (soaShrink combine soa . unsort . slv' varout subargs $ args)

-- prettyOpWithArgs :: forall env t. Val env -> Cluster op t -> Args env t -> Adoc
-- prettyOpWithArgs env (Op (SLVOp (SOp (SOAOp op soa) (SA _ unsort)) sa) _) args = prettyOpWithArgs env op (soaShrink combine soa . unsort . slv' varToOut sa $ args)
Expand Down
1 change: 1 addition & 0 deletions src/Data/Array/Accelerate/Trafo/Operation/LiveVars.hs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -230,6 +230,7 @@ class SLVOperation op where

newtype ShrinkOperation op f = ShrinkOperation (forall f' env' env. SubArgs f f' -> Args env' f' -> Args env f -> ShrunkOperation op env')

-- existential over f: otherwise, you couldn't change the non-array arguments. You need this e.g. for a Generate: smaller array means smaller function.
data ShrunkOperation op env where
ShrunkOperation :: op f -> Args env f -> ShrunkOperation op env

Expand Down
8 changes: 5 additions & 3 deletions src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Clustering.hs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
-Wno-overlapping-patterns
-Wno-incomplete-patterns
#-}
{-# LANGUAGE BlockArguments #-}

module Data.Array.Accelerate.Trafo.Partitioning.ILP.Clustering where

Expand Down Expand Up @@ -314,10 +315,11 @@ data FoldType op env


unfused :: forall op args env r. MakesILP op => op args -> Label -> LabelledArgsOp op env args -> (forall args'. Clustered op args' -> LabelledArgsOp op env args' -> r) -> r
unfused op l largs k = singleton l largs op $
\c@(Clustered (Op (SOp (SOAOp (_op :: op argsToo) soas) (SA sort _unsort)) _) b) ->
unfused op l largs k = singleton l largs op \case
c@(Clustered (Op (SLV (SOp (SOAOp (_op :: op argsToo) soas) (SA sort _unsort)) subargs) _l) _b) ->
case unsafeCoerce Refl of -- we know that `_op` is the same as `op`
(Refl :: args :~: argsToo) -> k c (sort $ soaExpand splitLabelledArgsOp soas largs)
(Refl :: args :~: argsToo) -> k c (slv louttovar subargs $ sort $ soaExpand splitLabelledArgsOp soas largs)
_ -> error "singleton gave fused"

louttovar :: LabelledArgOp op env (Out sh e) -> LabelledArgOp op env (Var' sh)
louttovar (LOp a (_,ls) b) = LOp (outvar a) (NotArr, ls) b -- unsafe marker: maybe this NotArr ends up a problem?
Expand Down
4 changes: 3 additions & 1 deletion src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Graph.hs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@ import Data.Array.Accelerate.Trafo.Partitioning.ILP.Solver
import Data.Array.Accelerate.Type
import Data.Array.Accelerate.Analysis.Hash.Exp

import Data.Array.Accelerate.Trafo.Operation.LiveVars

-- Data structures
-- In this file, order often subly matters.
-- To keep this clear, we use Set whenever it does not,
Expand Down Expand Up @@ -160,7 +162,7 @@ unOpLabels = mapArgs $ \(LOp arg l _) -> L arg l

type BackendCluster op = PreArgs (BackendClusterArg op)

class (Eq (BackendVar op), Ord (BackendVar op), Eq (BackendArg op), Show (BackendArg op), Ord (BackendArg op), Show (BackendVar op)) => MakesILP op where
class (ShrinkArg (BackendClusterArg op), Eq (BackendVar op), Ord (BackendVar op), Eq (BackendArg op), Show (BackendArg op), Ord (BackendArg op), Show (BackendVar op)) => MakesILP op where
-- Vars needed to express backend-specific fusion rules.
type BackendVar op
-- Information that the backend attaches to the argument for reconstruction,
Expand Down

0 comments on commit f82ed09

Please sign in to comment.