accelerate-llvm-native-1.4.0.0: src/Data/Array/Accelerate/LLVM/Native/CodeGen/Scan.hs
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RebindableSyntax #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ViewPatterns #-}
-- |
-- Module : Data.Array.Accelerate.LLVM.Native.CodeGen.Scan
-- Copyright : [2014..2020] The Accelerate Team
-- License : BSD3
--
-- Maintainer : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability : experimental
-- Portability : non-portable (GHC extensions)
--
module Data.Array.Accelerate.LLVM.Native.CodeGen.Scan
where
import Data.Array.Accelerate.AST ( Direction(..) )
import Data.Array.Accelerate.Representation.Array
import Data.Array.Accelerate.Representation.Shape
import Data.Array.Accelerate.Representation.Type
import Data.Array.Accelerate.Type
import Data.Array.Accelerate.LLVM.CodeGen.Arithmetic as A
import Data.Array.Accelerate.LLVM.CodeGen.Array
import Data.Array.Accelerate.LLVM.CodeGen.Base
import Data.Array.Accelerate.LLVM.CodeGen.Environment
import Data.Array.Accelerate.LLVM.CodeGen.Exp
import Data.Array.Accelerate.LLVM.CodeGen.Loop
import Data.Array.Accelerate.LLVM.CodeGen.Monad
import Data.Array.Accelerate.LLVM.CodeGen.Sugar
import Data.Array.Accelerate.LLVM.Compile.Cache
import Data.Array.Accelerate.LLVM.Native.CodeGen.Base
import Data.Array.Accelerate.LLVM.Native.CodeGen.Generate
import Data.Array.Accelerate.LLVM.Native.CodeGen.Loop
import Data.Array.Accelerate.LLVM.Native.Target ( Native )
import Control.Applicative
import Control.Monad
import Data.String ( fromString )
import Data.Coerce as Safe
import Prelude as P
-- 'Data.List.scanl' or 'Data.List.scanl1' style exclusive scan,
-- but with the restriction that the combination function must be associative
-- to enable efficient parallel implementation.
--
-- > scanl (+) 10 (use $ fromList (Z :. 10) [0..])
-- >
-- > ==> Array (Z :. 11) [10,10,11,13,16,20,25,31,38,46,55]
--
mkScan
:: UID
-> Gamma aenv
-> ArrayR (Array (sh, Int) e)
-> Direction
-> IRFun2 Native aenv (e -> e -> e)
-> Maybe (IRExp Native aenv e)
-> MIRDelayed Native aenv (Array (sh, Int) e)
-> CodeGen Native (IROpenAcc Native aenv (Array (sh, Int) e))
mkScan uid aenv aR dir combine seed arr
= foldr1 (+++) <$> sequence (codeScanS ++ codeScanP ++ codeScanFill)
where
codeScanS = [ mkScanS dir uid aenv aR combine seed arr ]
codeScanP = case aR of
ArrayR (ShapeRsnoc ShapeRz) eR -> [ mkScanP dir uid aenv eR combine seed arr ]
_ -> []
-- Input can be empty iff a seed is given. We then need to compile a fill kernel
codeScanFill = case seed of
Just s -> [ mkScanFill uid aenv aR s ]
Nothing -> []
-- Variant of 'scanl' where the final result is returned in a separate array.
--
-- > scanr' (+) 10 (use $ fromList (Z :. 10) [0..])
-- >
-- > ==> ( Array (Z :. 10) [10,10,11,13,16,20,25,31,38,46]
-- , Array Z [55]
-- )
--
mkScan'
:: UID
-> Gamma aenv
-> ArrayR (Array (sh, Int) e)
-> Direction
-> IRFun2 Native aenv (e -> e -> e)
-> IRExp Native aenv e
-> MIRDelayed Native aenv (Array (sh, Int) e)
-> CodeGen Native (IROpenAcc Native aenv (Array (sh, Int) e, Array sh e))
mkScan' uid aenv aR dir combine seed arr
| ArrayR (ShapeRsnoc ShapeRz) eR <- aR
= foldr1 (+++) <$> sequence [ mkScan'S dir uid aenv aR combine seed arr
, mkScan'P dir uid aenv eR combine seed arr
, mkScan'Fill uid aenv aR seed
]
--
| otherwise
= (+++) <$> mkScan'S dir uid aenv aR combine seed arr
<*> mkScan'Fill uid aenv aR seed
-- If the innermost dimension of an exclusive scan is empty, then we just fill
-- the result with the seed element.
--
mkScanFill
:: UID
-> Gamma aenv
-> ArrayR (Array sh e)
-> IRExp Native aenv e
-> CodeGen Native (IROpenAcc Native aenv (Array sh e))
mkScanFill uid aenv aR seed =
mkGenerate uid aenv aR (IRFun1 (const seed))
mkScan'Fill
:: UID
-> Gamma aenv
-> ArrayR (Array (sh, Int) e)
-> IRExp Native aenv e
-> CodeGen Native (IROpenAcc Native aenv (Array (sh, Int) e, Array sh e))
mkScan'Fill uid aenv aR seed =
Safe.coerce <$> mkScanFill uid aenv (reduceRank aR) seed
-- A single thread sequentially scans along an entire innermost dimension. For
-- inclusive scans we can assume that the innermost-dimension is at least one
-- element.
--
-- Note that we can use this both when there is a single thread, or in parallel
-- where threads are scheduled over the outer dimensions (segments).
--
mkScanS
:: Direction
-> UID
-> Gamma aenv
-> ArrayR (Array (sh, Int) e)
-> IRFun2 Native aenv (e -> e -> e)
-> MIRExp Native aenv e
-> MIRDelayed Native aenv (Array (sh, Int) e)
-> CodeGen Native (IROpenAcc Native aenv (Array (sh, Int) e))
mkScanS dir uid aenv aR combine mseed marr =
let
(start, end, paramGang) = gangParam shR
(arrOut, paramOut) = mutableArray aR "out"
(arrIn, paramIn) = delayedArray "in" marr
paramEnv = envParam aenv
ShapeRsnoc shR = arrayRshape aR
--
next i = case dir of
LeftToRight -> A.add numType i (liftInt 1)
RightToLeft -> A.sub numType i (liftInt 1)
in
makeOpenAcc uid "scanS" (paramGang ++ paramOut ++ paramIn ++ paramEnv) $ do
-- The dimensions of the input and output arrays are (almost) the same
-- but LLVM can't know that so make it explicit so that we reuse loop
-- variables and index calculations
shIn <- delayedExtent arrIn
let sz = indexHead shIn
shOut = case mseed of
Nothing -> shIn
Just{} -> indexCons (indexTail shIn) (indexHead (irArrayShape arrOut))
-- Loop over the outer dimensions
imapNestFromTo shR start end (indexTail shIn) $ \ix _ -> do
-- index i* is the index that we will read data from. Recall that the
-- supremum index is exclusive
i0 <- case dir of
LeftToRight -> return (liftInt 0)
RightToLeft -> A.sub numType sz (liftInt 1)
-- index j* is the index that we write to. Recall that for exclusive scans
-- the output array inner dimension is one larger than the input.
j0 <- case mseed of
Nothing -> return i0 -- merge 'i' and 'j' indices whenever we can
Just{} -> case dir of
LeftToRight -> return i0
RightToLeft -> return sz
-- Evaluate or read the initial element. Update the read-from index
-- appropriately.
(v0,i1) <- case mseed of
Just seed -> (,) <$> seed <*> pure i0
Nothing -> (,) <$> app1 (delayedIndex arrIn) (indexCons ix i0) <*> next i0
-- Write first element, then continue looping through the rest of
-- this innermost dimension
k0 <- intOfIndex (arrayRshape aR) shOut (indexCons ix j0)
j1 <- next j0
writeArray TypeInt arrOut k0 v0
void $ while (TupRunit `TupRpair` TupRsingle scalarTypeInt `TupRpair` TupRsingle scalarTypeInt `TupRpair` arrayRtype aR)
(\(A.untrip -> (i,_,_)) -> do
case dir of
LeftToRight -> A.lt singleType i sz
RightToLeft -> A.gte singleType i (liftInt 0))
(\(A.untrip -> (i,j,u)) -> do
v <- app1 (delayedIndex arrIn) (indexCons ix i)
w <- case dir of
LeftToRight -> app2 combine u v
RightToLeft -> app2 combine v u
k <- intOfIndex (arrayRshape aR) shOut (indexCons ix j)
writeArray TypeInt arrOut k w
A.trip <$> next i <*> next j <*> pure w)
(A.trip i1 j1 v0)
mkScan'S
:: Direction
-> UID
-> Gamma aenv
-> ArrayR (Array (sh, Int) e)
-> IRFun2 Native aenv (e -> e -> e)
-> IRExp Native aenv e
-> MIRDelayed Native aenv (Array (sh, Int) e)
-> CodeGen Native (IROpenAcc Native aenv (Array (sh, Int) e, Array sh e))
mkScan'S dir uid aenv aR combine seed marr =
let
(start, end, paramGang) = gangParam shR
(arrOut, paramOut) = mutableArray aR "out"
(arrSum, paramSum) = mutableArray (reduceRank aR) "sum"
(arrIn, paramIn) = delayedArray "in" marr
paramEnv = envParam aenv
ShapeRsnoc shR = arrayRshape aR
--
next i = case dir of
LeftToRight -> A.add numType i (liftInt 1)
RightToLeft -> A.sub numType i (liftInt 1)
in
makeOpenAcc uid "scanS" (paramGang ++ paramOut ++ paramSum ++ paramIn ++ paramEnv) $ do
shIn <- delayedExtent arrIn
let sz = indexHead shIn
shOut = shIn
imapNestFromTo shR start end (indexTail shIn) $ \ix ii -> do
-- index to read data from
i0 <- case dir of
LeftToRight -> return (liftInt 0)
RightToLeft -> A.sub numType sz (liftInt 1)
-- initial element
v0 <- seed
-- Loop through the input. Only at the top of the loop to we write the
-- carry-in value (i.e. value from the last loop iteration) to the output
-- array. This ensures correct behaviour if the input array was empty.
r <- while (TupRsingle scalarTypeInt `TupRpair` arrayRtype aR)
(\(A.unpair -> (i,_)) -> do
case dir of
LeftToRight -> A.lt singleType i sz
RightToLeft -> A.gte singleType i (liftInt 0))
(\(A.unpair -> (i,u)) -> do
k <- intOfIndex (arrayRshape aR) shOut (indexCons ix i)
writeArray TypeInt arrOut k u
v <- app1 (delayedIndex arrIn) (indexCons ix i)
w <- case dir of
LeftToRight -> app2 combine u v
RightToLeft -> app2 combine v u
A.pair <$> next i <*> pure w)
(A.pair i0 v0)
writeArray TypeInt arrSum ii (A.snd r)
mkScanP
:: Direction
-> UID
-> Gamma aenv
-> TypeR e
-> IRFun2 Native aenv (e -> e -> e)
-> MIRExp Native aenv e
-> MIRDelayed Native aenv (Vector e)
-> CodeGen Native (IROpenAcc Native aenv (Vector e))
mkScanP dir uid aenv eR combine mseed marr =
foldr1 (+++) <$> sequence [ mkScanP1 dir uid aenv eR combine mseed marr
, mkScanP2 dir uid aenv eR combine
, mkScanP3 dir uid aenv eR combine mseed
]
-- Parallel scan, step 1.
--
-- Threads scan a stripe of the input into a temporary array, incorporating the
-- initial element and any fused functions on the way. The final reduction
-- result of this chunk is written to a separate array.
--
mkScanP1
:: Direction
-> UID
-> Gamma aenv
-> TypeR e
-> IRFun2 Native aenv (e -> e -> e)
-> MIRExp Native aenv e
-> MIRDelayed Native aenv (Vector e)
-> CodeGen Native (IROpenAcc Native aenv (Vector e))
mkScanP1 dir uid aenv eR combine mseed marr =
let
(start, end, paramGang) = gangParam dim1
(arrOut, paramOut) = mutableArray (ArrayR dim1 eR) "out"
(arrTmp, paramTmp) = mutableArray (ArrayR dim1 eR) "tmp"
(arrIn, paramIn) = delayedArray "in" marr
paramEnv = envParam aenv
--
steps = local (TupRsingle scalarTypeInt) "ix.steps"
paramSteps = parameter (TupRsingle scalarTypeInt) "ix.steps"
piece = local (TupRsingle scalarTypeInt) "ix.piece"
paramPiece = parameter (TupRsingle scalarTypeInt) "ix.piece"
--
next i = case dir of
LeftToRight -> A.add numType i (liftInt 1)
RightToLeft -> A.sub numType i (liftInt 1)
firstPiece = case dir of
LeftToRight -> liftInt 0
RightToLeft -> steps
in
makeOpenAcc uid "scanP1" (paramGang ++ paramPiece ++ paramSteps ++ paramOut ++ paramTmp ++ paramIn ++ paramEnv) $ do
-- A thread scans a non-empty stripe of the input, storing the final
-- reduction result into a separate array.
--
-- For exclusive scans the first chunk must incorporate the initial element
-- into the input and output, while all other chunks increment their output
-- index by one.
--
-- index i* is the index that we read data from. Recall that the supremum
-- index is exclusive
i0 <- case dir of
LeftToRight -> return (indexHead start)
RightToLeft -> next (indexHead end)
-- index j* is the index that we write to. Recall that for exclusive scan
-- the output array is one larger than the input; the first piece uses
-- this spot to write the initial element, all other chunks shift by one.
j0 <- case mseed of
Nothing -> return i0
Just _ -> case dir of
LeftToRight -> if (TupRsingle scalarTypeInt, A.eq singleType piece firstPiece)
then return i0
else next i0
RightToLeft -> if (TupRsingle scalarTypeInt, A.eq singleType piece firstPiece)
then return (indexHead end)
else return i0
-- Evaluate/read the initial element for this piece. Update the read-from
-- index appropriately
(v0,i1) <- A.unpair <$> case mseed of
Just seed -> if (eR `TupRpair` TupRsingle scalarTypeInt, A.eq singleType piece firstPiece)
then A.pair <$> seed <*> pure i0
else A.pair <$> app1 (delayedLinearIndex arrIn) i0 <*> next i0
Nothing -> A.pair <$> app1 (delayedLinearIndex arrIn) i0 <*> next i0
-- Write first element
writeArray TypeInt arrOut j0 v0
j1 <- next j0
-- Continue looping through the rest of the input
let cont i =
case dir of
LeftToRight -> A.lt singleType i (indexHead end)
RightToLeft -> A.gte singleType i (indexHead start)
r <- while (TupRunit `TupRpair` TupRsingle scalarTypeInt `TupRpair` TupRsingle scalarTypeInt `TupRpair` eR)
(cont . A.fst3)
(\(A.untrip -> (i,j,v)) -> do
u <- app1 (delayedLinearIndex arrIn) i
v' <- case dir of
LeftToRight -> app2 combine v u
RightToLeft -> app2 combine u v
writeArray TypeInt arrOut j v'
A.trip <$> next i <*> next j <*> pure v')
(A.trip i1 j1 v0)
-- Final reduction result of this piece
writeArray TypeInt arrTmp piece (A.thd3 r)
-- Parallel scan, step 2.
--
-- A single thread performs an in-place inclusive scan of the partial block
-- sums. This forms the carry-in value which are added to the stripe partial
-- results in the final step.
--
mkScanP2
:: Direction
-> UID
-> Gamma aenv
-> TypeR e
-> IRFun2 Native aenv (e -> e -> e)
-> CodeGen Native (IROpenAcc Native aenv (Vector e))
mkScanP2 dir uid aenv eR combine =
let
(start, end, paramGang) = gangParam dim1
(arrTmp, paramTmp) = mutableArray (ArrayR dim1 eR) "tmp"
paramEnv = envParam aenv
--
cont i = case dir of
LeftToRight -> A.lt singleType i (indexHead end)
RightToLeft -> A.gte singleType i (indexHead start)
next i = case dir of
LeftToRight -> A.add numType i (liftInt 1)
RightToLeft -> A.sub numType i (liftInt 1)
in
makeOpenAcc uid "scanP2" (paramGang ++ paramTmp ++ paramEnv) $ do
i0 <- case dir of
LeftToRight -> return (indexHead start)
RightToLeft -> next (indexHead end)
v0 <- readArray TypeInt arrTmp i0
i1 <- next i0
void $ while (TupRsingle scalarTypeInt `TupRpair` eR)
(cont . A.fst)
(\(A.unpair -> (i,v)) -> do
u <- readArray TypeInt arrTmp i
i' <- next i
v' <- case dir of
LeftToRight -> app2 combine v u
RightToLeft -> app2 combine u v
writeArray TypeInt arrTmp i v'
return $ A.pair i' v')
(A.pair i1 v0)
-- Parallel scan, step 3.
--
-- Threads combine every element of the partial block results with the carry-in
-- value computed from step 2.
--
-- Note that first chunk does not need extra processing (has no carry-in value).
--
mkScanP3
:: Direction
-> UID
-> Gamma aenv
-> TypeR e
-> IRFun2 Native aenv (e -> e -> e)
-> MIRExp Native aenv e
-> CodeGen Native (IROpenAcc Native aenv (Vector e))
mkScanP3 dir uid aenv eR combine mseed =
let
(start, end, paramGang) = gangParam dim1
(arrOut, paramOut) = mutableArray (ArrayR dim1 eR) "out"
(arrTmp, paramTmp) = mutableArray (ArrayR dim1 eR) "tmp"
paramEnv = envParam aenv
--
steps = local (TupRsingle scalarTypeInt) "ix.steps"
paramSteps = parameter (TupRsingle scalarTypeInt) "ix.steps"
piece = local (TupRsingle scalarTypeInt) "ix.piece"
paramPiece = parameter (TupRsingle scalarTypeInt) "ix.piece"
--
next i = case dir of
LeftToRight -> A.add numType i (liftInt 1)
RightToLeft -> A.sub numType i (liftInt 1)
prev i = case dir of
LeftToRight -> A.sub numType i (liftInt 1)
RightToLeft -> A.add numType i (liftInt 1)
firstPiece = case dir of
LeftToRight -> liftInt 0
RightToLeft -> steps
in
makeOpenAcc uid "scanP3" (paramGang ++ paramPiece ++ paramSteps ++ paramOut ++ paramTmp ++ paramEnv) $ do
-- TODO: Don't schedule the "first" piece. In the scheduler this corresponds
-- to the split range with the smallest/largest linear index for left/right
-- scans respectively. For right scans this is not necessarily the last piece(?).
--
A.when (neq singleType piece firstPiece) $ do
-- Compute start and end indices, leaving space for the initial element
(inf,sup) <- case (dir, mseed) of
(LeftToRight, Just{}) -> (,) <$> next (indexHead start) <*> next (indexHead end)
_ -> (,) <$> pure (indexHead start) <*> pure (indexHead end)
-- Read in the carry in value for this piece
c <- readArray TypeInt arrTmp =<< prev piece
imapFromTo inf sup $ \i -> do
x <- readArray TypeInt arrOut i
y <- case dir of
LeftToRight -> app2 combine c x
RightToLeft -> app2 combine x c
writeArray TypeInt arrOut i y
mkScan'P
:: Direction
-> UID
-> Gamma aenv
-> TypeR e
-> IRFun2 Native aenv (e -> e -> e)
-> IRExp Native aenv e
-> MIRDelayed Native aenv (Vector e)
-> CodeGen Native (IROpenAcc Native aenv (Vector e, Scalar e))
mkScan'P dir uid aenv eR combine seed arr =
foldr1 (+++) <$> sequence [ mkScan'P1 dir uid aenv eR combine seed arr
, mkScan'P2 dir uid aenv eR combine
, mkScan'P3 dir uid aenv eR combine
]
-- Parallel scan', step 1
--
-- Threads scan a stripe of the input into a temporary array. Similar to
-- exclusive scan, the output indices are shifted by one relative to the input
-- indices to make space for the initial element.
--
mkScan'P1
:: Direction
-> UID
-> Gamma aenv
-> TypeR e
-> IRFun2 Native aenv (e -> e -> e)
-> IRExp Native aenv e
-> MIRDelayed Native aenv (Vector e)
-> CodeGen Native (IROpenAcc Native aenv (Vector e, Scalar e))
mkScan'P1 dir uid aenv eR combine seed marr =
let
(start, end, paramGang) = gangParam dim1
(arrOut, paramOut) = mutableArray (ArrayR dim1 eR) "out"
(arrTmp, paramTmp) = mutableArray (ArrayR dim1 eR) "tmp"
(arrIn, paramIn) = delayedArray "in" marr
paramEnv = envParam aenv
--
steps = local (TupRsingle scalarTypeInt) "ix.steps"
paramSteps = parameter (TupRsingle scalarTypeInt) "ix.steps"
piece = local (TupRsingle scalarTypeInt) "ix.piece"
paramPiece = parameter (TupRsingle scalarTypeInt) "ix.piece"
--
next i = case dir of
LeftToRight -> A.add numType i (liftInt 1)
RightToLeft -> A.sub numType i (liftInt 1)
firstPiece = case dir of
LeftToRight -> liftInt 0
RightToLeft -> steps
in
makeOpenAcc uid "scanP1" (paramGang ++ paramPiece ++ paramSteps ++ paramOut ++ paramTmp ++ paramIn ++ paramEnv) $ do
-- index i* is the index that we pull data from.
i0 <- case dir of
LeftToRight -> return (indexHead start)
RightToLeft -> next (indexHead end)
-- index j* is the index that we write results to. The first piece needs to
-- include the initial element, and all other chunks shift their results
-- across by one to make space.
j0 <- if (TupRsingle scalarTypeInt, A.eq singleType piece firstPiece)
then pure i0
else next i0
-- Evaluate/read the initial element. Update the read-from index
-- appropriately.
(v0,i1) <- A.unpair <$> if (eR `TupRpair` TupRsingle scalarTypeInt, A.eq singleType piece firstPiece)
then A.pair <$> seed <*> pure i0
else A.pair <$> app1 (delayedLinearIndex arrIn) i0 <*> pure j0
-- Write the first element
writeArray TypeInt arrOut j0 v0
j1 <- next j0
-- Continue looping through the rest of the input
let cont i =
case dir of
LeftToRight -> A.lt singleType i (indexHead end)
RightToLeft -> A.gte singleType i (indexHead start)
r <- while (TupRunit `TupRpair` TupRsingle scalarTypeInt `TupRpair` TupRsingle scalarTypeInt `TupRpair` eR)
(cont . A.fst3)
(\(A.untrip-> (i,j,v)) -> do
u <- app1 (delayedLinearIndex arrIn) i
v' <- case dir of
LeftToRight -> app2 combine v u
RightToLeft -> app2 combine u v
writeArray TypeInt arrOut j v'
A.trip <$> next i <*> next j <*> pure v')
(A.trip i1 j1 v0)
-- Write the final reduction result of this piece
writeArray TypeInt arrTmp piece (A.thd3 r)
-- Parallel scan', step 2
--
-- Identical to mkScanP2, except we store the total scan result into a separate
-- array (rather than discard it).
--
mkScan'P2
:: Direction
-> UID
-> Gamma aenv
-> TypeR e
-> IRFun2 Native aenv (e -> e -> e)
-> CodeGen Native (IROpenAcc Native aenv (Vector e, Scalar e))
mkScan'P2 dir uid aenv eR combine =
let
(start, end, paramGang) = gangParam dim1
(arrTmp, paramTmp) = mutableArray (ArrayR dim1 eR) "tmp"
(arrSum, paramSum) = mutableArray (ArrayR dim0 eR) "sum"
paramEnv = envParam aenv
--
cont i = case dir of
LeftToRight -> A.lt singleType i (indexHead end)
RightToLeft -> A.gte singleType i (indexHead start)
next i = case dir of
LeftToRight -> A.add numType i (liftInt 1)
RightToLeft -> A.sub numType i (liftInt 1)
in
makeOpenAcc uid "scanP2" (paramGang ++ paramSum ++ paramTmp ++ paramEnv) $ do
i0 <- case dir of
LeftToRight -> return (indexHead start)
RightToLeft -> next (indexHead end)
v0 <- readArray TypeInt arrTmp i0
i1 <- next i0
r <- while (TupRpair (TupRsingle scalarTypeInt) eR)
(cont . A.fst)
(\(A.unpair -> (i,v)) -> do
u <- readArray TypeInt arrTmp i
i' <- next i
v' <- case dir of
LeftToRight -> app2 combine v u
RightToLeft -> app2 combine u v
writeArray TypeInt arrTmp i v'
return $ A.pair i' v')
(A.pair i1 v0)
writeArray TypeInt arrSum (liftInt 0) (A.snd r)
-- Parallel scan', step 3
--
-- Similar to mkScanP3, except that indices are shifted by one since the output
-- array is the same size as the input (despite being an exclusive scan).
--
-- Note that the first chunk does not need to do any extra processing (has no
-- carry-in value).
--
mkScan'P3
:: Direction
-> UID
-> Gamma aenv
-> TypeR e
-> IRFun2 Native aenv (e -> e -> e)
-> CodeGen Native (IROpenAcc Native aenv (Vector e, Scalar e))
mkScan'P3 dir uid aenv eR combine =
let
(start, end, paramGang) = gangParam dim1
(arrOut, paramOut) = mutableArray (ArrayR dim1 eR) "out"
(arrTmp, paramTmp) = mutableArray (ArrayR dim1 eR) "tmp"
paramEnv = envParam aenv
--
steps = local (TupRsingle scalarTypeInt) "ix.steps"
paramSteps = parameter (TupRsingle scalarTypeInt) "ix.steps"
piece = local (TupRsingle scalarTypeInt) "ix.piece"
paramPiece = parameter (TupRsingle scalarTypeInt) "ix.piece"
--
next i = case dir of
LeftToRight -> A.add numType i (liftInt 1)
RightToLeft -> A.sub numType i (liftInt 1)
prev i = case dir of
LeftToRight -> A.sub numType i (liftInt 1)
RightToLeft -> A.add numType i (liftInt 1)
firstPiece = case dir of
LeftToRight -> liftInt 0
RightToLeft -> steps
in
makeOpenAcc uid "scanP3" (paramGang ++ paramPiece ++ paramSteps ++ paramOut ++ paramTmp ++ paramEnv) $ do
-- TODO: don't schedule the "first" piece.
--
A.when (neq singleType piece firstPiece) $ do
-- Compute start and end indices, leaving space for the initial element
inf <- next (indexHead start)
sup <- next (indexHead end)
-- Read the carry-in value for this piece
c <- readArray TypeInt arrTmp =<< prev piece
-- Apply the carry-in value to all elements of the output
imapFromTo inf sup $ \i -> do
x <- readArray TypeInt arrOut i
y <- case dir of
LeftToRight -> app2 combine c x
RightToLeft -> app2 combine x c
writeArray TypeInt arrOut i y