Loop fusion in Haskell

Roman Leshchinskiy
Programming Languages and Systems University of New South Wales

What is this about?
What I do Data Parallel Haskell compiles nested data-parallel programs to flat data-parallel ones lots of arrays and collective operations involved

What is this about?
What I do Data Parallel Haskell compiles nested data-parallel programs to flat data-parallel ones lots of arrays and collective operations involved

zipWith (-) (zipWith (*) (zipWith (-) (zipWith (-) (zipWith (*) (zipWith (-) (zipWith (-)

(replicate_s segd as1) xs) (replicate_s segd bs1) ys)) (replicate_s segd bs2) ys) (replicate_s segd as2) xs))

What is this about?
What I do Data Parallel Haskell compiles nested data-parallel programs to flat data-parallel ones lots of arrays and collective operations involved What other people do array programs with lots of collective operations

return . foldl’ hash 5381 . map toLower . filter isAlpha =<< readFile f

What is this about?
What I do Data Parallel Haskell compiles nested data-parallel programs to flat data-parallel ones lots of arrays and collective operations involved What other people do array programs with lots of collective operations What everybody wants no temporary arrays fused loops C-like speed

Loop fusion is easy!
foo xs = map (*5) (map (+3) xs)

Loop fusion is easy!
foo xs = map (*5) (map (+3) xs)

RULES
"map/map" map f (map g xs) = map (f . g) xs

Loop fusion is easy!
foo xs = map (*5) (map (+3) xs) bar ys = filter even (filter (<42) ys)

RULES
"map/map" map f (map g xs) = map (f . g) xs

Loop fusion is easy!
foo xs = map (*5) (map (+3) xs) bar ys = filter even (filter (<42) ys)

RULES
"map/map" "filter/filter" map f (map g xs) = map (f . g) xs filter f (filter g xs) = filter (λ x → f x && g x) xs

Loop fusion is easy!
foo xs = map (*5) (map (+3) xs) bar ys = filter even (filter (<42) ys) baz zs = map (+1) (filter even zs)

RULES
"map/map" "filter/filter" map f (map g xs) = map (f . g) xs filter f (filter g xs) = filter (λ x → f x && g x) xs

Loop fusion is easy!
foo xs = map (*5) (map (+3) xs) bar ys = filter even (filter (<42) ys) baz zs = map (+1) (filter even zs)

RULES
"map/map" "filter/filter" "map/filter" map f (map g xs) = map (f . g) xs filter f (filter g xs) = filter (λ x → f x && g x) xs map f (filter g xs) = mapFilter f g xs

Loop fusion is easy!
foo xs = map (*5) (map (+3) xs) bar ys = filter even (filter (<42) ys) baz zs = map (+1) (filter even zs)

RULES
"map/map" "filter/filter" map f (map g xs) = map (f . g) xs filter f (filter g xs) = filter (λ x → f x && g x) xs "map/filter" map f (filter g xs) = mapFilter f g xs "map/mapFilter" map f (mapFilter g h xs) = mapFilter (f . g) h xs "mapFilter/filter" mapFilter f g (filter h xs) = mapFilter (f λ x → g x && h x) xs ...

Loop fusion is easy!
foo xs = map (*5) (map (+3) xs) bar ys = filter even (filter (<42) ys) baz zs = map (+1) (filter even zs)

RULES
"map/map" "filter/filter" map f (map g xs) = map (f . g) xs filter f (filter g xs) = filter (λ x → f x && g x) xs "map/filter" map f (filter g xs) = mapFilter f g xs "map/mapFilter" map f (mapFilter g h xs) = mapFilter (f . g) h xs "mapFilter/filter" mapFilter f g (filter h xs) = mapFilter (f λ x → g x && h x) xs ...

AD B

EA ID

The challenge

use a constant number of rewrite rules don’t require new rules for new combinators make adding new combinators easy fuse everything! don’t require specialised compiler support handle both sequential and parallel loops

Sequential loops

Streams
data Step s a = Yield a s | Done data Stream a = ∃s. Stream (s → Step s a) s

Streams
stepper data Step s a = Yield a s | Done data Stream a = ∃s. Stream (s → Step s a) s

state

stepper produces next element and state from current state similar to an iterator actually encodes an anamorphism (unfold)

Streams
data Step s a = Yield a s | Done data Stream a = ∃s. Stream (s → Step s a) s stepper produces next element and state from current state similar to an iterator actually encodes an anamorphism (unfold)

sumS :: Num a ⇒ Stream a → a sumS (Stream step s) = go 0 s where go z s = case step s of Yield x s’ → go (z+x) s’ Done → z

Streams
data Step s a = Yield a s | Done data Stream a = ∃s. Stream (s → Step s a) s stepper produces next element and state from current state similar to an iterator actually encodes an anamorphism (unfold)

stream :: Array a → Stream a stream arr = Stream step 0 where step i | i < length arr = Yield (arr ! i) (i+1) | otherwise = Done

Streams
data Step s a = Yield a s | Done data Stream a = ∃s. Stream (s → Step s a) s stepper produces next element and state from current state similar to an iterator actually encodes an anamorphism (unfold)

mapS :: (a → b) → Stream a → Stream b mapS f (Stream step s) = Stream step’ s where step’ s = case step s of Yield x s’ → Yield (f x) s’ Done → Done

Streams
data Step s a = Yield a s | Done data Stream a = ∃s. Stream (s → Step s a) s stepper produces next element and state from current state similar to an iterator actually encodes an anamorphism (unfold)

unstream :: Stream a → Array a unstream (Stream step s) = <allocate, fill and freeze>

Stream fusion in three easy steps

Stream fusion in three easy steps
Step 1: implement array operations in terms of streams sum :: Num a ⇒ Array a → a sum = sumS . stream map :: (a → b) → Array a → Array b map f = unstream . mapS f . stream

Stream fusion in three easy steps
Step 1: implement array operations in terms of streams sum :: Num a ⇒ Array a → a sum = sumS . stream map :: (a → b) → Array a → Array b map f = unstream . mapS f . stream Step 2: inline them sumsq :: Num a ⇒ Array a → a sumsq = sum . map (λx -> x*x)

Stream fusion in three easy steps
Step 1: implement array operations in terms of streams sum :: Num a ⇒ Array a → a sum = sumS . stream map :: (a → b) → Array a → Array b map f = unstream . mapS f . stream Step 2: inline them sumsq :: Num a ⇒ Array a → a sumsq = sum . map (λx -> x*x) = sumS . stream . unstream . mapS f . stream

Stream fusion in three easy steps
Step 1: implement array operations in terms of streams sum :: Num a ⇒ Array a → a sum = sumS . stream map :: (a → b) → Array a → Array b map f = unstream . mapS f . stream Step 2: inline them sumsq :: Num a ⇒ Array a → a sumsq = sum . map (λx -> x*x) = sumS . stream . unstream . mapS f . stream

Stream fusion in three easy steps
Step 1: implement array operations in terms of streams sum :: Num a ⇒ Array a → a sum = sumS . stream map :: (a → b) → Array a → Array b map f = unstream . mapS f . stream Step 2: inline them sumsq :: Num a ⇒ Array a → a sumsq = sum . map (λx -> x*x) = sumS . stream . unstream . mapS f . stream Step 3: eliminate stream/unstream pairs "stream/unstream" stream (unstream s) = s

Stream fusion in three easy steps
Step 1: implement array operations in terms of streams sum :: Num a ⇒ Array a → a sum = sumS . stream map :: (a → b) → Array a → Array b map f = unstream . mapS f . stream Step 2: inline them sumsq :: Num a ⇒ Array a → a sumsq = sum . map (λx -> x*x) = sumS . mapS f . stream Step 3: eliminate stream/unstream pairs "stream/unstream" stream (unstream s) = s

Stream fusion in three easy steps
Step 1: implement array operations in terms of streams sum :: Num a ⇒ Array a → a sum = sumS . stream

map :: (a → b) → Array a → Array b map f = unstream . mapS f . stream Step 2: inline them

sumsq :: Num a ⇒ Array a → a sumsq = sum . map (λx -> x*x) = sumS . mapS f . stream Step 3: eliminate stream/unstream pairs "stream/unstream" stream (unstream s) = s

Le tG

H

C

do

th

e

re st

Optimising stream operations
sumsq xs = sumS (mapS square ( stream xs))

Optimising stream operations
sumsq xs = sumS (mapS square ( stream xs)) inline

stream :: Array a → Stream a stream arr = Stream step 0 where step i | i < length arr = Yield (arr ! i) (i+1) | otherwise = Done

Optimising stream operations
sumsq xs = sumS ( mapS square (Stream step1 0)) where step1 i = case i < length xs of True → Yield (xs ! i) (i+1) False → Done

Optimising stream operations
sumsq xs = sumS ( mapS square (Stream step1 0)) where step1 i = case i < inline length xs of True → Yield (xs ! i) (i+1) False → Done

mapS :: (a → b) → Stream a → b mapS f (Stream step s) = Stream step’ s where step’ s = case step s of Yield x s’ → Yield (f x) s’ Done → Done

Optimising stream operations
sumsq xs = sumS (Stream step2 0) where step1 i = case i < length xs of True → Yield (xs ! i) (i+1) False → Done step2 i = case step1 i of Yield x i’ → Yield (square x) i’ Done → Done

Optimising stream operations
sumsq xs = sumS (Stream step2 0) where step1 i = case i < length xs of inline True → Yield (xs ! i) (i+1) False → Done step2 i = case step1 i of Yield x i’ → Yield (square x) i’ Done → Done sumS :: Num a ⇒ Stream a → a sumS (Stream step s) = go 0 s where go z s = case step s of Yield x s’ → go (z+x) s’ Done → z

Optimising stream operations
sumsq xs = go 0 0 where step1 i = case i < length xs of True → Yield (xs ! i) (i+1) False → Done step2 i = case step1 i of Yield x i’ → Yield (square x) i’ Done → Done go z i = case step2 i of Yield x i’ → go (z+x) i’ Done → z

Optimising stream operations
sumsq xs = go 0 0 where step1 i = case i < length xs of True → Yield (xs ! i) (i+1) False → Done step2 i = case step1 i of Yield x i’ → Yield (square x) i’ Done → Done go z i = case step2 i of Yield x i’ → go (z+x) i’ Done → inline z

Optimising stream operations
sumsq xs = go 0 0 where step1 i = case i < length xs of True → Yield (xs ! i) (i+1) False → Done go z i = case (case step1 i Yield x i’ Done Yield x i’ → go Done → z of → Yield (square x) i’ → Done) of (z+x) i’

Optimising stream operations
sumsq xs = go 0 0 where step1 i = case i < length xs of True → Yield (xs ! i) (i+1) case of case False → Done go z i = case (case step1 i of Yield x i’ → Yield (square x) i’ Done → Done) of Yield x i’ → go (z+x) i’ Done → z

Optimising stream operations
sumsq xs = go 0 0 where step1 i = case i < length xs of True → Yield (xs ! i) (i+1) False → Done go z i = case step1 i of Yield x i’ → go (z + square x) i’ Done → z

Optimising stream operations
sumsq xs = go 0 0 where step1 i = case i < length xs of True → Yield (xs ! i) (i+1) False → Done go z i = case step1 i of Yield x i’ → go (z + square x) i’ Done → inline z

Optimising stream operations
sumsq xs = go 0 0 where go z i = case (case i < length xs of True → Yield (xs ! i) (i+1) False → Done) of Yield x i’ → go (z + square x) i’ Done → z

Optimising stream operations
case of case sumsq xs = go 0 0 where go z i = case (case i < length xs of True → Yield (xs ! i) (i+1) False → Done) of Yield x i’ → go (z + square x) i’ Done → z

Optimising stream operations
sumsq xs = go 0 0 where go z i = case i < length xs of True → go (z + square (xs ! i)) (i+1) False → z

Optimising stream operations
sumsq xs = go 0 0 where go z i = case i < length xs of True → go (z + square (xs ! i)) (i+1) False → z

optimal loop no Stream or Step values ever created only general-purpose optimisations will be optimised further (unboxing etc.) requires a great compiler (thanks GHC team!)

Why does it work?
sumsq xs = go 0 0 where step1 i = case i < length xs of True → Yield (xs ! i) (i+1) False → Done step2 i = case step1 i of Yield x i’ → Yield (square x) i’ Done → Done go z i = case step2 i of Yield x i’ → go (z+x) i’ Done → z

Why does it work?
sumsq xs = go 0 0 where step1 i = case i < length xs of ive True -recuYield (xs ! i) (i+1) → rs non False → Done step2 i = case step1 i of Yield x i’ → Yield (square x) i’ Done → Done go z i = case step2 i of Yield x i’ → go (z+x) i’ Done → z

Why does it work?
sumsq xs = go 0 0 where step1 i = case i < length xs of ive True -recuYield (xs ! i) (i+1) → rs non False → Done step2 i = case step1 i of ive Yield -recurs → Yield (square x) i’ x i’ non Done → Done go z i = case step2 i of Yield x i’ → go (z+x) i’ Done → z

Why does it work?
sumsq xs = go 0 0 where step1 i = case i < length xs of ive True -recuYield (xs ! i) (i+1) → rs non False → Done step2 i = case step1 i of ive Yield -recurs → Yield (square x) i’ x i’ non Done → Done go z i = case step2 i of Yield xursive→ go (z+x) i’ i’ rec Done → z

A slight problem

filterS :: (a → Bool) → Stream a → Stream a filterS f (Stream step s) = Stream step’ s where step’ s = case step s of Yield x s’ | f x → Yield x s’ | otherwise → step s’ Done → Done

A slight problem

filterS :: (a → Bool) → Stream a → Stream a filterS f (Stream step s) = Stream step’ s where step’ s = case step s of Yield x s’ | f x cursive → Yield x s’ re | otherwise → step s’ Done → Done

Extending streams
Idea: allow a loop iteration not to produce an element

Extending streams
Idea: allow a loop iteration not to produce an element data Step s a = Yield a s | Skip s | Done

Extending streams
Idea: allow a loop iteration not to produce an element data Step s a = Yield a s | Skip s | Done filterS :: (a → Bool) → Stream a → Stream a filterS f (Stream step s) = Stream step’ s where step’ s = case step s of Yield x s’ | f x → Yield x s’ | otherwise → Skip s’ Skip s’ → Skip s’ Done → Done

Extending streams
Idea: allow a loop iteration not to produce an element data Step s a = Yield a s | Skip s | Done filterS :: (a → Bool) → Stream a → Stream a filterS f (Stream step s) = Stream step’ s where step’ s = case step s of Yield x s’ | f x → Yield x s’ ve ursi Skip s’ | otherwiseec → n-r Skip s’ no → Skip s’ Done → Done

Stream fusion - summary

encode loops by streams implement array operations in terms of streams eliminate stream/unstream pairs (temporaries) stream producers are non-recursive standard optimisations remove overhead (loop fusion)

Stream fusion - summary

encode loops by streams implement array operations in terms of streams eliminate stream/unstream pairs (temporaries) stream producers are non-recursive standard optimisations remove overhead (loop fusion)

Standard optimisations: inlining, case-of-case, worker/wrapper transformation, SpecConstr, LiberateCase, specialisation ...

Parallel loops

DPH on multicores
Evaluation strategy after vectorisation operations are data parallel and flat executed by a gang of worker threads essentially fork-join parallelism

DPH on multicores
Evaluation strategy after vectorisation operations are data parallel and flat executed by a gang of worker threads essentially fork-join parallelism

mapP :: (a → b) → Array a → Array b mapP f xs = <split xs across workers> <map f over each chunk> <collect local results>

DPH on multicores
Evaluation strategy after vectorisation operations are data parallel and flat executed by a gang of worker threads essentially fork-join parallelism f is sequential mapP :: (a → b) → Array a → Array b mapP f xs = <split xs across workers> <map f over each chunk> <collect local results>

DPH on multicores
Evaluation strategy after vectorisation operations are data parallel and flat executed by a gang of worker threads essentially fork-join parallelism

mapP :: (a → b) → Array a → Array b mapP f xs = <split xs across workers> <map f over each chunk> <collect local results> sumP :: Num a ⇒ Array a → a sumP xs = <split xs across workers> <sum each chunk> <reduce local sums>

DPH on multicores
Evaluation strategy after vectorisation operations are data parallel and flat executed by a gang of worker threads essentially fork-join parallelism

sumsqP = sumP . mapP square

DPH on multicores
Evaluation strategy after vectorisation operations are data parallel and flat executed by a gang of worker threads essentially fork-join parallelism

sumsqP xs = <split xs across workers> <map square over each chunk> <collect local results> <split results across workers> <sum each chunk> <reduce local sums>

DPH on multicores
Evaluation strategy after vectorisation operations are data parallel and flat executed by a gang of worker threads essentially fork-join parallelism

sumsqP xs = <split xs across workers> <map square over each chunk> <collect local results> <split results across workers> <sum each chunk> <reduce local sums>

DPH on multicores
Evaluation strategy after vectorisation operations are data parallel and flat executed by a gang of worker threads essentially fork-join parallelism

sumsqP xs = <split xs across workers> <map square over each chunk> <collect local results> <split results across workers> <sum each chunk> <reduce local sums>

Distributed types
Idea: let’s make the evaluation strategy explicit! (Keller 1999)

Distributed types
Idea: let’s make the evaluation strategy explicit! (Keller 1999) data Dist a Dist (Array a) Dist Double a is distributed across threads each thread has a local array (chunk) each thread has a local Double

Distributed types
Idea: let’s make the evaluation strategy explicit! (Keller 1999) data Dist a Dist (Array a) Dist Double splitD joinD a is distributed across threads each thread has a local array (chunk) each thread has a local Double distribute an array across threads collect thread-local chunks

splitD :: Array a → Dist (Array a) joinD :: Dist (Array a) → Array a

Distributed types
Idea: let’s make the evaluation strategy explicit! (Keller 1999) data Dist a Dist (Array a) Dist Double splitD joinD mapD sumD splitD joinD mapD sumD :: :: :: :: a is distributed across threads each thread has a local array (chunk) each thread has a local Double distribute an array across threads collect thread-local chunks execute a sequential operation in each thread compute sum of local values Array a → Dist (Array a) Dist (Array a) → Array a (a → b) → Dist a → Dist b Num a ⇒ Dist a → a

Programming with distributed types

mapP f xs = <split xs across workers> <map f over each chunk> <collect local results>

Programming with distributed types

mapP f = joinD . mapD (map f) . splitD

-- collect -- map f over chunks -- split

Programming with distributed types

mapP f = joinD . mapD (map f) . splitD

-- collect -- map f over chunks -- split

sumP xs = <split xs across workers> <sum each chunk> <reduce local sums>

Programming with distributed types

mapP f = joinD . mapD (map f) . splitD sumP = sumD . mapD sum . splitD

-- collect -- map f over chunks -- split -- reduce -- sum each chunk -- split

Fusing distributed types

sumsqP = sumP . mapP square

Fusing distributed types

sumsqP = . . . . .

sumD mapD sum splitD joinD mapD (map square) splitD

-------

reduce sum each chunk split collect map square over chunks split

Fusing distributed types

sumsqP = . . . . .

sumD mapD sum splitD joinD mapD (map square) splitD

-------

reduce sum each chunk split collect map square over chunks split

RULES
splitD (joinD xs) = xs

Fusing distributed types

sumsqP = . . .

sumD mapD sum mapD (map square) splitD

-----

reduce sum each chunk map square over chunks split

RULES
splitD (joinD xs) = xs

Fusing distributed types

sumsqP = . . .

sumD mapD sum mapD (map square) splitD

-----

reduce sum each chunk map square over chunks split

RULES
splitD (joinD xs) = xs mapD f (mapD g xs) = mapD (f . g) xs

Fusing distributed types

sumsqP = sumD -- reduce . mapD (sum . map square) -- work . splitD -- split

RULES
splitD (joinD xs) = xs mapD f (mapD g xs) = mapD (f . g) xs

Fusing distributed types

sumsqP = sumD -- reduce . mapD (sum . map square) -- work . splitD -- split stream fusion

RULES
splitD (joinD xs) = xs mapD f (mapD g xs) = mapD (f . g) xs

Distributed types on multicores

data Dist a splitD joinD mapD splitD/joinD mapD/mapD

a is distributed across threads distribute xs across threads collect thread-local chunks execute a sequential operation in each thread eliminate communication eliminate synchronisation

Distributed types on clusters

data Dist a splitD joinD mapD splitD/joinD mapD/mapD

a is distributed across nodes scatter gather execute operation on each node eliminate communication eliminate synchronisation

Distributed types on GPUs

data Dist a splitD joinD mapD splitD/joinD mapD/mapD

a is in GPU memory CPU −→ GPU transfer GPU −→ CPU transfer execute kernel on the GPU eliminate memory transfers (communication) fuse kernels (synchronisation)

Distribured types – summary

encode parallel loops as split/work/join eliminate unnecessary split/join pairs fuse sequential work (stream fusion) very general mechanism for fusing parallel computations applicable to a wide range of architectures again, no specialised compiler support

1

Obligatory benchmark

1

2

4

8

sumsq, Haskell dotp, C

sumsq, C smvm, Haskell

dotp, Haskell smvm, C

Runtime @ greyarea
10000

1000

100

10

1

1

2

4

8

16

32

64

sumsq, Haskell dotp, C

sumsq, C smvm, Haskell

dotp, Haskell smvm, C

Parting thoughts
it’s nice, it’s easy to use, it works high-level functional programs compiled to highly efficient code even parallel ones! rewrite rules + great optimiser = win DPH doesn’t require any special-purpose optimisations try this in an imperative language...

Parting thoughts
it’s nice, it’s easy to use, it works high-level functional programs compiled to highly efficient code even parallel ones! rewrite rules + great optimiser = win DPH doesn’t require any special-purpose optimisations ’ try this inoant imperative language... d n

Parting thoughts
it’s nice, it’s easy to use, it works high-level functional programs compiled to highly efficient code even parallel ones! rewrite rules + great optimiser = win DPH doesn’t require any special-purpose optimisations ’ try this inoant imperative language... d n

Stream fusion: dph, bytestring, vector, uvector Distributed types: dph

Sign up to vote on this title
UsefulNot useful