{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
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 )
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
-> 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
uid Gamma aenv
aenv ArrayR (Array sh e)
aR IRFun2 Native aenv (e -> e -> e)
f Maybe (IRExp Native aenv e)
z MIRDelayed Native aenv (Array (sh, Int) e)
arr =
IROpenAcc Native aenv (Array sh e)
-> IROpenAcc Native aenv (Array sh e)
-> IROpenAcc Native aenv (Array sh e)
forall aenv a.
IROpenAcc Native aenv a
-> IROpenAcc Native aenv a -> IROpenAcc Native aenv a
(+++) (IROpenAcc Native aenv (Array sh e)
-> IROpenAcc Native aenv (Array sh e)
-> IROpenAcc Native aenv (Array sh e))
-> CodeGen Native (IROpenAcc Native aenv (Array sh e))
-> CodeGen
Native
(IROpenAcc Native aenv (Array sh e)
-> IROpenAcc Native aenv (Array sh e))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> case ArrayR (Array sh e)
aR of
ArrayR ShapeR sh
ShapeRz TypeR e
eR -> 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 (Scalar e))
forall aenv e.
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 (Scalar e))
mkFoldAll UID
uid Gamma aenv
aenv TypeR e
eR IRFun2 Native aenv (e -> e -> e)
IRFun2 Native aenv (e -> e -> e)
f Maybe (IRExp Native aenv e)
MIRExp Native aenv e
z MIRDelayed Native aenv (Array (sh, Int) e)
MIRDelayed Native aenv (Vector e)
arr
ArrayR (Array sh e)
_ -> 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))
forall aenv sh e.
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
uid Gamma aenv
aenv ArrayR (Array sh e)
aR IRFun2 Native aenv (e -> e -> e)
f Maybe (IRExp Native aenv e)
z MIRDelayed Native aenv (Array (sh, Int) e)
arr
CodeGen
Native
(IROpenAcc Native aenv (Array sh e)
-> IROpenAcc Native aenv (Array sh e))
-> CodeGen Native (IROpenAcc Native aenv (Array sh e))
-> CodeGen Native (IROpenAcc Native aenv (Array sh e))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> case Maybe (IRExp Native aenv e)
z of
Just IRExp Native aenv e
z' -> UID
-> Gamma aenv
-> ArrayR (Array sh e)
-> IRExp Native aenv e
-> CodeGen Native (IROpenAcc Native aenv (Array sh e))
forall aenv sh e.
UID
-> Gamma aenv
-> ArrayR (Array sh e)
-> IRExp Native aenv e
-> CodeGen Native (IROpenAcc Native aenv (Array sh e))
mkFoldFill UID
uid Gamma aenv
aenv ArrayR (Array sh e)
aR IRExp Native aenv e
z'
Maybe (IRExp Native aenv e)
Nothing -> IROpenAcc Native aenv (Array sh e)
-> CodeGen Native (IROpenAcc Native aenv (Array sh e))
forall (m :: * -> *) a. Monad m => a -> m a
return ([Kernel Native aenv (Array sh e)]
-> IROpenAcc Native aenv (Array sh e)
forall arch aenv arrs.
[Kernel arch aenv arrs] -> IROpenAcc arch aenv arrs
IROpenAcc [])
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
-> 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
uid Gamma aenv
aenv aR :: ArrayR (Array sh e)
aR@(ArrayR ShapeR sh
shR TypeR e
eR) IRFun2 Native aenv (e -> e -> e)
combine MIRExp Native aenv e
mseed MIRDelayed Native aenv (Array (sh, Int) e)
mdelayed =
let
(Operands sh
start, Operands sh
end, [Parameter]
paramGang) = ShapeR sh -> (Operands sh, Operands sh, [Parameter])
forall sh. ShapeR sh -> (Operands sh, Operands sh, [Parameter])
gangParam ShapeR sh
shR
(IRArray (Array sh e)
arrOut, [Parameter]
paramOut) = ArrayR (Array sh e)
-> Name (Array sh e) -> (IRArray (Array sh e), [Parameter])
forall sh e.
ArrayR (Array sh e)
-> Name (Array sh e) -> (IRArray (Array sh e), [Parameter])
mutableArray ArrayR (Array sh e)
aR Name (Array sh e)
"out"
(IRDelayed Native aenv (Array (sh, Int) e)
arrIn, [Parameter]
paramIn) = Name (Array (sh, Int) e)
-> MIRDelayed Native aenv (Array (sh, Int) e)
-> (IRDelayed Native aenv (Array (sh, Int) e), [Parameter])
forall sh e arch aenv.
Name (Array sh e)
-> MIRDelayed arch aenv (Array sh e)
-> (IRDelayed arch aenv (Array sh e), [Parameter])
delayedArray Name (Array (sh, Int) e)
"in" MIRDelayed Native aenv (Array (sh, Int) e)
mdelayed
paramEnv :: [Parameter]
paramEnv = Gamma aenv -> [Parameter]
forall aenv. Gamma aenv -> [Parameter]
envParam Gamma aenv
aenv
zero :: Operands Int
zero = Int -> Operands Int
liftInt Int
0
in
UID
-> Label
-> [Parameter]
-> CodeGen Native ()
-> CodeGen Native (IROpenAcc Native aenv (Array sh e))
forall aenv a.
UID
-> Label
-> [Parameter]
-> CodeGen Native ()
-> CodeGen Native (IROpenAcc Native aenv a)
makeOpenAcc UID
uid Label
"fold" ([Parameter]
paramGang [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramOut [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramIn [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramEnv) (CodeGen Native ()
-> CodeGen Native (IROpenAcc Native aenv (Array sh e)))
-> CodeGen Native ()
-> CodeGen Native (IROpenAcc Native aenv (Array sh e))
forall a b. (a -> b) -> a -> b
$ do
Operands Int
sz <- Operands (sh, Int) -> Operands Int
forall sh sz. Operands (sh, sz) -> Operands sz
indexHead (Operands (sh, Int) -> Operands Int)
-> CodeGen Native (Operands (sh, Int))
-> CodeGen Native (Operands Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IRDelayed Native aenv (Array (sh, Int) e)
-> CodeGen Native (Operands (sh, Int))
forall arch aenv sh e.
IRDelayed arch aenv (Array sh e) -> IRExp arch aenv sh
delayedExtent IRDelayed Native aenv (Array (sh, Int) e)
arrIn
ShapeR sh
-> Operands sh
-> Operands sh
-> Operands sh
-> (Operands sh -> Operands Int -> CodeGen Native ())
-> CodeGen Native ()
forall sh.
ShapeR sh
-> Operands sh
-> Operands sh
-> Operands sh
-> (Operands sh -> Operands Int -> CodeGen Native ())
-> CodeGen Native ()
imapNestFromTo ShapeR sh
shR Operands sh
start Operands sh
end (IRArray (Array sh e) -> Operands sh
forall sh e. IRArray (Array sh e) -> Operands sh
irArrayShape IRArray (Array sh e)
arrOut) ((Operands sh -> Operands Int -> CodeGen Native ())
-> CodeGen Native ())
-> (Operands sh -> Operands Int -> CodeGen Native ())
-> CodeGen Native ()
forall a b. (a -> b) -> a -> b
$ \Operands sh
ix Operands Int
i -> do
Operands e
r <- case MIRExp Native aenv e
mseed of
Just IRExp Native aenv e
seed -> do Operands e
z <- IRExp Native aenv e
seed
TypeR e
-> Operands Int
-> Operands Int
-> (Operands e -> Operands e -> CodeGen Native (Operands e))
-> Operands e
-> (Operands Int -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
forall e.
TypeR e
-> Operands Int
-> Operands Int
-> (Operands e -> Operands e -> CodeGen Native (Operands e))
-> Operands e
-> (Operands Int -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
reduceFromTo TypeR e
eR Operands Int
zero Operands Int
sz (IRFun2 Native aenv (e -> e -> e)
-> Operands e -> Operands e -> IRExp Native aenv e
forall arch env aenv a b c.
IROpenFun2 arch env aenv (a -> b -> c)
-> Operands a -> Operands b -> IROpenExp arch ((env, a), b) aenv c
app2 IRFun2 Native aenv (e -> e -> e)
combine) Operands e
Operands e
z (IROpenFun1 Native () aenv ((sh, Int) -> e)
-> Operands (sh, Int) -> IRExp Native aenv e
forall arch env aenv a b.
IROpenFun1 arch env aenv (a -> b)
-> Operands a -> IROpenExp arch (env, a) aenv b
app1 (IRDelayed Native aenv (Array (sh, Int) e)
-> IROpenFun1 Native () aenv ((sh, Int) -> e)
forall arch aenv sh e.
IRDelayed arch aenv (Array sh e) -> IRFun1 arch aenv (sh -> e)
delayedIndex IRDelayed Native aenv (Array (sh, Int) e)
arrIn) (Operands (sh, Int) -> IRExp Native aenv e)
-> (Operands Int -> Operands (sh, Int))
-> Operands Int
-> IRExp Native aenv e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Operands sh -> Operands Int -> Operands (sh, Int)
forall sh sz. Operands sh -> Operands sz -> Operands (sh, sz)
indexCons Operands sh
ix)
MIRExp Native aenv e
Nothing -> TypeR e
-> Operands Int
-> Operands Int
-> (Operands e -> Operands e -> CodeGen Native (Operands e))
-> (Operands Int -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
forall e.
TypeR e
-> Operands Int
-> Operands Int
-> (Operands e -> Operands e -> CodeGen Native (Operands e))
-> (Operands Int -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
reduce1FromTo TypeR e
eR Operands Int
zero Operands Int
sz (IRFun2 Native aenv (e -> e -> e)
-> Operands e -> Operands e -> IRExp Native aenv e
forall arch env aenv a b c.
IROpenFun2 arch env aenv (a -> b -> c)
-> Operands a -> Operands b -> IROpenExp arch ((env, a), b) aenv c
app2 IRFun2 Native aenv (e -> e -> e)
combine) (IROpenFun1 Native () aenv ((sh, Int) -> e)
-> Operands (sh, Int) -> IRExp Native aenv e
forall arch env aenv a b.
IROpenFun1 arch env aenv (a -> b)
-> Operands a -> IROpenExp arch (env, a) aenv b
app1 (IRDelayed Native aenv (Array (sh, Int) e)
-> IROpenFun1 Native () aenv ((sh, Int) -> e)
forall arch aenv sh e.
IRDelayed arch aenv (Array sh e) -> IRFun1 arch aenv (sh -> e)
delayedIndex IRDelayed Native aenv (Array (sh, Int) e)
arrIn) (Operands (sh, Int) -> IRExp Native aenv e)
-> (Operands Int -> Operands (sh, Int))
-> Operands Int
-> IRExp Native aenv e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Operands sh -> Operands Int -> Operands (sh, Int)
forall sh sz. Operands sh -> Operands sz -> Operands (sh, sz)
indexCons Operands sh
ix)
IntegralType Int
-> IRArray (Array sh e)
-> Operands Int
-> Operands e
-> CodeGen Native ()
forall int sh e arch.
IntegralType int
-> IRArray (Array sh e)
-> Operands int
-> Operands e
-> CodeGen arch ()
writeArray IntegralType Int
TypeInt IRArray (Array sh e)
arrOut Operands Int
i Operands e
Operands e
r
CodeGen Native ()
forall arch. HasCallStack => CodeGen arch ()
return_
mkFoldAll
:: 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 (Scalar e))
mkFoldAll :: 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 (Scalar e))
mkFoldAll UID
uid Gamma aenv
aenv TypeR e
eR IRFun2 Native aenv (e -> e -> e)
combine MIRExp Native aenv e
mseed MIRDelayed Native aenv (Vector e)
mdelayed =
(IROpenAcc Native aenv (Scalar e)
-> IROpenAcc Native aenv (Scalar e)
-> IROpenAcc Native aenv (Scalar e))
-> [IROpenAcc Native aenv (Scalar e)]
-> IROpenAcc Native aenv (Scalar e)
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 IROpenAcc Native aenv (Scalar e)
-> IROpenAcc Native aenv (Scalar e)
-> IROpenAcc Native aenv (Scalar e)
forall aenv a.
IROpenAcc Native aenv a
-> IROpenAcc Native aenv a -> IROpenAcc Native aenv a
(+++) ([IROpenAcc Native aenv (Scalar e)]
-> IROpenAcc Native aenv (Scalar e))
-> CodeGen Native [IROpenAcc Native aenv (Scalar e)]
-> CodeGen Native (IROpenAcc Native aenv (Scalar e))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [CodeGen Native (IROpenAcc Native aenv (Scalar e))]
-> CodeGen Native [IROpenAcc Native aenv (Scalar e)]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [ 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 (Scalar e))
forall aenv e.
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 (Scalar e))
mkFoldAllS UID
uid Gamma aenv
aenv TypeR e
eR IRFun2 Native aenv (e -> e -> e)
combine MIRExp Native aenv e
mseed MIRDelayed Native aenv (Vector e)
mdelayed
, UID
-> Gamma aenv
-> TypeR e
-> IRFun2 Native aenv (e -> e -> e)
-> MIRDelayed Native aenv (Vector e)
-> CodeGen Native (IROpenAcc Native aenv (Scalar e))
forall aenv e.
UID
-> Gamma aenv
-> TypeR e
-> IRFun2 Native aenv (e -> e -> e)
-> MIRDelayed Native aenv (Vector e)
-> CodeGen Native (IROpenAcc Native aenv (Scalar e))
mkFoldAllP1 UID
uid Gamma aenv
aenv TypeR e
eR IRFun2 Native aenv (e -> e -> e)
combine MIRDelayed Native aenv (Vector e)
mdelayed
, UID
-> Gamma aenv
-> TypeR e
-> IRFun2 Native aenv (e -> e -> e)
-> MIRExp Native aenv e
-> CodeGen Native (IROpenAcc Native aenv (Scalar e))
forall aenv e.
UID
-> Gamma aenv
-> TypeR e
-> IRFun2 Native aenv (e -> e -> e)
-> MIRExp Native aenv e
-> CodeGen Native (IROpenAcc Native aenv (Scalar e))
mkFoldAllP2 UID
uid Gamma aenv
aenv TypeR e
eR IRFun2 Native aenv (e -> e -> e)
combine MIRExp Native aenv e
mseed
]
mkFoldAllS
:: 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 (Scalar e))
mkFoldAllS :: 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 (Scalar e))
mkFoldAllS UID
uid Gamma aenv
aenv TypeR e
eR IRFun2 Native aenv (e -> e -> e)
combine MIRExp Native aenv e
mseed MIRDelayed Native aenv (Vector e)
mdelayed =
let
(Operands DIM1
start, Operands DIM1
end, [Parameter]
paramGang) = ShapeR DIM1 -> (Operands DIM1, Operands DIM1, [Parameter])
forall sh. ShapeR sh -> (Operands sh, Operands sh, [Parameter])
gangParam ShapeR DIM1
dim1
(IRArray (Scalar e)
arrOut, [Parameter]
paramOut) = ArrayR (Scalar e)
-> Name (Scalar e) -> (IRArray (Scalar e), [Parameter])
forall sh e.
ArrayR (Array sh e)
-> Name (Array sh e) -> (IRArray (Array sh e), [Parameter])
mutableArray (ShapeR () -> TypeR e -> ArrayR (Scalar e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR ()
dim0 TypeR e
eR) Name (Scalar e)
"out"
(IRDelayed Native aenv (Vector e)
arrIn, [Parameter]
paramIn) = Name (Vector e)
-> MIRDelayed Native aenv (Vector e)
-> (IRDelayed Native aenv (Vector e), [Parameter])
forall sh e arch aenv.
Name (Array sh e)
-> MIRDelayed arch aenv (Array sh e)
-> (IRDelayed arch aenv (Array sh e), [Parameter])
delayedArray Name (Vector e)
"in" MIRDelayed Native aenv (Vector e)
mdelayed
paramEnv :: [Parameter]
paramEnv = Gamma aenv -> [Parameter]
forall aenv. Gamma aenv -> [Parameter]
envParam Gamma aenv
aenv
zero :: Operands Int
zero = Int -> Operands Int
liftInt Int
0
in
UID
-> Label
-> [Parameter]
-> CodeGen Native ()
-> CodeGen Native (IROpenAcc Native aenv (Scalar e))
forall aenv a.
UID
-> Label
-> [Parameter]
-> CodeGen Native ()
-> CodeGen Native (IROpenAcc Native aenv a)
makeOpenAcc UID
uid Label
"foldAllS" ([Parameter]
paramGang [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramOut [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramIn [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramEnv) (CodeGen Native ()
-> CodeGen Native (IROpenAcc Native aenv (Scalar e)))
-> CodeGen Native ()
-> CodeGen Native (IROpenAcc Native aenv (Scalar e))
forall a b. (a -> b) -> a -> b
$ do
Operands e
r <- case MIRExp Native aenv e
mseed of
Just CodeGen Native (Operands e)
seed -> do Operands e
z <- CodeGen Native (Operands e)
seed
TypeR e
-> Operands Int
-> Operands Int
-> (Operands e -> Operands e -> CodeGen Native (Operands e))
-> Operands e
-> (Operands Int -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
forall e.
TypeR e
-> Operands Int
-> Operands Int
-> (Operands e -> Operands e -> CodeGen Native (Operands e))
-> Operands e
-> (Operands Int -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
reduceFromTo TypeR e
eR (Operands DIM1 -> Operands Int
forall sh sz. Operands (sh, sz) -> Operands sz
indexHead Operands DIM1
start) (Operands DIM1 -> Operands Int
forall sh sz. Operands (sh, sz) -> Operands sz
indexHead Operands DIM1
end) (IRFun2 Native aenv (e -> e -> e)
-> Operands e -> Operands e -> CodeGen Native (Operands e)
forall arch env aenv a b c.
IROpenFun2 arch env aenv (a -> b -> c)
-> Operands a -> Operands b -> IROpenExp arch ((env, a), b) aenv c
app2 IRFun2 Native aenv (e -> e -> e)
combine) Operands e
z (IROpenFun1 Native () aenv (Int -> e)
-> Operands Int -> CodeGen Native (Operands e)
forall arch env aenv a b.
IROpenFun1 arch env aenv (a -> b)
-> Operands a -> IROpenExp arch (env, a) aenv b
app1 (IRDelayed Native aenv (Vector e)
-> IROpenFun1 Native () aenv (Int -> e)
forall arch aenv sh e.
IRDelayed arch aenv (Array sh e) -> IRFun1 arch aenv (Int -> e)
delayedLinearIndex IRDelayed Native aenv (Vector e)
arrIn))
MIRExp Native aenv e
Nothing -> TypeR e
-> Operands Int
-> Operands Int
-> (Operands e -> Operands e -> CodeGen Native (Operands e))
-> (Operands Int -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
forall e.
TypeR e
-> Operands Int
-> Operands Int
-> (Operands e -> Operands e -> CodeGen Native (Operands e))
-> (Operands Int -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
reduce1FromTo TypeR e
eR (Operands DIM1 -> Operands Int
forall sh sz. Operands (sh, sz) -> Operands sz
indexHead Operands DIM1
start) (Operands DIM1 -> Operands Int
forall sh sz. Operands (sh, sz) -> Operands sz
indexHead Operands DIM1
end) (IRFun2 Native aenv (e -> e -> e)
-> Operands e -> Operands e -> CodeGen Native (Operands e)
forall arch env aenv a b c.
IROpenFun2 arch env aenv (a -> b -> c)
-> Operands a -> Operands b -> IROpenExp arch ((env, a), b) aenv c
app2 IRFun2 Native aenv (e -> e -> e)
combine) (IROpenFun1 Native () aenv (Int -> e)
-> Operands Int -> CodeGen Native (Operands e)
forall arch env aenv a b.
IROpenFun1 arch env aenv (a -> b)
-> Operands a -> IROpenExp arch (env, a) aenv b
app1 (IRDelayed Native aenv (Vector e)
-> IROpenFun1 Native () aenv (Int -> e)
forall arch aenv sh e.
IRDelayed arch aenv (Array sh e) -> IRFun1 arch aenv (Int -> e)
delayedLinearIndex IRDelayed Native aenv (Vector e)
arrIn))
IntegralType Int
-> IRArray (Scalar e)
-> Operands Int
-> Operands e
-> CodeGen Native ()
forall int sh e arch.
IntegralType int
-> IRArray (Array sh e)
-> Operands int
-> Operands e
-> CodeGen arch ()
writeArray IntegralType Int
TypeInt IRArray (Scalar e)
arrOut Operands Int
zero Operands e
r
CodeGen Native ()
forall arch. HasCallStack => CodeGen arch ()
return_
mkFoldAllP1
:: UID
-> Gamma aenv
-> TypeR e
-> IRFun2 Native aenv (e -> e -> e)
-> MIRDelayed Native aenv (Vector e)
-> CodeGen Native (IROpenAcc Native aenv (Scalar e))
mkFoldAllP1 :: UID
-> Gamma aenv
-> TypeR e
-> IRFun2 Native aenv (e -> e -> e)
-> MIRDelayed Native aenv (Vector e)
-> CodeGen Native (IROpenAcc Native aenv (Scalar e))
mkFoldAllP1 UID
uid Gamma aenv
aenv TypeR e
eR IRFun2 Native aenv (e -> e -> e)
combine MIRDelayed Native aenv (Vector e)
mdelayed =
let
(Operands DIM1
start, Operands DIM1
end, [Parameter]
paramGang) = ShapeR DIM1 -> (Operands DIM1, Operands DIM1, [Parameter])
forall sh. ShapeR sh -> (Operands sh, Operands sh, [Parameter])
gangParam ShapeR DIM1
dim1
(IRArray (Vector e)
arrTmp, [Parameter]
paramTmp) = ArrayR (Vector e)
-> Name (Vector e) -> (IRArray (Vector e), [Parameter])
forall sh e.
ArrayR (Array sh e)
-> Name (Array sh e) -> (IRArray (Array sh e), [Parameter])
mutableArray (ShapeR DIM1 -> TypeR e -> ArrayR (Vector e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR DIM1
dim1 TypeR e
eR) Name (Vector e)
"tmp"
(IRDelayed Native aenv (Vector e)
arrIn, [Parameter]
paramIn) = Name (Vector e)
-> MIRDelayed Native aenv (Vector e)
-> (IRDelayed Native aenv (Vector e), [Parameter])
forall sh e arch aenv.
Name (Array sh e)
-> MIRDelayed arch aenv (Array sh e)
-> (IRDelayed arch aenv (Array sh e), [Parameter])
delayedArray Name (Vector e)
"in" MIRDelayed Native aenv (Vector e)
mdelayed
piece :: Operands Int
piece = TypeR Int -> Name Int -> Operands Int
forall a. TypeR a -> Name a -> Operands a
local (ScalarType Int -> TypeR Int
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle ScalarType Int
scalarTypeInt) Name Int
"ix.piece"
paramPiece :: [Parameter]
paramPiece = TypeR Int -> Name Int -> [Parameter]
forall t. TypeR t -> Name t -> [Parameter]
parameter (ScalarType Int -> TypeR Int
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle ScalarType Int
scalarTypeInt) Name Int
"ix.piece"
paramEnv :: [Parameter]
paramEnv = Gamma aenv -> [Parameter]
forall aenv. Gamma aenv -> [Parameter]
envParam Gamma aenv
aenv
in
UID
-> Label
-> [Parameter]
-> CodeGen Native ()
-> CodeGen Native (IROpenAcc Native aenv (Scalar e))
forall aenv a.
UID
-> Label
-> [Parameter]
-> CodeGen Native ()
-> CodeGen Native (IROpenAcc Native aenv a)
makeOpenAcc UID
uid Label
"foldAllP1" ([Parameter]
paramGang [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramPiece [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramTmp [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramIn [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramEnv) (CodeGen Native ()
-> CodeGen Native (IROpenAcc Native aenv (Scalar e)))
-> CodeGen Native ()
-> CodeGen Native (IROpenAcc Native aenv (Scalar e))
forall a b. (a -> b) -> a -> b
$ do
Operands e
r <- TypeR e
-> Operands Int
-> Operands Int
-> (Operands e -> Operands e -> CodeGen Native (Operands e))
-> (Operands Int -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
forall e.
TypeR e
-> Operands Int
-> Operands Int
-> (Operands e -> Operands e -> CodeGen Native (Operands e))
-> (Operands Int -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
reduce1FromTo TypeR e
eR (Operands DIM1 -> Operands Int
forall sh sz. Operands (sh, sz) -> Operands sz
indexHead Operands DIM1
start) (Operands DIM1 -> Operands Int
forall sh sz. Operands (sh, sz) -> Operands sz
indexHead Operands DIM1
end) (IRFun2 Native aenv (e -> e -> e)
-> Operands e -> Operands e -> CodeGen Native (Operands e)
forall arch env aenv a b c.
IROpenFun2 arch env aenv (a -> b -> c)
-> Operands a -> Operands b -> IROpenExp arch ((env, a), b) aenv c
app2 IRFun2 Native aenv (e -> e -> e)
combine) (IROpenFun1 Native () aenv (Int -> e)
-> Operands Int -> CodeGen Native (Operands e)
forall arch env aenv a b.
IROpenFun1 arch env aenv (a -> b)
-> Operands a -> IROpenExp arch (env, a) aenv b
app1 (IRDelayed Native aenv (Vector e)
-> IROpenFun1 Native () aenv (Int -> e)
forall arch aenv sh e.
IRDelayed arch aenv (Array sh e) -> IRFun1 arch aenv (Int -> e)
delayedLinearIndex IRDelayed Native aenv (Vector e)
arrIn))
IntegralType Int
-> IRArray (Vector e)
-> Operands Int
-> Operands e
-> CodeGen Native ()
forall int sh e arch.
IntegralType int
-> IRArray (Array sh e)
-> Operands int
-> Operands e
-> CodeGen arch ()
writeArray IntegralType Int
TypeInt IRArray (Vector e)
arrTmp Operands Int
piece Operands e
r
CodeGen Native ()
forall arch. HasCallStack => CodeGen arch ()
return_
mkFoldAllP2
:: UID
-> Gamma aenv
-> TypeR e
-> IRFun2 Native aenv (e -> e -> e)
-> MIRExp Native aenv e
-> CodeGen Native (IROpenAcc Native aenv (Scalar e))
mkFoldAllP2 :: UID
-> Gamma aenv
-> TypeR e
-> IRFun2 Native aenv (e -> e -> e)
-> MIRExp Native aenv e
-> CodeGen Native (IROpenAcc Native aenv (Scalar e))
mkFoldAllP2 UID
uid Gamma aenv
aenv TypeR e
eR IRFun2 Native aenv (e -> e -> e)
combine MIRExp Native aenv e
mseed =
let
(Operands DIM1
start, Operands DIM1
end, [Parameter]
paramGang) = ShapeR DIM1 -> (Operands DIM1, Operands DIM1, [Parameter])
forall sh. ShapeR sh -> (Operands sh, Operands sh, [Parameter])
gangParam ShapeR DIM1
dim1
(IRArray (Array DIM1 e)
arrTmp, [Parameter]
paramTmp) = ArrayR (Array DIM1 e)
-> Name (Array DIM1 e) -> (IRArray (Array DIM1 e), [Parameter])
forall sh e.
ArrayR (Array sh e)
-> Name (Array sh e) -> (IRArray (Array sh e), [Parameter])
mutableArray (ShapeR DIM1 -> TypeR e -> ArrayR (Array DIM1 e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR DIM1
dim1 TypeR e
eR) Name (Array DIM1 e)
"tmp"
(IRArray (Scalar e)
arrOut, [Parameter]
paramOut) = ArrayR (Scalar e)
-> Name (Scalar e) -> (IRArray (Scalar e), [Parameter])
forall sh e.
ArrayR (Array sh e)
-> Name (Array sh e) -> (IRArray (Array sh e), [Parameter])
mutableArray (ShapeR () -> TypeR e -> ArrayR (Scalar e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR ()
dim0 TypeR e
eR) Name (Scalar e)
"out"
paramEnv :: [Parameter]
paramEnv = Gamma aenv -> [Parameter]
forall aenv. Gamma aenv -> [Parameter]
envParam Gamma aenv
aenv
zero :: Operands Int
zero = Int -> Operands Int
liftInt Int
0
in
UID
-> Label
-> [Parameter]
-> CodeGen Native ()
-> CodeGen Native (IROpenAcc Native aenv (Scalar e))
forall aenv a.
UID
-> Label
-> [Parameter]
-> CodeGen Native ()
-> CodeGen Native (IROpenAcc Native aenv a)
makeOpenAcc UID
uid Label
"foldAllP2" ([Parameter]
paramGang [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramTmp [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramOut [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramEnv) (CodeGen Native ()
-> CodeGen Native (IROpenAcc Native aenv (Scalar e)))
-> CodeGen Native ()
-> CodeGen Native (IROpenAcc Native aenv (Scalar e))
forall a b. (a -> b) -> a -> b
$ do
Operands e
r <- case MIRExp Native aenv e
mseed of
Just CodeGen Native (Operands e)
seed -> do Operands e
z <- CodeGen Native (Operands e)
seed
TypeR e
-> Operands Int
-> Operands Int
-> (Operands e -> Operands e -> CodeGen Native (Operands e))
-> Operands e
-> (Operands Int -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
forall e.
TypeR e
-> Operands Int
-> Operands Int
-> (Operands e -> Operands e -> CodeGen Native (Operands e))
-> Operands e
-> (Operands Int -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
reduceFromTo TypeR e
eR (Operands DIM1 -> Operands Int
forall sh sz. Operands (sh, sz) -> Operands sz
indexHead Operands DIM1
start) (Operands DIM1 -> Operands Int
forall sh sz. Operands (sh, sz) -> Operands sz
indexHead Operands DIM1
end) (IRFun2 Native aenv (e -> e -> e)
-> Operands e -> Operands e -> CodeGen Native (Operands e)
forall arch env aenv a b c.
IROpenFun2 arch env aenv (a -> b -> c)
-> Operands a -> Operands b -> IROpenExp arch ((env, a), b) aenv c
app2 IRFun2 Native aenv (e -> e -> e)
combine) Operands e
z (IntegralType Int
-> IRArray (Array DIM1 e)
-> Operands Int
-> CodeGen Native (Operands e)
forall int sh e arch.
IntegralType int
-> IRArray (Array sh e)
-> Operands int
-> CodeGen arch (Operands e)
readArray IntegralType Int
TypeInt IRArray (Array DIM1 e)
arrTmp)
MIRExp Native aenv e
Nothing -> TypeR e
-> Operands Int
-> Operands Int
-> (Operands e -> Operands e -> CodeGen Native (Operands e))
-> (Operands Int -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
forall e.
TypeR e
-> Operands Int
-> Operands Int
-> (Operands e -> Operands e -> CodeGen Native (Operands e))
-> (Operands Int -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
reduce1FromTo TypeR e
eR (Operands DIM1 -> Operands Int
forall sh sz. Operands (sh, sz) -> Operands sz
indexHead Operands DIM1
start) (Operands DIM1 -> Operands Int
forall sh sz. Operands (sh, sz) -> Operands sz
indexHead Operands DIM1
end) (IRFun2 Native aenv (e -> e -> e)
-> Operands e -> Operands e -> CodeGen Native (Operands e)
forall arch env aenv a b c.
IROpenFun2 arch env aenv (a -> b -> c)
-> Operands a -> Operands b -> IROpenExp arch ((env, a), b) aenv c
app2 IRFun2 Native aenv (e -> e -> e)
combine) (IntegralType Int
-> IRArray (Array DIM1 e)
-> Operands Int
-> CodeGen Native (Operands e)
forall int sh e arch.
IntegralType int
-> IRArray (Array sh e)
-> Operands int
-> CodeGen arch (Operands e)
readArray IntegralType Int
TypeInt IRArray (Array DIM1 e)
arrTmp)
IntegralType Int
-> IRArray (Scalar e)
-> Operands Int
-> Operands e
-> CodeGen Native ()
forall int sh e arch.
IntegralType int
-> IRArray (Array sh e)
-> Operands int
-> Operands e
-> CodeGen arch ()
writeArray IntegralType Int
TypeInt IRArray (Scalar e)
arrOut Operands Int
zero Operands e
r
CodeGen Native ()
forall arch. HasCallStack => CodeGen arch ()
return_
mkFoldFill
:: UID
-> Gamma aenv
-> ArrayR (Array sh e)
-> IRExp Native aenv e
-> CodeGen Native (IROpenAcc Native aenv (Array sh e))
mkFoldFill :: UID
-> Gamma aenv
-> ArrayR (Array sh e)
-> IRExp Native aenv e
-> CodeGen Native (IROpenAcc Native aenv (Array sh e))
mkFoldFill UID
uid Gamma aenv
aenv ArrayR (Array sh e)
aR IRExp Native aenv e
seed =
UID
-> Gamma aenv
-> ArrayR (Array sh e)
-> IRFun1 Native aenv (sh -> e)
-> CodeGen Native (IROpenAcc Native aenv (Array sh e))
forall aenv sh e.
UID
-> Gamma aenv
-> ArrayR (Array sh e)
-> IRFun1 Native aenv (sh -> e)
-> CodeGen Native (IROpenAcc Native aenv (Array sh e))
mkGenerate UID
uid Gamma aenv
aenv ArrayR (Array sh e)
aR ((Operands sh -> IRExp Native aenv e)
-> IRFun1 Native aenv (sh -> e)
forall a arch env aenv b.
(Operands a -> IROpenExp arch (env, a) aenv b)
-> IROpenFun1 arch env aenv (a -> b)
IRFun1 (IRExp Native aenv e -> Operands sh -> IRExp Native aenv e
forall a b. a -> b -> a
const IRExp Native aenv e
seed))
reduceFromTo
:: TypeR e
-> Operands Int
-> Operands Int
-> (Operands e -> Operands e -> CodeGen Native (Operands e))
-> Operands e
-> (Operands Int -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
reduceFromTo :: TypeR e
-> Operands Int
-> Operands Int
-> (Operands e -> Operands e -> CodeGen Native (Operands e))
-> Operands e
-> (Operands Int -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
reduceFromTo TypeR e
eR Operands Int
m Operands Int
n Operands e -> Operands e -> CodeGen Native (Operands e)
f Operands e
z Operands Int -> CodeGen Native (Operands e)
get =
TypeR e
-> Operands Int
-> Operands Int
-> Operands e
-> (Operands Int -> Operands e -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
forall a.
TypeR a
-> Operands Int
-> Operands Int
-> Operands a
-> (Operands Int -> Operands a -> CodeGen Native (Operands a))
-> CodeGen Native (Operands a)
iterFromTo TypeR e
eR Operands Int
m Operands Int
n Operands e
z ((Operands Int -> Operands e -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e))
-> (Operands Int -> Operands e -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
forall a b. (a -> b) -> a -> b
$ \Operands Int
i Operands e
acc -> do
Operands e
x <- Operands Int -> CodeGen Native (Operands e)
get Operands Int
i
Operands e
y <- Operands e -> Operands e -> CodeGen Native (Operands e)
f Operands e
acc Operands e
x
Operands e -> CodeGen Native (Operands e)
forall (m :: * -> *) a. Monad m => a -> m a
return Operands e
y
reduce1FromTo
:: TypeR e
-> Operands Int
-> Operands Int
-> (Operands e -> Operands e -> CodeGen Native (Operands e))
-> (Operands Int -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
reduce1FromTo :: TypeR e
-> Operands Int
-> Operands Int
-> (Operands e -> Operands e -> CodeGen Native (Operands e))
-> (Operands Int -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
reduce1FromTo TypeR e
eR Operands Int
m Operands Int
n Operands e -> Operands e -> CodeGen Native (Operands e)
f Operands Int -> CodeGen Native (Operands e)
get = do
Operands e
z <- Operands Int -> CodeGen Native (Operands e)
get Operands Int
m
Operands Int
m1 <- NumType Int
-> Operands Int -> Operands Int -> CodeGen Native (Operands Int)
forall a arch.
NumType a -> Operands a -> Operands a -> CodeGen arch (Operands a)
add NumType Int
forall a. IsNum a => NumType a
numType Operands Int
m (NumType Int -> Operand Int -> Operands Int
forall (dict :: * -> *) a.
(IROP dict, HasCallStack) =>
dict a -> Operand a -> Operands a
ir NumType Int
forall a. IsNum a => NumType a
numType (NumType Int -> Int -> Operand Int
forall a. NumType a -> a -> Operand a
num NumType Int
forall a. IsNum a => NumType a
numType Int
1))
TypeR e
-> Operands Int
-> Operands Int
-> (Operands e -> Operands e -> CodeGen Native (Operands e))
-> Operands e
-> (Operands Int -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
forall e.
TypeR e
-> Operands Int
-> Operands Int
-> (Operands e -> Operands e -> CodeGen Native (Operands e))
-> Operands e
-> (Operands Int -> CodeGen Native (Operands e))
-> CodeGen Native (Operands e)
reduceFromTo TypeR e
eR Operands Int
m1 Operands Int
n Operands e -> Operands e -> CodeGen Native (Operands e)
f Operands e
z Operands Int -> CodeGen Native (Operands e)
get