accelerate-llvm-native-1.4.0.0: src/Data/Array/Accelerate/LLVM/Native/CodeGen/Fold.hs
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
-- |
-- Module : Data.Array.Accelerate.LLVM.Native.CodeGen.Fold
-- 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.Fold
where
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.Constant
import Data.Array.Accelerate.LLVM.CodeGen.Environment
import Data.Array.Accelerate.LLVM.CodeGen.Exp
import Data.Array.Accelerate.LLVM.CodeGen.IR
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 Prelude as P hiding ( length )
-- Reduce an array along the innermost dimension. The reduction
-- function must be associative to allow for an efficient parallel
-- implementation. When an initial value is given, the input can be
-- empty. The initial element does not need to be a neutral element of
-- the operator. When no initial value is given, the array must be
-- non-empty
--
mkFold
:: UID
-> Gamma aenv
-> ArrayR (Array sh e)
-> 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 e))
mkFold uid aenv aR f z arr =
(+++) <$> case aR of
ArrayR ShapeRz eR -> mkFoldAll uid aenv eR f z arr
_ -> mkFoldDim uid aenv aR f z arr
<*> case z of
Just z' -> mkFoldFill uid aenv aR z'
Nothing -> return (IROpenAcc [])
-- Reduce a multidimensional (>1) array along the innermost dimension.
--
-- For simplicity, each element of the output (reduction along the entire length
-- of an innermost-dimension index) is computed by a single thread.
--
mkFoldDim
:: UID
-> Gamma aenv
-> ArrayR (Array sh 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 e))
mkFoldDim uid aenv aR@(ArrayR shR eR) combine mseed mdelayed =
let
(start, end, paramGang) = gangParam shR
(arrOut, paramOut) = mutableArray aR "out"
(arrIn, paramIn) = delayedArray "in" mdelayed
paramEnv = envParam aenv
zero = liftInt 0
in
makeOpenAcc uid "fold" (paramGang ++ paramOut ++ paramIn ++ paramEnv) $ do
sz <- indexHead <$> delayedExtent arrIn
imapNestFromTo shR start end (irArrayShape arrOut) $ \ix i -> do
r <- case mseed of
Just seed -> do z <- seed
reduceFromTo eR zero sz (app2 combine) z (app1 (delayedIndex arrIn) . indexCons ix)
Nothing -> reduce1FromTo eR zero sz (app2 combine) (app1 (delayedIndex arrIn) . indexCons ix)
writeArray TypeInt arrOut i r
-- Reduce an array to single element.
--
-- Since reductions consume arrays that have been fused into them,
-- a parallel fold requires two passes. At an example, take vector dot
-- product:
--
-- > dotp xs ys = fold (+) 0 (zipWith (*) xs ys)
--
-- 1. The first pass reads in the fused array data, in this case corresponding
-- to the function (\i -> (xs!i) * (ys!i)).
--
-- 2. The second pass reads in the manifest array data from the first step and
-- directly reduces the array. This second step should be small and so is
-- usually just done by a single core.
--
-- Note that the first step is split into two kernels, the second of which
-- reads a carry-in value of that thread's partial reduction, so that
-- threads can still participate in work-stealing. These kernels must not
-- be invoked over empty ranges.
--
-- The final step is sequential reduction of the partial results. If this
-- is an exclusive reduction, the seed element is included at this point.
--
mkFoldAll
:: UID
-> Gamma aenv -- ^ array environment
-> TypeR e
-> IRFun2 Native aenv (e -> e -> e) -- ^ combination function
-> MIRExp Native aenv e -- ^ seed element, if this is an exclusive reduction
-> MIRDelayed Native aenv (Vector e) -- ^ input data
-> CodeGen Native (IROpenAcc Native aenv (Scalar e))
mkFoldAll uid aenv eR combine mseed mdelayed =
foldr1 (+++) <$> sequence [ mkFoldAllS uid aenv eR combine mseed mdelayed
, mkFoldAllP1 uid aenv eR combine mdelayed
, mkFoldAllP2 uid aenv eR combine mseed
]
-- Sequential reduction of an entire array to a single element
--
mkFoldAllS
:: UID
-> Gamma aenv -- ^ array environment
-> TypeR e
-> IRFun2 Native aenv (e -> e -> e) -- ^ combination function
-> MIRExp Native aenv e -- ^ seed element, if this is an exclusive reduction
-> MIRDelayed Native aenv (Vector e) -- ^ input data
-> CodeGen Native (IROpenAcc Native aenv (Scalar e))
mkFoldAllS uid aenv eR combine mseed mdelayed =
let
(start, end, paramGang) = gangParam dim1
(arrOut, paramOut) = mutableArray (ArrayR dim0 eR) "out"
(arrIn, paramIn) = delayedArray "in" mdelayed
paramEnv = envParam aenv
zero = liftInt 0
in
makeOpenAcc uid "foldAllS" (paramGang ++ paramOut ++ paramIn ++ paramEnv) $ do
r <- case mseed of
Just seed -> do z <- seed
reduceFromTo eR (indexHead start) (indexHead end) (app2 combine) z (app1 (delayedLinearIndex arrIn))
Nothing -> reduce1FromTo eR (indexHead start) (indexHead end) (app2 combine) (app1 (delayedLinearIndex arrIn))
writeArray TypeInt arrOut zero r
-- Parallel reduction of an entire array to a single element, step 1.
--
-- Threads reduce each stripe of the input into a temporary array, incorporating
-- any fused functions on the way.
--
mkFoldAllP1
:: UID
-> Gamma aenv -- ^ array environment
-> TypeR e
-> IRFun2 Native aenv (e -> e -> e) -- ^ combination function
-> MIRDelayed Native aenv (Vector e) -- ^ input data
-> CodeGen Native (IROpenAcc Native aenv (Scalar e))
mkFoldAllP1 uid aenv eR combine mdelayed =
let
(start, end, paramGang) = gangParam dim1
(arrTmp, paramTmp) = mutableArray (ArrayR dim1 eR) "tmp"
(arrIn, paramIn) = delayedArray "in" mdelayed
piece = local (TupRsingle scalarTypeInt) "ix.piece"
paramPiece = parameter (TupRsingle scalarTypeInt) "ix.piece"
paramEnv = envParam aenv
in
makeOpenAcc uid "foldAllP1" (paramGang ++ paramPiece ++ paramTmp ++ paramIn ++ paramEnv) $ do
-- A thread reduces a sequential (non-empty) stripe of the input and stores
-- that value into a temporary array at a specific index. This method thus
-- supports non-commutative operators because the order of operations
-- remains left-to-right.
--
r <- reduce1FromTo eR (indexHead start) (indexHead end) (app2 combine) (app1 (delayedLinearIndex arrIn))
writeArray TypeInt arrTmp piece r
-- Parallel reduction of an entire array to a single element, step 2.
--
-- A single thread reduces the temporary array to a single element.
--
-- During execution, we choose a stripe size in phase 1 so that the temporary is
-- small-ish and thus suitable for sequential reduction. An alternative would be
-- to keep the stripe size constant and, for if the partial reductions array is
-- large, continuing reducing it in parallel.
--
mkFoldAllP2
:: UID
-> Gamma aenv -- ^ array environment
-> TypeR e
-> IRFun2 Native aenv (e -> e -> e) -- ^ combination function
-> MIRExp Native aenv e -- ^ seed element, if this is an exclusive reduction
-> CodeGen Native (IROpenAcc Native aenv (Scalar e))
mkFoldAllP2 uid aenv eR combine mseed =
let
(start, end, paramGang) = gangParam dim1
(arrTmp, paramTmp) = mutableArray (ArrayR dim1 eR) "tmp"
(arrOut, paramOut) = mutableArray (ArrayR dim0 eR) "out"
paramEnv = envParam aenv
zero = liftInt 0
in
makeOpenAcc uid "foldAllP2" (paramGang ++ paramTmp ++ paramOut ++ paramEnv) $ do
r <- case mseed of
Just seed -> do z <- seed
reduceFromTo eR (indexHead start) (indexHead end) (app2 combine) z (readArray TypeInt arrTmp)
Nothing -> reduce1FromTo eR (indexHead start) (indexHead end) (app2 combine) (readArray TypeInt arrTmp)
writeArray TypeInt arrOut zero r
-- Exclusive reductions over empty arrays (of any dimension) fill the lower
-- dimensions with the initial element
--
mkFoldFill
:: UID
-> Gamma aenv
-> ArrayR (Array sh e)
-> IRExp Native aenv e
-> CodeGen Native (IROpenAcc Native aenv (Array sh e))
mkFoldFill uid aenv aR seed =
mkGenerate uid aenv aR (IRFun1 (const seed))
-- Reduction loops
-- ---------------
-- Reduction of a (possibly empty) index space.
--
reduceFromTo
:: TypeR e
-> Operands Int -- ^ starting index
-> Operands Int -- ^ final index (exclusive)
-> (Operands e -> Operands e -> CodeGen Native (Operands e)) -- ^ combination function
-> Operands e -- ^ initial value
-> (Operands Int -> CodeGen Native (Operands e)) -- ^ function to retrieve element at index
-> CodeGen Native (Operands e)
reduceFromTo eR m n f z get =
iterFromTo eR m n z $ \i acc -> do
x <- get i
y <- f acc x
return y
-- Reduction of an array over a _non-empty_ index space. The array must
-- contain at least one element.
--
reduce1FromTo
:: TypeR e
-> Operands Int -- ^ starting index
-> Operands Int -- ^ final index
-> (Operands e -> Operands e -> CodeGen Native (Operands e)) -- ^ combination function
-> (Operands Int -> CodeGen Native (Operands e)) -- ^ function to retrieve element at index
-> CodeGen Native (Operands e)
reduce1FromTo eR m n f get = do
z <- get m
m1 <- add numType m (ir numType (num numType 1))
reduceFromTo eR m1 n f z get