-- | Interpreter for the IR.
--
-- Running an IR program requires specifying the name of an entry point along
-- with values for the entry point parameters. Top level statements
-- ('progConsts') are evaluated first, and then the specified entry point is
-- invoked with the provided values.
--
-- The goal of this interpreter is not performance, but operational clarity.
-- Hence you should not expect programs to run fast at all.
module Futhark.IR.Run (runSOACS, runGPU) where

import Control.Monad (foldM, forM, when, zipWithM, zipWithM_, (>=>))
import Control.Monad.Error.Class
import Control.Monad.Except (ExceptT, runExceptT)
import Control.Monad.IO.Class
import Control.Monad.Reader (MonadReader, ReaderT, asks, local, runReaderT)
import Data.Int qualified as I
import Data.List qualified as L
import Data.List.NonEmpty qualified as NE
import Data.Map qualified as M
import Data.Maybe (mapMaybe)
import Data.Text qualified as T
import Data.Text.IO qualified as TIO
import Data.Vector.Storable qualified as SVec
import Data.Vector.Storable.Mutable qualified as MSVec
import Foreign.Storable (Storable)
import Futhark.Data qualified as V
import Futhark.IR
import Futhark.IR.GPU
  ( GPU,
    HostOp (..),
    SizeClass (..),
    SizeOp (..),
  )
import Futhark.IR.SOACS (HistOp (..), Reduce (..), SOAC (..), SOACS, Scan (..), ScremaForm (..), flatMapNonuniform)
import Futhark.IR.SegOp qualified as Seg
import Futhark.Util (showText)
import Language.Futhark.Primitive qualified as P
import Numeric.Half qualified as H
import System.IO (stderr)

data Val
  = PrimVal PrimValue
  | ArrayValue [Int] ArrayValues
  | AccValue Accumulator

-- The operator of an accumulator is looked up in 'interpAccOps' by
-- its certificate.
data Accumulator = Accumulator
  { Accumulator -> VName
accCert :: VName,
    Accumulator -> [Int]
accShape :: [Int],
    Accumulator -> [Val]
accArrays :: [Val]
  }

data ArrayValues
  = I8ArrayValues (MSVec.IOVector I.Int8)
  | I16ArrayValues (MSVec.IOVector I.Int16)
  | I32ArrayValues (MSVec.IOVector I.Int32)
  | I64ArrayValues (MSVec.IOVector I.Int64)
  | F16ArrayValues (MSVec.IOVector H.Half)
  | F32ArrayValues (MSVec.IOVector Float)
  | F64ArrayValues (MSVec.IOVector Double)
  | BoolArrayValues (MSVec.IOVector Bool)
  | UnitArrayValues (MSVec.IOVector ())

arrayValuesType :: ArrayValues -> PrimType
arrayValuesType :: ArrayValues -> PrimType
arrayValuesType I8ArrayValues {} = IntType -> PrimType
IntType IntType
Int8
arrayValuesType I16ArrayValues {} = IntType -> PrimType
IntType IntType
Int16
arrayValuesType I32ArrayValues {} = IntType -> PrimType
IntType IntType
Int32
arrayValuesType I64ArrayValues {} = IntType -> PrimType
IntType IntType
Int64
arrayValuesType F16ArrayValues {} = FloatType -> PrimType
FloatType FloatType
Float16
arrayValuesType F32ArrayValues {} = FloatType -> PrimType
FloatType FloatType
Float32
arrayValuesType F64ArrayValues {} = FloatType -> PrimType
FloatType FloatType
Float64
arrayValuesType BoolArrayValues {} = PrimType
Bool
arrayValuesType UnitArrayValues {} = PrimType
Unit

data KernelResultValue
  = KernelValue Val
  | KernelTile [(Int, Int)] Val
  | KernelRegTile [(Int, Int, Int)] Val

type Env = M.Map VName Val

type FunEnv rep = M.Map Name (FunDef rep)

type OpEvaluator rep =
  FunEnv rep -> Env -> Op rep -> InterpM rep [Val]

data InterpEnv rep = InterpEnv
  { forall rep. InterpEnv rep -> OpEvaluator rep
interpOpEvaluator :: OpEvaluator rep,
    -- | Accumulator operators, keyed by certificate, together with the
    -- environments of the 'WithAcc' that defined them. The operator
    -- receives (indices <> old values <> new values).
    forall rep.
InterpEnv rep -> Map VName (FunEnv rep, Env, Lambda rep)
interpAccOps :: M.Map VName (FunEnv rep, Env, Lambda rep)
  }

newtype InterpM rep a = InterpM
  { forall rep a.
InterpM rep a -> ReaderT (InterpEnv rep) (ExceptT Text IO) a
unInterpM :: ReaderT (InterpEnv rep) (ExceptT T.Text IO) a
  }
  deriving
    ((forall a b. (a -> b) -> InterpM rep a -> InterpM rep b)
-> (forall a b. a -> InterpM rep b -> InterpM rep a)
-> Functor (InterpM rep)
forall a b. a -> InterpM rep b -> InterpM rep a
forall a b. (a -> b) -> InterpM rep a -> InterpM rep b
forall rep a b. a -> InterpM rep b -> InterpM rep a
forall rep a b. (a -> b) -> InterpM rep a -> InterpM rep b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall rep a b. (a -> b) -> InterpM rep a -> InterpM rep b
fmap :: forall a b. (a -> b) -> InterpM rep a -> InterpM rep b
$c<$ :: forall rep a b. a -> InterpM rep b -> InterpM rep a
<$ :: forall a b. a -> InterpM rep b -> InterpM rep a
Functor, Functor (InterpM rep)
Functor (InterpM rep) =>
(forall a. a -> InterpM rep a)
-> (forall a b.
    InterpM rep (a -> b) -> InterpM rep a -> InterpM rep b)
-> (forall a b c.
    (a -> b -> c) -> InterpM rep a -> InterpM rep b -> InterpM rep c)
-> (forall a b. InterpM rep a -> InterpM rep b -> InterpM rep b)
-> (forall a b. InterpM rep a -> InterpM rep b -> InterpM rep a)
-> Applicative (InterpM rep)
forall rep. Functor (InterpM rep)
forall a. a -> InterpM rep a
forall rep a. a -> InterpM rep a
forall a b. InterpM rep a -> InterpM rep b -> InterpM rep a
forall a b. InterpM rep a -> InterpM rep b -> InterpM rep b
forall a b. InterpM rep (a -> b) -> InterpM rep a -> InterpM rep b
forall rep a b. InterpM rep a -> InterpM rep b -> InterpM rep a
forall rep a b. InterpM rep a -> InterpM rep b -> InterpM rep b
forall rep a b.
InterpM rep (a -> b) -> InterpM rep a -> InterpM rep b
forall a b c.
(a -> b -> c) -> InterpM rep a -> InterpM rep b -> InterpM rep c
forall rep a b c.
(a -> b -> c) -> InterpM rep a -> InterpM rep b -> InterpM rep c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall rep a. a -> InterpM rep a
pure :: forall a. a -> InterpM rep a
$c<*> :: forall rep a b.
InterpM rep (a -> b) -> InterpM rep a -> InterpM rep b
<*> :: forall a b. InterpM rep (a -> b) -> InterpM rep a -> InterpM rep b
$cliftA2 :: forall rep a b c.
(a -> b -> c) -> InterpM rep a -> InterpM rep b -> InterpM rep c
liftA2 :: forall a b c.
(a -> b -> c) -> InterpM rep a -> InterpM rep b -> InterpM rep c
$c*> :: forall rep a b. InterpM rep a -> InterpM rep b -> InterpM rep b
*> :: forall a b. InterpM rep a -> InterpM rep b -> InterpM rep b
$c<* :: forall rep a b. InterpM rep a -> InterpM rep b -> InterpM rep a
<* :: forall a b. InterpM rep a -> InterpM rep b -> InterpM rep a
Applicative, Applicative (InterpM rep)
Applicative (InterpM rep) =>
(forall a b.
 InterpM rep a -> (a -> InterpM rep b) -> InterpM rep b)
-> (forall a b. InterpM rep a -> InterpM rep b -> InterpM rep b)
-> (forall a. a -> InterpM rep a)
-> Monad (InterpM rep)
forall rep. Applicative (InterpM rep)
forall a. a -> InterpM rep a
forall rep a. a -> InterpM rep a
forall a b. InterpM rep a -> InterpM rep b -> InterpM rep b
forall a b. InterpM rep a -> (a -> InterpM rep b) -> InterpM rep b
forall rep a b. InterpM rep a -> InterpM rep b -> InterpM rep b
forall rep a b.
InterpM rep a -> (a -> InterpM rep b) -> InterpM rep b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall rep a b.
InterpM rep a -> (a -> InterpM rep b) -> InterpM rep b
>>= :: forall a b. InterpM rep a -> (a -> InterpM rep b) -> InterpM rep b
$c>> :: forall rep a b. InterpM rep a -> InterpM rep b -> InterpM rep b
>> :: forall a b. InterpM rep a -> InterpM rep b -> InterpM rep b
$creturn :: forall rep a. a -> InterpM rep a
return :: forall a. a -> InterpM rep a
Monad, MonadReader (InterpEnv rep), MonadError T.Text, Monad (InterpM rep)
Monad (InterpM rep) =>
(forall a. IO a -> InterpM rep a) -> MonadIO (InterpM rep)
forall rep. Monad (InterpM rep)
forall a. IO a -> InterpM rep a
forall rep a. IO a -> InterpM rep a
forall (m :: * -> *).
Monad m =>
(forall a. IO a -> m a) -> MonadIO m
$cliftIO :: forall rep a. IO a -> InterpM rep a
liftIO :: forall a. IO a -> InterpM rep a
MonadIO)

interpError :: T.Text -> InterpM rep a
interpError :: forall rep a. Text -> InterpM rep a
interpError = Text -> InterpM rep a
forall a. Text -> InterpM rep a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError

newArrayValue :: [Int] -> PrimType -> [PrimValue] -> InterpM rep Val
newArrayValue :: forall rep. [Int] -> PrimType -> [PrimValue] -> InterpM rep Val
newArrayValue [Int]
shape PrimType
element_type [PrimValue]
values =
  [Int] -> ArrayValues -> Val
ArrayValue [Int]
shape (ArrayValues -> Val) -> InterpM rep ArrayValues -> InterpM rep Val
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PrimType -> InterpM rep ArrayValues
forall {f :: * -> *}. MonadIO f => PrimType -> f ArrayValues
newValues PrimType
element_type
  where
    newValues :: PrimType -> f ArrayValues
newValues (IntType IntType
Int8) = IOVector Int8 -> ArrayValues
I8ArrayValues (IOVector Int8 -> ArrayValues)
-> f (IOVector Int8) -> f ArrayValues
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PrimValue -> Int8) -> f (IOVector Int8)
forall {m :: * -> *} {b}.
(MonadIO m, Storable b) =>
(PrimValue -> b) -> m (MVector RealWorld b)
newPrimVector PrimValue -> Int8
expectInt8
    newValues (IntType IntType
Int16) = IOVector Int16 -> ArrayValues
I16ArrayValues (IOVector Int16 -> ArrayValues)
-> f (IOVector Int16) -> f ArrayValues
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PrimValue -> Int16) -> f (IOVector Int16)
forall {m :: * -> *} {b}.
(MonadIO m, Storable b) =>
(PrimValue -> b) -> m (MVector RealWorld b)
newPrimVector PrimValue -> Int16
expectInt16
    newValues (IntType IntType
Int32) = IOVector Int32 -> ArrayValues
I32ArrayValues (IOVector Int32 -> ArrayValues)
-> f (IOVector Int32) -> f ArrayValues
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PrimValue -> Int32) -> f (IOVector Int32)
forall {m :: * -> *} {b}.
(MonadIO m, Storable b) =>
(PrimValue -> b) -> m (MVector RealWorld b)
newPrimVector PrimValue -> Int32
expectInt32
    newValues (IntType IntType
Int64) = IOVector Int64 -> ArrayValues
I64ArrayValues (IOVector Int64 -> ArrayValues)
-> f (IOVector Int64) -> f ArrayValues
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PrimValue -> Int64) -> f (IOVector Int64)
forall {m :: * -> *} {b}.
(MonadIO m, Storable b) =>
(PrimValue -> b) -> m (MVector RealWorld b)
newPrimVector PrimValue -> Int64
expectInt64
    newValues (FloatType FloatType
Float16) = IOVector Half -> ArrayValues
F16ArrayValues (IOVector Half -> ArrayValues)
-> f (IOVector Half) -> f ArrayValues
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PrimValue -> Half) -> f (IOVector Half)
forall {m :: * -> *} {b}.
(MonadIO m, Storable b) =>
(PrimValue -> b) -> m (MVector RealWorld b)
newPrimVector PrimValue -> Half
expectFloat16
    newValues (FloatType FloatType
Float32) = IOVector Float -> ArrayValues
F32ArrayValues (IOVector Float -> ArrayValues)
-> f (IOVector Float) -> f ArrayValues
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PrimValue -> Float) -> f (IOVector Float)
forall {m :: * -> *} {b}.
(MonadIO m, Storable b) =>
(PrimValue -> b) -> m (MVector RealWorld b)
newPrimVector PrimValue -> Float
expectFloat32
    newValues (FloatType FloatType
Float64) = IOVector Double -> ArrayValues
F64ArrayValues (IOVector Double -> ArrayValues)
-> f (IOVector Double) -> f ArrayValues
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PrimValue -> Double) -> f (IOVector Double)
forall {m :: * -> *} {b}.
(MonadIO m, Storable b) =>
(PrimValue -> b) -> m (MVector RealWorld b)
newPrimVector PrimValue -> Double
expectFloat64
    newValues PrimType
Bool = IOVector Bool -> ArrayValues
BoolArrayValues (IOVector Bool -> ArrayValues)
-> f (IOVector Bool) -> f ArrayValues
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PrimValue -> Bool) -> f (IOVector Bool)
forall {m :: * -> *} {b}.
(MonadIO m, Storable b) =>
(PrimValue -> b) -> m (MVector RealWorld b)
newPrimVector PrimValue -> Bool
expectBool
    newValues PrimType
Unit = IOVector () -> ArrayValues
UnitArrayValues (IOVector () -> ArrayValues) -> f (IOVector ()) -> f ArrayValues
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PrimValue -> ()) -> f (IOVector ())
forall {m :: * -> *} {b}.
(MonadIO m, Storable b) =>
(PrimValue -> b) -> m (MVector RealWorld b)
newPrimVector PrimValue -> ()
expectUnit

    newPrimVector :: (PrimValue -> b) -> m (MVector RealWorld b)
newPrimVector PrimValue -> b
unwrap = IO (MVector RealWorld b) -> m (MVector RealWorld b)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MVector RealWorld b) -> m (MVector RealWorld b))
-> IO (MVector RealWorld b) -> m (MVector RealWorld b)
forall a b. (a -> b) -> a -> b
$ do
      vector <- Int -> IO (MVector (PrimState IO) b)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
Int -> m (MVector (PrimState m) a)
MSVec.new ([PrimValue] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [PrimValue]
values)
      zipWithM_ (MSVec.write vector) [0 ..] (map unwrap values)
      pure vector

    expectInt8 :: PrimValue -> Int8
expectInt8 (IntValue (Int8Value Int8
element)) = Int8
element
    expectInt8 PrimValue
_ = [Char] -> Int8
forall a. HasCallStack => [Char] -> a
error [Char]
"expected an i8 value"
    expectInt16 :: PrimValue -> Int16
expectInt16 (IntValue (Int16Value Int16
element)) = Int16
element
    expectInt16 PrimValue
_ = [Char] -> Int16
forall a. HasCallStack => [Char] -> a
error [Char]
"expected an i16 value"
    expectInt32 :: PrimValue -> Int32
expectInt32 (IntValue (Int32Value Int32
element)) = Int32
element
    expectInt32 PrimValue
_ = [Char] -> Int32
forall a. HasCallStack => [Char] -> a
error [Char]
"expected an i32 value"
    expectInt64 :: PrimValue -> Int64
expectInt64 (IntValue (Int64Value Int64
element)) = Int64
element
    expectInt64 PrimValue
_ = [Char] -> Int64
forall a. HasCallStack => [Char] -> a
error [Char]
"expected an i64 value"
    expectFloat16 :: PrimValue -> Half
expectFloat16 (FloatValue (Float16Value Half
element)) = Half
element
    expectFloat16 PrimValue
_ = [Char] -> Half
forall a. HasCallStack => [Char] -> a
error [Char]
"expected an f16 value"
    expectFloat32 :: PrimValue -> Float
expectFloat32 (FloatValue (Float32Value Float
element)) = Float
element
    expectFloat32 PrimValue
_ = [Char] -> Float
forall a. HasCallStack => [Char] -> a
error [Char]
"expected an f32 value"
    expectFloat64 :: PrimValue -> Double
expectFloat64 (FloatValue (Float64Value Double
element)) = Double
element
    expectFloat64 PrimValue
_ = [Char] -> Double
forall a. HasCallStack => [Char] -> a
error [Char]
"expected an f64 value"
    expectBool :: PrimValue -> Bool
expectBool (BoolValue Bool
element) = Bool
element
    expectBool PrimValue
_ = [Char] -> Bool
forall a. HasCallStack => [Char] -> a
error [Char]
"expected a bool value"
    expectUnit :: PrimValue -> ()
expectUnit PrimValue
UnitValue = ()
    expectUnit PrimValue
_ = [Char] -> ()
forall a. HasCallStack => [Char] -> a
error [Char]
"expected a unit value"

arrayValues :: ArrayValues -> InterpM rep [PrimValue]
arrayValues :: forall rep. ArrayValues -> InterpM rep [PrimValue]
arrayValues ArrayValues
values =
  (Int -> InterpM rep PrimValue) -> [Int] -> InterpM rep [PrimValue]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (ArrayValues -> Int -> InterpM rep PrimValue
forall rep. ArrayValues -> Int -> InterpM rep PrimValue
readArrayValue ArrayValues
values) [Int
0 .. ArrayValues -> Int
arrayValuesLength ArrayValues
values Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]

arrayValuesLength :: ArrayValues -> Int
arrayValuesLength :: ArrayValues -> Int
arrayValuesLength (I8ArrayValues IOVector Int8
values) = IOVector Int8 -> Int
forall a s. Storable a => MVector s a -> Int
MSVec.length IOVector Int8
values
arrayValuesLength (I16ArrayValues IOVector Int16
values) = IOVector Int16 -> Int
forall a s. Storable a => MVector s a -> Int
MSVec.length IOVector Int16
values
arrayValuesLength (I32ArrayValues IOVector Int32
values) = IOVector Int32 -> Int
forall a s. Storable a => MVector s a -> Int
MSVec.length IOVector Int32
values
arrayValuesLength (I64ArrayValues IOVector Int64
values) = IOVector Int64 -> Int
forall a s. Storable a => MVector s a -> Int
MSVec.length IOVector Int64
values
arrayValuesLength (F16ArrayValues IOVector Half
values) = IOVector Half -> Int
forall a s. Storable a => MVector s a -> Int
MSVec.length IOVector Half
values
arrayValuesLength (F32ArrayValues IOVector Float
values) = IOVector Float -> Int
forall a s. Storable a => MVector s a -> Int
MSVec.length IOVector Float
values
arrayValuesLength (F64ArrayValues IOVector Double
values) = IOVector Double -> Int
forall a s. Storable a => MVector s a -> Int
MSVec.length IOVector Double
values
arrayValuesLength (BoolArrayValues IOVector Bool
values) = IOVector Bool -> Int
forall a s. Storable a => MVector s a -> Int
MSVec.length IOVector Bool
values
arrayValuesLength (UnitArrayValues IOVector ()
values) = IOVector () -> Int
forall a s. Storable a => MVector s a -> Int
MSVec.length IOVector ()
values

readArrayValue :: ArrayValues -> Int -> InterpM rep PrimValue
readArrayValue :: forall rep. ArrayValues -> Int -> InterpM rep PrimValue
readArrayValue ArrayValues
vector Int
index
  | Int
index Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
index Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= ArrayValues -> Int
arrayValuesLength ArrayValues
vector =
      Text -> InterpM rep PrimValue
forall rep a. Text -> InterpM rep a
interpError Text
"array index out of bounds"
readArrayValue (I8ArrayValues IOVector Int8
values) Int
index =
  IntValue -> PrimValue
IntValue (IntValue -> PrimValue) -> (Int8 -> IntValue) -> Int8 -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int8 -> IntValue
Int8Value (Int8 -> PrimValue) -> InterpM rep Int8 -> InterpM rep PrimValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Int8 -> InterpM rep Int8
forall a. IO a -> InterpM rep a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (MVector (PrimState IO) Int8 -> Int -> IO Int8
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> m a
MSVec.read IOVector Int8
MVector (PrimState IO) Int8
values Int
index)
readArrayValue (I16ArrayValues IOVector Int16
values) Int
index =
  IntValue -> PrimValue
IntValue (IntValue -> PrimValue)
-> (Int16 -> IntValue) -> Int16 -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int16 -> IntValue
Int16Value (Int16 -> PrimValue) -> InterpM rep Int16 -> InterpM rep PrimValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Int16 -> InterpM rep Int16
forall a. IO a -> InterpM rep a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (MVector (PrimState IO) Int16 -> Int -> IO Int16
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> m a
MSVec.read IOVector Int16
MVector (PrimState IO) Int16
values Int
index)
readArrayValue (I32ArrayValues IOVector Int32
values) Int
index =
  IntValue -> PrimValue
IntValue (IntValue -> PrimValue)
-> (Int32 -> IntValue) -> Int32 -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int32 -> IntValue
Int32Value (Int32 -> PrimValue) -> InterpM rep Int32 -> InterpM rep PrimValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Int32 -> InterpM rep Int32
forall a. IO a -> InterpM rep a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (MVector (PrimState IO) Int32 -> Int -> IO Int32
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> m a
MSVec.read IOVector Int32
MVector (PrimState IO) Int32
values Int
index)
readArrayValue (I64ArrayValues IOVector Int64
values) Int
index =
  IntValue -> PrimValue
IntValue (IntValue -> PrimValue)
-> (Int64 -> IntValue) -> Int64 -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> IntValue
Int64Value (Int64 -> PrimValue) -> InterpM rep Int64 -> InterpM rep PrimValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Int64 -> InterpM rep Int64
forall a. IO a -> InterpM rep a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (MVector (PrimState IO) Int64 -> Int -> IO Int64
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> m a
MSVec.read IOVector Int64
MVector (PrimState IO) Int64
values Int
index)
readArrayValue (F16ArrayValues IOVector Half
values) Int
index =
  FloatValue -> PrimValue
FloatValue (FloatValue -> PrimValue)
-> (Half -> FloatValue) -> Half -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Half -> FloatValue
Float16Value (Half -> PrimValue) -> InterpM rep Half -> InterpM rep PrimValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Half -> InterpM rep Half
forall a. IO a -> InterpM rep a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (MVector (PrimState IO) Half -> Int -> IO Half
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> m a
MSVec.read IOVector Half
MVector (PrimState IO) Half
values Int
index)
readArrayValue (F32ArrayValues IOVector Float
values) Int
index =
  FloatValue -> PrimValue
FloatValue (FloatValue -> PrimValue)
-> (Float -> FloatValue) -> Float -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Float -> FloatValue
Float32Value (Float -> PrimValue) -> InterpM rep Float -> InterpM rep PrimValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Float -> InterpM rep Float
forall a. IO a -> InterpM rep a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (MVector (PrimState IO) Float -> Int -> IO Float
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> m a
MSVec.read IOVector Float
MVector (PrimState IO) Float
values Int
index)
readArrayValue (F64ArrayValues IOVector Double
values) Int
index =
  FloatValue -> PrimValue
FloatValue (FloatValue -> PrimValue)
-> (Double -> FloatValue) -> Double -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> FloatValue
Float64Value (Double -> PrimValue)
-> InterpM rep Double -> InterpM rep PrimValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Double -> InterpM rep Double
forall a. IO a -> InterpM rep a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (MVector (PrimState IO) Double -> Int -> IO Double
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> m a
MSVec.read IOVector Double
MVector (PrimState IO) Double
values Int
index)
readArrayValue (BoolArrayValues IOVector Bool
values) Int
index =
  Bool -> PrimValue
BoolValue (Bool -> PrimValue) -> InterpM rep Bool -> InterpM rep PrimValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Bool -> InterpM rep Bool
forall a. IO a -> InterpM rep a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (MVector (PrimState IO) Bool -> Int -> IO Bool
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> m a
MSVec.read IOVector Bool
MVector (PrimState IO) Bool
values Int
index)
readArrayValue (UnitArrayValues IOVector ()
values) Int
index =
  PrimValue
UnitValue PrimValue -> InterpM rep () -> InterpM rep PrimValue
forall a b. a -> InterpM rep b -> InterpM rep a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ IO () -> InterpM rep ()
forall a. IO a -> InterpM rep a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (MVector (PrimState IO) () -> Int -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> m a
MSVec.read IOVector ()
MVector (PrimState IO) ()
values Int
index)

writeArrayValue :: ArrayValues -> Int -> PrimValue -> InterpM rep ()
writeArrayValue :: forall rep. ArrayValues -> Int -> PrimValue -> InterpM rep ()
writeArrayValue (I8ArrayValues IOVector Int8
values) Int
index (IntValue (Int8Value Int8
element)) =
  IO () -> InterpM rep ()
forall a. IO a -> InterpM rep a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> InterpM rep ()) -> IO () -> InterpM rep ()
forall a b. (a -> b) -> a -> b
$ MVector (PrimState IO) Int8 -> Int -> Int8 -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MSVec.write IOVector Int8
MVector (PrimState IO) Int8
values Int
index Int8
element
writeArrayValue (I16ArrayValues IOVector Int16
values) Int
index (IntValue (Int16Value Int16
element)) =
  IO () -> InterpM rep ()
forall a. IO a -> InterpM rep a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> InterpM rep ()) -> IO () -> InterpM rep ()
forall a b. (a -> b) -> a -> b
$ MVector (PrimState IO) Int16 -> Int -> Int16 -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MSVec.write IOVector Int16
MVector (PrimState IO) Int16
values Int
index Int16
element
writeArrayValue (I32ArrayValues IOVector Int32
values) Int
index (IntValue (Int32Value Int32
element)) =
  IO () -> InterpM rep ()
forall a. IO a -> InterpM rep a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> InterpM rep ()) -> IO () -> InterpM rep ()
forall a b. (a -> b) -> a -> b
$ MVector (PrimState IO) Int32 -> Int -> Int32 -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MSVec.write IOVector Int32
MVector (PrimState IO) Int32
values Int
index Int32
element
writeArrayValue (I64ArrayValues IOVector Int64
values) Int
index (IntValue (Int64Value Int64
element)) =
  IO () -> InterpM rep ()
forall a. IO a -> InterpM rep a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> InterpM rep ()) -> IO () -> InterpM rep ()
forall a b. (a -> b) -> a -> b
$ MVector (PrimState IO) Int64 -> Int -> Int64 -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MSVec.write IOVector Int64
MVector (PrimState IO) Int64
values Int
index Int64
element
writeArrayValue (F16ArrayValues IOVector Half
values) Int
index (FloatValue (Float16Value Half
element)) =
  IO () -> InterpM rep ()
forall a. IO a -> InterpM rep a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> InterpM rep ()) -> IO () -> InterpM rep ()
forall a b. (a -> b) -> a -> b
$ MVector (PrimState IO) Half -> Int -> Half -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MSVec.write IOVector Half
MVector (PrimState IO) Half
values Int
index Half
element
writeArrayValue (F32ArrayValues IOVector Float
values) Int
index (FloatValue (Float32Value Float
element)) =
  IO () -> InterpM rep ()
forall a. IO a -> InterpM rep a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> InterpM rep ()) -> IO () -> InterpM rep ()
forall a b. (a -> b) -> a -> b
$ MVector (PrimState IO) Float -> Int -> Float -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MSVec.write IOVector Float
MVector (PrimState IO) Float
values Int
index Float
element
writeArrayValue (F64ArrayValues IOVector Double
values) Int
index (FloatValue (Float64Value Double
element)) =
  IO () -> InterpM rep ()
forall a. IO a -> InterpM rep a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> InterpM rep ()) -> IO () -> InterpM rep ()
forall a b. (a -> b) -> a -> b
$ MVector (PrimState IO) Double -> Int -> Double -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MSVec.write IOVector Double
MVector (PrimState IO) Double
values Int
index Double
element
writeArrayValue (BoolArrayValues IOVector Bool
values) Int
index (BoolValue Bool
element) =
  IO () -> InterpM rep ()
forall a. IO a -> InterpM rep a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> InterpM rep ()) -> IO () -> InterpM rep ()
forall a b. (a -> b) -> a -> b
$ MVector (PrimState IO) Bool -> Int -> Bool -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MSVec.write IOVector Bool
MVector (PrimState IO) Bool
values Int
index Bool
element
writeArrayValue (UnitArrayValues IOVector ()
values) Int
index PrimValue
UnitValue =
  IO () -> InterpM rep ()
forall a. IO a -> InterpM rep a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> InterpM rep ()) -> IO () -> InterpM rep ()
forall a b. (a -> b) -> a -> b
$ MVector (PrimState IO) () -> Int -> () -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MSVec.write IOVector ()
MVector (PrimState IO) ()
values Int
index ()
writeArrayValue ArrayValues
_ Int
_ PrimValue
_ =
  [Char] -> InterpM rep ()
forall a. HasCallStack => [Char] -> a
error [Char]
"array element type mismatch"

cloneArrayValues :: ArrayValues -> IO ArrayValues
cloneArrayValues :: ArrayValues -> IO ArrayValues
cloneArrayValues (I8ArrayValues IOVector Int8
values) = IOVector Int8 -> ArrayValues
I8ArrayValues (IOVector Int8 -> ArrayValues)
-> IO (IOVector Int8) -> IO ArrayValues
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) Int8 -> IO (MVector (PrimState IO) Int8)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> m (MVector (PrimState m) a)
MSVec.clone IOVector Int8
MVector (PrimState IO) Int8
values
cloneArrayValues (I16ArrayValues IOVector Int16
values) = IOVector Int16 -> ArrayValues
I16ArrayValues (IOVector Int16 -> ArrayValues)
-> IO (IOVector Int16) -> IO ArrayValues
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) Int16 -> IO (MVector (PrimState IO) Int16)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> m (MVector (PrimState m) a)
MSVec.clone IOVector Int16
MVector (PrimState IO) Int16
values
cloneArrayValues (I32ArrayValues IOVector Int32
values) = IOVector Int32 -> ArrayValues
I32ArrayValues (IOVector Int32 -> ArrayValues)
-> IO (IOVector Int32) -> IO ArrayValues
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) Int32 -> IO (MVector (PrimState IO) Int32)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> m (MVector (PrimState m) a)
MSVec.clone IOVector Int32
MVector (PrimState IO) Int32
values
cloneArrayValues (I64ArrayValues IOVector Int64
values) = IOVector Int64 -> ArrayValues
I64ArrayValues (IOVector Int64 -> ArrayValues)
-> IO (IOVector Int64) -> IO ArrayValues
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) Int64 -> IO (MVector (PrimState IO) Int64)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> m (MVector (PrimState m) a)
MSVec.clone IOVector Int64
MVector (PrimState IO) Int64
values
cloneArrayValues (F16ArrayValues IOVector Half
values) = IOVector Half -> ArrayValues
F16ArrayValues (IOVector Half -> ArrayValues)
-> IO (IOVector Half) -> IO ArrayValues
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) Half -> IO (MVector (PrimState IO) Half)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> m (MVector (PrimState m) a)
MSVec.clone IOVector Half
MVector (PrimState IO) Half
values
cloneArrayValues (F32ArrayValues IOVector Float
values) = IOVector Float -> ArrayValues
F32ArrayValues (IOVector Float -> ArrayValues)
-> IO (IOVector Float) -> IO ArrayValues
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) Float -> IO (MVector (PrimState IO) Float)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> m (MVector (PrimState m) a)
MSVec.clone IOVector Float
MVector (PrimState IO) Float
values
cloneArrayValues (F64ArrayValues IOVector Double
values) = IOVector Double -> ArrayValues
F64ArrayValues (IOVector Double -> ArrayValues)
-> IO (IOVector Double) -> IO ArrayValues
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) Double -> IO (MVector (PrimState IO) Double)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> m (MVector (PrimState m) a)
MSVec.clone IOVector Double
MVector (PrimState IO) Double
values
cloneArrayValues (BoolArrayValues IOVector Bool
values) = IOVector Bool -> ArrayValues
BoolArrayValues (IOVector Bool -> ArrayValues)
-> IO (IOVector Bool) -> IO ArrayValues
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) Bool -> IO (MVector (PrimState IO) Bool)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> m (MVector (PrimState m) a)
MSVec.clone IOVector Bool
MVector (PrimState IO) Bool
values
cloneArrayValues (UnitArrayValues IOVector ()
values) = IOVector () -> ArrayValues
UnitArrayValues (IOVector () -> ArrayValues) -> IO (IOVector ()) -> IO ArrayValues
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) () -> IO (MVector (PrimState IO) ())
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> m (MVector (PrimState m) a)
MSVec.clone IOVector ()
MVector (PrimState IO) ()
values

evalStms :: FunEnv rep -> Env -> Stms rep -> InterpM rep Env
evalStms :: forall rep. FunEnv rep -> Env -> Stms rep -> InterpM rep Env
evalStms FunEnv rep
funs Env
env Stms rep
stms =
  (Env -> Stm rep -> InterpM rep Env)
-> Env -> [Stm rep] -> InterpM rep Env
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (FunEnv rep -> Env -> Stm rep -> InterpM rep Env
forall rep. FunEnv rep -> Env -> Stm rep -> InterpM rep Env
evalStm FunEnv rep
funs) Env
env (Stms rep -> [Stm rep]
forall rep. Stms rep -> [Stm rep]
stmsToList Stms rep
stms)

evalBody :: FunEnv rep -> Env -> Body rep -> InterpM rep [Val]
evalBody :: forall rep. FunEnv rep -> Env -> Body rep -> InterpM rep [Val]
evalBody FunEnv rep
funs Env
env (Body BodyDec rep
_ Stms rep
stms [SubExpRes]
results) = do
  env' <- FunEnv rep -> Env -> Stms rep -> InterpM rep Env
forall rep. FunEnv rep -> Env -> Stms rep -> InterpM rep Env
evalStms FunEnv rep
funs Env
env Stms rep
stms
  mapM (evalSubExp env' . resSubExp) results

evalKernelBody ::
  FunEnv rep -> Env -> Seg.KernelBody rep -> InterpM rep [KernelResultValue]
evalKernelBody :: forall rep.
FunEnv rep
-> Env -> KernelBody rep -> InterpM rep [KernelResultValue]
evalKernelBody FunEnv rep
funs Env
env (Body BodyDec rep
_ Stms rep
stms [KernelResult]
results) = do
  env' <- FunEnv rep -> Env -> Stms rep -> InterpM rep Env
forall rep. FunEnv rep -> Env -> Stms rep -> InterpM rep Env
evalStms FunEnv rep
funs Env
env Stms rep
stms
  mapM (evalKernelResult env') results
  where
    evalKernelResult :: Env -> KernelResult -> InterpM rep KernelResultValue
evalKernelResult Env
env' (Seg.Returns ResultManifest
_ Certs
_ SubExp
result) =
      Val -> KernelResultValue
KernelValue (Val -> KernelResultValue)
-> InterpM rep Val -> InterpM rep KernelResultValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> SubExp -> InterpM rep Val
forall rep. Env -> SubExp -> InterpM rep Val
evalSubExp Env
env' SubExp
result
    evalKernelResult Env
env' (Seg.TileReturns Certs
_ [(SubExp, SubExp)]
dimensions VName
tile) =
      [(Int, Int)] -> Val -> KernelResultValue
KernelTile ([(Int, Int)] -> Val -> KernelResultValue)
-> InterpM rep [(Int, Int)]
-> InterpM rep (Val -> KernelResultValue)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((SubExp, SubExp) -> InterpM rep (Int, Int))
-> [(SubExp, SubExp)] -> InterpM rep [(Int, Int)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Env -> (SubExp, SubExp) -> InterpM rep (Int, Int)
forall {rep}. Env -> (SubExp, SubExp) -> InterpM rep (Int, Int)
evalPair Env
env') [(SubExp, SubExp)]
dimensions InterpM rep (Val -> KernelResultValue)
-> InterpM rep Val -> InterpM rep KernelResultValue
forall a b. InterpM rep (a -> b) -> InterpM rep a -> InterpM rep b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Env -> VName -> InterpM rep Val
forall rep. Env -> VName -> InterpM rep Val
evalVar Env
env' VName
tile
    evalKernelResult Env
env' (Seg.RegTileReturns Certs
_ [(SubExp, SubExp, SubExp)]
dimensions VName
tile) =
      [(Int, Int, Int)] -> Val -> KernelResultValue
KernelRegTile ([(Int, Int, Int)] -> Val -> KernelResultValue)
-> InterpM rep [(Int, Int, Int)]
-> InterpM rep (Val -> KernelResultValue)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((SubExp, SubExp, SubExp) -> InterpM rep (Int, Int, Int))
-> [(SubExp, SubExp, SubExp)] -> InterpM rep [(Int, Int, Int)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Env -> (SubExp, SubExp, SubExp) -> InterpM rep (Int, Int, Int)
forall {rep}.
Env -> (SubExp, SubExp, SubExp) -> InterpM rep (Int, Int, Int)
evalTriple Env
env') [(SubExp, SubExp, SubExp)]
dimensions InterpM rep (Val -> KernelResultValue)
-> InterpM rep Val -> InterpM rep KernelResultValue
forall a b. InterpM rep (a -> b) -> InterpM rep a -> InterpM rep b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Env -> VName -> InterpM rep Val
forall rep. Env -> VName -> InterpM rep Val
evalVar Env
env' VName
tile

    evalPair :: Env -> (SubExp, SubExp) -> InterpM rep (Int, Int)
evalPair Env
env' (SubExp
size, SubExp
tile_size) =
      (,) (Int -> Int -> (Int, Int))
-> InterpM rep Int -> InterpM rep (Int -> (Int, Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> SubExp -> InterpM rep Int
forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env' SubExp
size InterpM rep (Int -> (Int, Int))
-> InterpM rep Int -> InterpM rep (Int, Int)
forall a b. InterpM rep (a -> b) -> InterpM rep a -> InterpM rep b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Env -> SubExp -> InterpM rep Int
forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env' SubExp
tile_size
    evalTriple :: Env -> (SubExp, SubExp, SubExp) -> InterpM rep (Int, Int, Int)
evalTriple Env
env' (SubExp
size, SubExp
block_tile, SubExp
reg_tile) =
      (,,)
        (Int -> Int -> Int -> (Int, Int, Int))
-> InterpM rep Int -> InterpM rep (Int -> Int -> (Int, Int, Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> SubExp -> InterpM rep Int
forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env' SubExp
size
        InterpM rep (Int -> Int -> (Int, Int, Int))
-> InterpM rep Int -> InterpM rep (Int -> (Int, Int, Int))
forall a b. InterpM rep (a -> b) -> InterpM rep a -> InterpM rep b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Env -> SubExp -> InterpM rep Int
forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env' SubExp
block_tile
        InterpM rep (Int -> (Int, Int, Int))
-> InterpM rep Int -> InterpM rep (Int, Int, Int)
forall a b. InterpM rep (a -> b) -> InterpM rep a -> InterpM rep b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Env -> SubExp -> InterpM rep Int
forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env' SubExp
reg_tile

-- Evaluate the expression then bind the pattern names to its results
evalStm :: FunEnv rep -> Env -> Stm rep -> InterpM rep Env
evalStm :: forall rep. FunEnv rep -> Env -> Stm rep -> InterpM rep Env
evalStm FunEnv rep
funs Env
env (Let Pat (LetDec rep)
pat StmAux (ExpDec rep)
aux Exp rep
expression) = do
  vals <-
    FunEnv rep -> Env -> Exp rep -> InterpM rep [Val]
forall rep. FunEnv rep -> Env -> Exp rep -> InterpM rep [Val]
evalExp FunEnv rep
funs Env
env Exp rep
expression InterpM rep [Val]
-> (Text -> InterpM rep [Val]) -> InterpM rep [Val]
forall a. InterpM rep a -> (Text -> InterpM rep a) -> InterpM rep a
forall e (m :: * -> *) a.
MonadError e m =>
m a -> (e -> m a) -> m a
`catchError` \Text
err ->
      Text -> InterpM rep [Val]
forall rep a. Text -> InterpM rep a
interpError (Text -> InterpM rep [Val]) -> Text -> InterpM rep [Val]
forall a b. (a -> b) -> a -> b
$ Provenance -> Text -> Text
addProvenance (StmAux (ExpDec rep) -> Provenance
forall dec. StmAux dec -> Provenance
stmAuxLoc StmAux (ExpDec rep)
aux) Text
err

  let names = (PatElem (LetDec rep) -> VName)
-> [PatElem (LetDec rep)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map PatElem (LetDec rep) -> VName
forall dec. PatElem dec -> VName
patElemName ([PatElem (LetDec rep)] -> [VName])
-> [PatElem (LetDec rep)] -> [VName]
forall a b. (a -> b) -> a -> b
$ Pat (LetDec rep) -> [PatElem (LetDec rep)]
forall dec. Pat dec -> [PatElem dec]
patElems Pat (LetDec rep)
pat
  pure $ M.union (M.fromList $ zip names vals) env

addProvenance :: Provenance -> T.Text -> T.Text
addProvenance :: Provenance -> Text -> Text
addProvenance (Provenance [Loc]
locations Loc
location) Text
message
  | Loc
location Loc -> Loc -> Bool
forall a. Eq a => a -> a -> Bool
== Loc
forall a. Monoid a => a
mempty = Text
message
  | Bool
otherwise =
      Text
message
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\n"
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> [Text] -> Text
prettyStacktrace
          Int
0
          ((Loc -> Text) -> [Loc] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Loc -> Text
forall a. Located a => a -> Text
locText ([Loc] -> [Text]) -> [Loc] -> [Text]
forall a b. (a -> b) -> a -> b
$ Loc
location Loc -> [Loc] -> [Loc]
forall a. a -> [a] -> [a]
: [Loc] -> [Loc]
forall a. [a] -> [a]
reverse [Loc]
locations)

-- Produce one Val per pattern element the expression is expected to bind.
evalExp :: FunEnv rep -> Env -> Exp rep -> InterpM rep [Val]
evalExp :: forall rep. FunEnv rep -> Env -> Exp rep -> InterpM rep [Val]
evalExp FunEnv rep
_ Env
env (BasicOp BasicOp
op) = Env -> BasicOp -> InterpM rep [Val]
forall rep. Env -> BasicOp -> InterpM rep [Val]
evalBasicOp Env
env BasicOp
op
evalExp FunEnv rep
funs Env
env (Match [SubExp]
ses [Case (Body rep)]
cases Body rep
default_body MatchDec (BranchType rep)
_) = do
  values <- (SubExp -> InterpM rep PrimValue)
-> [SubExp] -> InterpM rep [PrimValue]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Env -> SubExp -> InterpM rep Val
forall rep. Env -> SubExp -> InterpM rep Val
evalSubExp Env
env (SubExp -> InterpM rep Val)
-> (Val -> InterpM rep PrimValue)
-> SubExp
-> InterpM rep PrimValue
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Val -> InterpM rep PrimValue
forall rep. Val -> InterpM rep PrimValue
expectPrimVal) [SubExp]
ses
  evalBody funs env $ selectCase values cases
  where
    selectCase :: [PrimValue] -> [Case (Body rep)] -> Body rep
selectCase [PrimValue]
values (Case [Maybe PrimValue]
patterns Body rep
body : [Case (Body rep)]
remaining)
      | [Maybe PrimValue] -> [PrimValue] -> Bool
forall {b}. Eq b => [Maybe b] -> [b] -> Bool
matches [Maybe PrimValue]
patterns [PrimValue]
values = Body rep
body
      | Bool
otherwise = [PrimValue] -> [Case (Body rep)] -> Body rep
selectCase [PrimValue]
values [Case (Body rep)]
remaining
    selectCase [PrimValue]
_ [] = Body rep
default_body

    matches :: [Maybe b] -> [b] -> Bool
matches [Maybe b]
patterns [b]
values = [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ([Bool] -> Bool) -> [Bool] -> Bool
forall a b. (a -> b) -> a -> b
$ (Maybe b -> b -> Bool) -> [Maybe b] -> [b] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Maybe b -> b -> Bool
forall {a}. Eq a => Maybe a -> a -> Bool
matchesValue [Maybe b]
patterns [b]
values

    matchesValue :: Maybe a -> a -> Bool
matchesValue Maybe a
Nothing a
_ = Bool
True
    matchesValue (Just a
expected) a
actual = a
expected a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
actual
evalExp FunEnv rep
funs Env
env (Loop [(FParam rep, SubExp)]
params_and_inits (ForLoop VName
iterator IntType
int_type SubExp
bound_exp) Body rep
body) = do
  initial_values <- ((FParam rep, SubExp) -> InterpM rep Val)
-> [(FParam rep, SubExp)] -> InterpM rep [Val]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Env -> SubExp -> InterpM rep Val
forall rep. Env -> SubExp -> InterpM rep Val
evalSubExp Env
env (SubExp -> InterpM rep Val)
-> ((FParam rep, SubExp) -> SubExp)
-> (FParam rep, SubExp)
-> InterpM rep Val
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FParam rep, SubExp) -> SubExp
forall a b. (a, b) -> b
snd) [(FParam rep, SubExp)]
params_and_inits
  bound_value <- evalSubExp env bound_exp >>= expectPrimVal
  bound <- expectInt bound_value
  runIterations 0 bound initial_values
  where
    param_names :: [VName]
param_names = ((FParam rep, SubExp) -> VName)
-> [(FParam rep, SubExp)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (FParam rep -> VName
forall dec. Param dec -> VName
paramName (FParam rep -> VName)
-> ((FParam rep, SubExp) -> FParam rep)
-> (FParam rep, SubExp)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FParam rep, SubExp) -> FParam rep
forall a b. (a, b) -> a
fst) [(FParam rep, SubExp)]
params_and_inits

    runIterations :: t -> t -> [Val] -> InterpM rep [Val]
runIterations t
iteration t
bound [Val]
current_values
      | t
iteration t -> t -> Bool
forall a. Ord a => a -> a -> Bool
>= t
bound =
          [Val] -> InterpM rep [Val]
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Val]
current_values
      | Bool
otherwise = do
          let iterator_value :: Val
iterator_value =
                PrimValue -> Val
PrimVal (PrimValue -> Val) -> PrimValue -> Val
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
IntValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> t -> IntValue
forall int. Integral int => IntType -> int -> IntValue
P.intValue IntType
int_type t
iteration
              loop_bindings :: Env
loop_bindings =
                [(VName, Val)] -> Env
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Val)] -> Env) -> [(VName, Val)] -> Env
forall a b. (a -> b) -> a -> b
$
                  (VName
iterator, Val
iterator_value)
                    (VName, Val) -> [(VName, Val)] -> [(VName, Val)]
forall a. a -> [a] -> [a]
: [VName] -> [Val] -> [(VName, Val)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
param_names [Val]
current_values
              iteration_env :: Env
iteration_env =
                Env -> Env -> Env
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union Env
loop_bindings Env
env

          next_values <- FunEnv rep -> Env -> Body rep -> InterpM rep [Val]
forall rep. FunEnv rep -> Env -> Body rep -> InterpM rep [Val]
evalBody FunEnv rep
funs Env
iteration_env Body rep
body
          runIterations (iteration + 1) bound next_values
evalExp FunEnv rep
funs Env
env (Loop [(FParam rep, SubExp)]
params_and_inits (WhileLoop VName
condition) Body rep
body) = do
  initial_values <- ((FParam rep, SubExp) -> InterpM rep Val)
-> [(FParam rep, SubExp)] -> InterpM rep [Val]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Env -> SubExp -> InterpM rep Val
forall rep. Env -> SubExp -> InterpM rep Val
evalSubExp Env
env (SubExp -> InterpM rep Val)
-> ((FParam rep, SubExp) -> SubExp)
-> (FParam rep, SubExp)
-> InterpM rep Val
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FParam rep, SubExp) -> SubExp
forall a b. (a, b) -> b
snd) [(FParam rep, SubExp)]
params_and_inits
  runWhile initial_values
  where
    param_names :: [VName]
param_names = ((FParam rep, SubExp) -> VName)
-> [(FParam rep, SubExp)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (FParam rep -> VName
forall dec. Param dec -> VName
paramName (FParam rep -> VName)
-> ((FParam rep, SubExp) -> FParam rep)
-> (FParam rep, SubExp)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FParam rep, SubExp) -> FParam rep
forall a b. (a, b) -> a
fst) [(FParam rep, SubExp)]
params_and_inits

    runWhile :: [Val] -> InterpM rep [Val]
runWhile [Val]
current_values = do
      let loop_env :: Env
loop_env =
            Env -> Env -> Env
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union
              ([(VName, Val)] -> Env
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Val)] -> Env) -> [(VName, Val)] -> Env
forall a b. (a -> b) -> a -> b
$ [VName] -> [Val] -> [(VName, Val)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
param_names [Val]
current_values)
              Env
env

      condition_value <- Env -> VName -> InterpM rep Val
forall rep. Env -> VName -> InterpM rep Val
evalVar Env
loop_env VName
condition

      case condition_value of
        PrimVal (BoolValue Bool
False) ->
          [Val] -> InterpM rep [Val]
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Val]
current_values
        PrimVal (BoolValue Bool
True) ->
          [Val] -> InterpM rep [Val]
runWhile ([Val] -> InterpM rep [Val])
-> InterpM rep [Val] -> InterpM rep [Val]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< FunEnv rep -> Env -> Body rep -> InterpM rep [Val]
forall rep. FunEnv rep -> Env -> Body rep -> InterpM rep [Val]
evalBody FunEnv rep
funs Env
loop_env Body rep
body
        Val
_ ->
          [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"while-loop condition is not boolean"
evalExp FunEnv rep
funs Env
env (Apply Name
fname [(SubExp, Diet)]
args [(RetType rep, RetAls)]
_ Safety
_) = do
  arg_vals <- ((SubExp, Diet) -> InterpM rep Val)
-> [(SubExp, Diet)] -> InterpM rep [Val]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Env -> SubExp -> InterpM rep Val
forall rep. Env -> SubExp -> InterpM rep Val
evalSubExp Env
env (SubExp -> InterpM rep Val)
-> ((SubExp, Diet) -> SubExp) -> (SubExp, Diet) -> InterpM rep Val
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SubExp, Diet) -> SubExp
forall a b. (a, b) -> a
fst) [(SubExp, Diet)]
args

  case M.lookup fname funs of
    Just FunDef rep
callee ->
      let bindings :: Env
bindings = [(VName, Val)] -> Env
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Val)] -> Env) -> [(VName, Val)] -> Env
forall a b. (a -> b) -> a -> b
$ [VName] -> [Val] -> [(VName, Val)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((FParam rep -> VName) -> [FParam rep] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map FParam rep -> VName
forall dec. Param dec -> VName
paramName ([FParam rep] -> [VName]) -> [FParam rep] -> [VName]
forall a b. (a -> b) -> a -> b
$ FunDef rep -> [FParam rep]
forall rep. FunDef rep -> [FParam rep]
funDefParams FunDef rep
callee) [Val]
arg_vals
       in FunEnv rep -> Env -> Body rep -> InterpM rep [Val]
forall rep. FunEnv rep -> Env -> Body rep -> InterpM rep [Val]
evalBody FunEnv rep
funs (Env -> Env -> Env
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union Env
bindings Env
env) (FunDef rep -> Body rep
forall rep. FunDef rep -> Body rep
funDefBody FunDef rep
callee)
    Maybe (FunDef rep)
Nothing ->
      Name -> [Val] -> InterpM rep [Val]
forall rep. Name -> [Val] -> InterpM rep [Val]
evalPrimitiveFunction Name
fname [Val]
arg_vals
evalExp FunEnv rep
funs Env
env (Op Op rep
op) = do
  eval_op <- (InterpEnv rep -> FunEnv rep -> Env -> Op rep -> InterpM rep [Val])
-> InterpM rep (FunEnv rep -> Env -> Op rep -> InterpM rep [Val])
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks InterpEnv rep -> FunEnv rep -> Env -> Op rep -> InterpM rep [Val]
forall rep. InterpEnv rep -> OpEvaluator rep
interpOpEvaluator
  eval_op funs env op -- map/reduction/scan
evalExp FunEnv rep
funs Env
env (WithAcc [WithAccInput rep]
inputs Lambda rep
lambda) =
  FunEnv rep
-> Env -> [WithAccInput rep] -> Lambda rep -> InterpM rep [Val]
forall rep.
FunEnv rep
-> Env -> [WithAccInput rep] -> Lambda rep -> InterpM rep [Val]
evalWithAcc FunEnv rep
funs Env
env [WithAccInput rep]
inputs Lambda rep
lambda

evalPrimitiveFunction ::
  Name ->
  [Val] ->
  InterpM rep [Val]
evalPrimitiveFunction :: forall rep. Name -> [Val] -> InterpM rep [Val]
evalPrimitiveFunction Name
fname [Val]
args =
  case Text
-> Map Text ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
-> Maybe ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (Name -> Text
nameToText Name
fname) Map Text ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
P.primFuns of
    Maybe ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
Nothing ->
      [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InterpM rep [Val]) -> [Char] -> InterpM rep [Val]
forall a b. (a -> b) -> a -> b
$ [Char]
"function not found: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Name -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Name
fname
    Just ([PrimType]
_, PrimType
_, [PrimValue] -> Maybe PrimValue
function) -> do
      values <- (Val -> InterpM rep PrimValue) -> [Val] -> InterpM rep [PrimValue]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Val -> InterpM rep PrimValue
forall rep. Val -> InterpM rep PrimValue
expectPrimVal [Val]
args
      case function values of
        Just PrimValue
result -> [Val] -> InterpM rep [Val]
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> Val
PrimVal PrimValue
result]
        Maybe PrimValue
Nothing ->
          [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InterpM rep [Val]) -> [Char] -> InterpM rep [Val]
forall a b. (a -> b) -> a -> b
$ [Char]
"invalid arguments to primitive function: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Name -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Name
fname

evalWithAcc ::
  FunEnv rep ->
  Env ->
  [WithAccInput rep] ->
  Lambda rep ->
  InterpM rep [Val]
evalWithAcc :: forall rep.
FunEnv rep
-> Env -> [WithAccInput rep] -> Lambda rep -> InterpM rep [Val]
evalWithAcc FunEnv rep
funs Env
env [WithAccInput rep]
inputs Lambda rep
lambda = do
  evaluated_inputs <- (WithAccInput rep
 -> InterpM rep ([Int], [Val], Maybe (Lambda rep, [SubExp])))
-> [WithAccInput rep]
-> InterpM rep [([Int], [Val], Maybe (Lambda rep, [SubExp]))]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM WithAccInput rep
-> InterpM rep ([Int], [Val], Maybe (Lambda rep, [SubExp]))
forall {t :: * -> *} {c} {rep}.
Traversable t =>
(Shape, t VName, c) -> InterpM rep ([Int], t Val, c)
evaluateInput [WithAccInput rep]
inputs

  let accumulator_count = [WithAccInput rep] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [WithAccInput rep]
inputs
      (certificate_params, accumulator_params) =
        splitAt accumulator_count $ lambdaParams lambda
      certificates = (Param (LParamInfo rep) -> VName)
-> [Param (LParamInfo rep)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (LParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName [Param (LParamInfo rep)]
certificate_params
      accumulators = (VName -> ([Int], [Val], Maybe (Lambda rep, [SubExp])) -> Val)
-> [VName]
-> [([Int], [Val], Maybe (Lambda rep, [SubExp]))]
-> [Val]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith VName -> ([Int], [Val], Maybe (Lambda rep, [SubExp])) -> Val
forall {c}. VName -> ([Int], [Val], c) -> Val
mkAccumulator [VName]
certificates [([Int], [Val], Maybe (Lambda rep, [SubExp]))]
evaluated_inputs
      operators =
        [(VName, (FunEnv rep, Env, Lambda rep))]
-> Map VName (FunEnv rep, Env, Lambda rep)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, (FunEnv rep, Env, Lambda rep))]
 -> Map VName (FunEnv rep, Env, Lambda rep))
-> ([(VName, ([Int], [Val], Maybe (Lambda rep, [SubExp])))]
    -> [(VName, (FunEnv rep, Env, Lambda rep))])
-> [(VName, ([Int], [Val], Maybe (Lambda rep, [SubExp])))]
-> Map VName (FunEnv rep, Env, Lambda rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, ([Int], [Val], Maybe (Lambda rep, [SubExp])))
 -> Maybe (VName, (FunEnv rep, Env, Lambda rep)))
-> [(VName, ([Int], [Val], Maybe (Lambda rep, [SubExp])))]
-> [(VName, (FunEnv rep, Env, Lambda rep))]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (VName, ([Int], [Val], Maybe (Lambda rep, [SubExp])))
-> Maybe (VName, (FunEnv rep, Env, Lambda rep))
forall {a} {a} {b} {c} {b}.
(a, (a, b, Maybe (c, b))) -> Maybe (a, (FunEnv rep, Env, c))
accOperator ([(VName, ([Int], [Val], Maybe (Lambda rep, [SubExp])))]
 -> Map VName (FunEnv rep, Env, Lambda rep))
-> [(VName, ([Int], [Val], Maybe (Lambda rep, [SubExp])))]
-> Map VName (FunEnv rep, Env, Lambda rep)
forall a b. (a -> b) -> a -> b
$ [VName]
-> [([Int], [Val], Maybe (Lambda rep, [SubExp]))]
-> [(VName, ([Int], [Val], Maybe (Lambda rep, [SubExp])))]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
certificates [([Int], [Val], Maybe (Lambda rep, [SubExp]))]
evaluated_inputs
      -- Certificates are bound to their accumulator so that zero-iteration
      -- maps can recover it from the result type 'Acc c ...'.
      bindings =
        [(VName, Val)] -> Env
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Val)] -> Env) -> [(VName, Val)] -> Env
forall a b. (a -> b) -> a -> b
$
          [VName] -> [Val] -> [(VName, Val)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
certificates [Val]
accumulators
            [(VName, Val)] -> [(VName, Val)] -> [(VName, Val)]
forall a. Semigroup a => a -> a -> a
<> [VName] -> [Val] -> [(VName, Val)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Param (LParamInfo rep) -> VName)
-> [Param (LParamInfo rep)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (LParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName [Param (LParamInfo rep)]
accumulator_params) [Val]
accumulators

  results <-
    local (withAccOps operators) $
      evalBody funs (M.union bindings env) (lambdaBody lambda)
  pure $
    concatMap (\([Int]
_, [Val]
arrays, Maybe (Lambda rep, [SubExp])
_) -> [Val]
arrays) evaluated_inputs
      <> drop accumulator_count results
  where
    mkAccumulator :: VName -> ([Int], [Val], c) -> Val
mkAccumulator VName
certificate ([Int]
index_shape, [Val]
arrays, c
_) =
      Accumulator -> Val
AccValue (Accumulator -> Val) -> Accumulator -> Val
forall a b. (a -> b) -> a -> b
$ VName -> [Int] -> [Val] -> Accumulator
Accumulator VName
certificate [Int]
index_shape [Val]
arrays

    accOperator :: (a, (a, b, Maybe (c, b))) -> Maybe (a, (FunEnv rep, Env, c))
accOperator (a
certificate, (a
_, b
_, Just (c
operator_lambda, b
_))) =
      (a, (FunEnv rep, Env, c)) -> Maybe (a, (FunEnv rep, Env, c))
forall a. a -> Maybe a
Just (a
certificate, (FunEnv rep
funs, Env
env, c
operator_lambda))
    accOperator (a, (a, b, Maybe (c, b)))
_ = Maybe (a, (FunEnv rep, Env, c))
forall a. Maybe a
Nothing

    withAccOps :: Map VName (FunEnv rep, Env, Lambda rep)
-> InterpEnv rep -> InterpEnv rep
withAccOps Map VName (FunEnv rep, Env, Lambda rep)
new_operators InterpEnv rep
interp_env =
      InterpEnv rep
interp_env {interpAccOps = new_operators <> interpAccOps interp_env}

    evaluateInput :: (Shape, t VName, c) -> InterpM rep ([Int], t Val, c)
evaluateInput (Shape
shape, t VName
array_names, c
operator) = do
      index_shape <- Env -> Shape -> InterpM rep [Int]
forall rep. Env -> Shape -> InterpM rep [Int]
evalShape Env
env Shape
shape
      arrays <- mapM (evalVar env) array_names
      pure (index_shape, arrays, operator)

evalVar :: Env -> VName -> InterpM rep Val
evalVar :: forall rep. Env -> VName -> InterpM rep Val
evalVar Env
env VName
v =
  InterpM rep Val
-> (Val -> InterpM rep Val) -> Maybe Val -> InterpM rep Val
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ([Char] -> InterpM rep Val
forall a. HasCallStack => [Char] -> a
error ([Char] -> InterpM rep Val) -> [Char] -> InterpM rep Val
forall a b. (a -> b) -> a -> b
$ [Char]
"unbound variable: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> VName -> [Char]
forall a. Pretty a => a -> [Char]
prettyString VName
v) Val -> InterpM rep Val
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Val -> InterpM rep Val) -> Maybe Val -> InterpM rep Val
forall a b. (a -> b) -> a -> b
$ VName -> Env -> Maybe Val
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Env
env

evalSubExp :: Env -> SubExp -> InterpM rep Val
evalSubExp :: forall rep. Env -> SubExp -> InterpM rep Val
evalSubExp Env
_ (Constant PrimValue
pv) = Val -> InterpM rep Val
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Val -> InterpM rep Val) -> Val -> InterpM rep Val
forall a b. (a -> b) -> a -> b
$ PrimValue -> Val
PrimVal PrimValue
pv
evalSubExp Env
env (Var VName
v) = Env -> VName -> InterpM rep Val
forall rep. Env -> VName -> InterpM rep Val
evalVar Env
env VName
v

expectPrimVal :: Val -> InterpM rep PrimValue
expectPrimVal :: forall rep. Val -> InterpM rep PrimValue
expectPrimVal (PrimVal PrimValue
pv) = PrimValue -> InterpM rep PrimValue
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure PrimValue
pv
expectPrimVal Val
_ = [Char] -> InterpM rep PrimValue
forall a. HasCallStack => [Char] -> a
error [Char]
"expected a primitive value"

expectInt :: PrimValue -> InterpM rep Int
expectInt :: forall rep. PrimValue -> InterpM rep Int
expectInt (IntValue IntValue
i) = Int -> InterpM rep Int
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> InterpM rep Int) -> Int -> InterpM rep Int
forall a b. (a -> b) -> a -> b
$ IntValue -> Int
forall int. Integral int => IntValue -> int
P.valueIntegral IntValue
i
expectInt PrimValue
_ = [Char] -> InterpM rep Int
forall a. HasCallStack => [Char] -> a
error [Char]
"expected an integer value"

evalInt :: Env -> SubExp -> InterpM rep Int
evalInt :: forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env = Env -> SubExp -> InterpM rep Val
forall rep. Env -> SubExp -> InterpM rep Val
evalSubExp Env
env (SubExp -> InterpM rep Val)
-> (Val -> InterpM rep Int) -> SubExp -> InterpM rep Int
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Val -> InterpM rep PrimValue
forall rep. Val -> InterpM rep PrimValue
expectPrimVal (Val -> InterpM rep PrimValue)
-> (PrimValue -> InterpM rep Int) -> Val -> InterpM rep Int
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> PrimValue -> InterpM rep Int
forall rep. PrimValue -> InterpM rep Int
expectInt

-- Safe division-like operations yield zero on a zero divisor, like the code generators.
evalBinOp :: BinOp -> PrimValue -> PrimValue -> Maybe PrimValue
evalBinOp :: BinOp -> PrimValue -> PrimValue -> Maybe PrimValue
evalBinOp BinOp
op PrimValue
x PrimValue
y
  | BinOp -> Bool
safeDivision BinOp
op, PrimValue -> Bool
P.zeroIsh PrimValue
y = PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue) -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimValue
P.blankPrimValue (PrimType -> PrimValue) -> PrimType -> PrimValue
forall a b. (a -> b) -> a -> b
$ BinOp -> PrimType
P.binOpType BinOp
op
  | Bool
otherwise = BinOp -> PrimValue -> PrimValue -> Maybe PrimValue
P.doBinOp BinOp
op PrimValue
x PrimValue
y
  where
    safeDivision :: BinOp -> Bool
safeDivision (UDiv IntType
_ Safety
Safe) = Bool
True
    safeDivision (UCeilDiv IntType
_ Safety
Safe) = Bool
True
    safeDivision (SDiv IntType
_ Safety
Safe) = Bool
True
    safeDivision (SCeilDiv IntType
_ Safety
Safe) = Bool
True
    safeDivision (UMod IntType
_ Safety
Safe) = Bool
True
    safeDivision (SMod IntType
_ Safety
Safe) = Bool
True
    safeDivision (SQuot IntType
_ Safety
Safe) = Bool
True
    safeDivision (SRem IntType
_ Safety
Safe) = Bool
True
    safeDivision BinOp
_ = Bool
False

evalBasicOp :: Env -> BasicOp -> InterpM rep [Val]
evalBasicOp :: forall rep. Env -> BasicOp -> InterpM rep [Val]
evalBasicOp Env
env (SubExp SubExp
se) = Val -> [Val]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Val -> [Val]) -> InterpM rep Val -> InterpM rep [Val]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> SubExp -> InterpM rep Val
forall rep. Env -> SubExp -> InterpM rep Val
evalSubExp Env
env SubExp
se
evalBasicOp Env
env (BinOp BinOp
op SubExp
x SubExp
y) = do
  xv <- Val -> InterpM rep PrimValue
forall rep. Val -> InterpM rep PrimValue
expectPrimVal (Val -> InterpM rep PrimValue)
-> InterpM rep Val -> InterpM rep PrimValue
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Env -> SubExp -> InterpM rep Val
forall rep. Env -> SubExp -> InterpM rep Val
evalSubExp Env
env SubExp
x
  yv <- expectPrimVal =<< evalSubExp env y
  case evalBinOp op xv yv of
    Just PrimValue
result -> [Val] -> InterpM rep [Val]
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> Val
PrimVal PrimValue
result]
    Maybe PrimValue
Nothing -> Text -> InterpM rep [Val]
forall rep a. Text -> InterpM rep a
interpError Text
"invalid binary operation"
evalBasicOp Env
env (UnOp UnOp
op SubExp
x) = do
  xv <- Val -> InterpM rep PrimValue
forall rep. Val -> InterpM rep PrimValue
expectPrimVal (Val -> InterpM rep PrimValue)
-> InterpM rep Val -> InterpM rep PrimValue
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Env -> SubExp -> InterpM rep Val
forall rep. Env -> SubExp -> InterpM rep Val
evalSubExp Env
env SubExp
x
  case P.doUnOp op xv of
    Just PrimValue
result -> [Val] -> InterpM rep [Val]
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> Val
PrimVal PrimValue
result]
    Maybe PrimValue
Nothing -> [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"invalid unary operation"
evalBasicOp Env
env (CmpOp CmpOp
op SubExp
x SubExp
y) = do
  xv <- Val -> InterpM rep PrimValue
forall rep. Val -> InterpM rep PrimValue
expectPrimVal (Val -> InterpM rep PrimValue)
-> InterpM rep Val -> InterpM rep PrimValue
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Env -> SubExp -> InterpM rep Val
forall rep. Env -> SubExp -> InterpM rep Val
evalSubExp Env
env SubExp
x
  yv <- expectPrimVal =<< evalSubExp env y
  case P.doCmpOp op xv yv of
    Just Bool
result -> [Val] -> InterpM rep [Val]
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> Val
PrimVal (PrimValue -> Val) -> PrimValue -> Val
forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue Bool
result]
    Maybe Bool
Nothing -> [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"invalid comparison operation"
evalBasicOp Env
env (ConvOp ConvOp
op SubExp
x) = do
  xv <- Val -> InterpM rep PrimValue
forall rep. Val -> InterpM rep PrimValue
expectPrimVal (Val -> InterpM rep PrimValue)
-> InterpM rep Val -> InterpM rep PrimValue
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Env -> SubExp -> InterpM rep Val
forall rep. Env -> SubExp -> InterpM rep Val
evalSubExp Env
env SubExp
x
  case P.doConvOp op xv of
    Just PrimValue
result -> [Val] -> InterpM rep [Val]
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> Val
PrimVal PrimValue
result]
    Maybe PrimValue
Nothing -> [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"invalid conversion operation"
evalBasicOp Env
env (ArrayLit [SubExp]
elements (Prim PrimType
element_type)) = do
  values <- (SubExp -> InterpM rep Val) -> [SubExp] -> InterpM rep [Val]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Env -> SubExp -> InterpM rep Val
forall rep. Env -> SubExp -> InterpM rep Val
evalSubExp Env
env) [SubExp]
elements
  primitive_values <- mapM expectPrimVal values
  pure <$> newArrayValue [length elements] element_type primitive_values
evalBasicOp Env
env (ArrayLit [SubExp]
elements (Array PrimType
element_type Shape
row_shape_exps NoMode
_)) = do
  row_shape <- Env -> Shape -> InterpM rep [Int]
forall rep. Env -> Shape -> InterpM rep [Int]
evalShape Env
env Shape
row_shape_exps
  rows <- mapM (evalSubExp env) elements
  row_values <- mapM valElements rows
  pure
    <$> newArrayValue
      (length elements : row_shape)
      element_type
      (concat row_values)
evalBasicOp Env
_ (ArrayLit [SubExp]
_ Acc {}) =
  [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"accumulator array literals are not implemented"
evalBasicOp Env
_ (ArrayLit [SubExp]
_ Mem {}) =
  [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"memory array literals are unsupported"
evalBasicOp Env
_ (ArrayVal [PrimValue]
values PrimType
element_type) =
  Val -> [Val]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Val -> [Val]) -> InterpM rep Val -> InterpM rep [Val]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int] -> PrimType -> [PrimValue] -> InterpM rep Val
forall rep. [Int] -> PrimType -> [PrimValue] -> InterpM rep Val
newArrayValue [[PrimValue] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [PrimValue]
values] PrimType
element_type [PrimValue]
values
evalBasicOp Env
env (Assert SubExp
condition ErrorMsg SubExp
message) = do
  condition_value <- Env -> SubExp -> InterpM rep Val
forall rep. Env -> SubExp -> InterpM rep Val
evalSubExp Env
env SubExp
condition InterpM rep Val
-> (Val -> InterpM rep PrimValue) -> InterpM rep PrimValue
forall a b. InterpM rep a -> (a -> InterpM rep b) -> InterpM rep b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Val -> InterpM rep PrimValue
forall rep. Val -> InterpM rep PrimValue
expectPrimVal
  case condition_value of
    BoolValue Bool
True -> [Val] -> InterpM rep [Val]
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> Val
PrimVal PrimValue
UnitValue]
    BoolValue Bool
False -> Text -> InterpM rep [Val]
forall rep a. Text -> InterpM rep a
interpError (Text -> InterpM rep [Val])
-> InterpM rep Text -> InterpM rep [Val]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Env -> ErrorMsg SubExp -> InterpM rep Text
forall rep. Env -> ErrorMsg SubExp -> InterpM rep Text
evalErrorMsg Env
env ErrorMsg SubExp
message
    PrimValue
_ -> [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"assert condition is not boolean"
evalBasicOp Env
env (Index VName
array_name Slice SubExp
slice_exp) = do
  array <- Env -> VName -> InterpM rep Val
forall rep. Env -> VName -> InterpM rep Val
evalVar Env
env VName
array_name
  slice <- evalSlice env slice_exp
  case array of
    ArrayValue [Int]
shape ArrayValues
values ->
      Val -> [Val]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Val -> [Val]) -> InterpM rep Val -> InterpM rep [Val]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int] -> ArrayValues -> Slice Int -> InterpM rep Val
forall rep. [Int] -> ArrayValues -> Slice Int -> InterpM rep Val
indexArray [Int]
shape ArrayValues
values Slice Int
slice
    Val
_ ->
      [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"cannot index a non-array value"
evalBasicOp Env
env (Reshape VName
array_name NewShape SubExp
reshape) = do
  array <- Env -> VName -> InterpM rep Val
forall rep. Env -> VName -> InterpM rep Val
evalVar Env
env VName
array_name
  dimensions <- evalShape env $ newShape reshape
  case array of
    ArrayValue [Int]
old_shape ArrayValues
values ->
      case NewShape SubExp -> ReshapeKind
reshapeKind NewShape SubExp
reshape of
        ReshapeKind
ReshapeCoerce
          | [Int]
dimensions [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
== [Int]
old_shape ->
              [Val] -> InterpM rep [Val]
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [[Int] -> ArrayValues -> Val
ArrayValue [Int]
dimensions ArrayValues
values]
          | Bool
otherwise ->
              [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"coercion to a different shape"
        ReshapeKind
ReshapeArbitrary
          | [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
dimensions Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== ArrayValues -> Int
arrayValuesLength ArrayValues
values ->
              [Val] -> InterpM rep [Val]
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [[Int] -> ArrayValues -> Val
ArrayValue [Int]
dimensions ArrayValues
values]
          | Bool
otherwise ->
              [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"reshape element count mismatch"
    Val
_ ->
      [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"cannot reshape a non-array value"
evalBasicOp Env
env (Opaque OpaqueOp
OpaqueNil SubExp
se) =
  Val -> [Val]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Val -> [Val]) -> InterpM rep Val -> InterpM rep [Val]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> SubExp -> InterpM rep Val
forall rep. Env -> SubExp -> InterpM rep Val
evalSubExp Env
env SubExp
se
evalBasicOp Env
env (Opaque (OpaqueTrace Text
t) SubExp
se) = do
  IO () -> InterpM rep ()
forall a. IO a -> InterpM rep a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> InterpM rep ()) -> IO () -> InterpM rep ()
forall a b. (a -> b) -> a -> b
$ Handle -> Text -> IO ()
TIO.hPutStrLn Handle
stderr Text
t
  Val -> [Val]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Val -> [Val]) -> InterpM rep Val -> InterpM rep [Val]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> SubExp -> InterpM rep Val
forall rep. Env -> SubExp -> InterpM rep Val
evalSubExp Env
env SubExp
se
evalBasicOp Env
env (Manifest VName
array_name [Int]
_) = do
  array <- Env -> VName -> InterpM rep Val
forall rep. Env -> VName -> InterpM rep Val
evalVar Env
env VName
array_name
  case array of
    ArrayValue [Int]
shape ArrayValues
values ->
      Val -> [Val]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Val -> [Val]) -> (ArrayValues -> Val) -> ArrayValues -> [Val]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> ArrayValues -> Val
ArrayValue [Int]
shape (ArrayValues -> [Val])
-> InterpM rep ArrayValues -> InterpM rep [Val]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO ArrayValues -> InterpM rep ArrayValues
forall a. IO a -> InterpM rep a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (ArrayValues -> IO ArrayValues
cloneArrayValues ArrayValues
values)
    Val
_ ->
      [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"cannot manifest a non-array value"
evalBasicOp Env
env (Iota SubExp
count_sub_exp SubExp
start_sub_exp SubExp
stride_sub_exp IntType
int_type) = do
  count <- Env -> SubExp -> InterpM rep Int
forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env SubExp
count_sub_exp
  stride <- evalInt env stride_sub_exp
  start <- evalInt env start_sub_exp

  if count < 0
    then interpError "iota length cannot be negative"
    else do
      let values = do
            i <- [Int
0 .. Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
            pure $ IntValue $ P.intValue int_type $ start + i * stride
      pure <$> newArrayValue [count] (IntType int_type) values
evalBasicOp Env
env (Replicate (Shape [SubExp]
shape_exps) SubExp
val_exp) = do
  dimensions <- (SubExp -> InterpM rep Int) -> [SubExp] -> InterpM rep [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Env -> SubExp -> InterpM rep Int
forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env) [SubExp]
shape_exps
  if any (< 0) dimensions
    then interpError " replicate dimensions cannot be negative"
    else do
      val <- evalSubExp env val_exp
      let copies = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
dimensions

      case (dimensions, val) of
        ([], ArrayValue [Int]
shape ArrayValues
values) ->
          Val -> [Val]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Val -> [Val]) -> (ArrayValues -> Val) -> ArrayValues -> [Val]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> ArrayValues -> Val
ArrayValue [Int]
shape (ArrayValues -> [Val])
-> InterpM rep ArrayValues -> InterpM rep [Val]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO ArrayValues -> InterpM rep ArrayValues
forall a. IO a -> InterpM rep a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (ArrayValues -> IO ArrayValues
cloneArrayValues ArrayValues
values)
        ([], Val
_) -> [Val] -> InterpM rep [Val]
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Val
val]
        ([Int]
_, PrimVal PrimValue
primitive_value) ->
          Val -> [Val]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Val -> [Val]) -> InterpM rep Val -> InterpM rep [Val]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int] -> PrimType -> [PrimValue] -> InterpM rep Val
forall rep. [Int] -> PrimType -> [PrimValue] -> InterpM rep Val
newArrayValue [Int]
dimensions (PrimValue -> PrimType
P.primValueType PrimValue
primitive_value) (Int -> PrimValue -> [PrimValue]
forall a. Int -> a -> [a]
replicate Int
copies PrimValue
primitive_value)
        ([Int]
_, ArrayValue [Int]
old_shape ArrayValues
values) -> do
          primitive_values <- ArrayValues -> InterpM rep [PrimValue]
forall rep. ArrayValues -> InterpM rep [PrimValue]
arrayValues ArrayValues
values
          pure <$> newArrayValue (dimensions <> old_shape) (arrayValuesType values) (concat $ replicate copies primitive_values)
        ([Int]
_, AccValue {}) ->
          [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"cannot replicate an accumulator value"
evalBasicOp Env
env (Rearrange VName
array_name [Int]
permutation) = do
  array <- Env -> VName -> InterpM rep Val
forall rep. Env -> VName -> InterpM rep Val
evalVar Env
env VName
array_name
  case array of
    ArrayValue [Int]
old_shape ArrayValues
values -> do
      let new_shape :: [Int]
new_shape = [Int] -> [Int] -> [Int]
forall a. [Int] -> [a] -> [a]
rearrangeShape [Int]
permutation [Int]
old_shape
          oldCoordinate :: [a] -> [a]
oldCoordinate = [Int] -> [a] -> [a]
forall a. [Int] -> [a] -> [a]
rearrangeShape ([Int] -> [a] -> [a]) -> [Int] -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ [Int] -> [Int]
rearrangeInverse [Int]
permutation
      new_values <-
        ([Int] -> InterpM rep PrimValue)
-> [[Int]] -> InterpM rep [PrimValue]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM
          (ArrayValues -> Int -> InterpM rep PrimValue
forall rep. ArrayValues -> Int -> InterpM rep PrimValue
readArrayValue ArrayValues
values (Int -> InterpM rep PrimValue)
-> ([Int] -> Int) -> [Int] -> InterpM rep PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [Int] -> Int
linearIndex [Int]
old_shape ([Int] -> Int) -> ([Int] -> [Int]) -> [Int] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [Int]
forall a. [a] -> [a]
oldCoordinate)
          ([Int] -> [[Int]]
allCoordinates [Int]
new_shape)
      pure <$> newArrayValue new_shape (arrayValuesType values) new_values
    Val
_ -> [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"cannot rearrange a non-array value"
evalBasicOp Env
env (Concat Int
concat_dim NonEmpty VName
array_names SubExp
result_size_exp) = do
  arrays <- (VName -> InterpM rep ([Int], ArrayValues))
-> [VName] -> InterpM rep [([Int], ArrayValues)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> InterpM rep ([Int], ArrayValues)
forall {rep}. VName -> InterpM rep ([Int], ArrayValues)
lookupArray ([VName] -> InterpM rep [([Int], ArrayValues)])
-> [VName] -> InterpM rep [([Int], ArrayValues)]
forall a b. (a -> b) -> a -> b
$ NonEmpty VName -> [VName]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty VName
array_names
  declared_size <- evalInt env result_size_exp

  case arrays of
    [] ->
      [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"concat requires at least one array"
    ([Int]
first_shape, ArrayValues
first_values) : [([Int], ArrayValues)]
_
      | Int
declared_size Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [([Int], ArrayValues)] -> Int
forall {b}. [([Int], b)] -> Int
actualSize [([Int], ArrayValues)]
arrays ->
          [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"concat result size mismatch"
      | Bool
otherwise -> do
          let result_shape :: [Int]
result_shape =
                Int -> Int -> [Int] -> [Int]
forall a. Int -> a -> [a] -> [a]
replaceAt Int
concat_dim Int
declared_size [Int]
first_shape

          result_values <- ([Int] -> InterpM rep PrimValue)
-> [[Int]] -> InterpM rep [PrimValue]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([([Int], ArrayValues)] -> [Int] -> InterpM rep PrimValue
forall {rep}.
[([Int], ArrayValues)] -> [Int] -> InterpM rep PrimValue
valueAt [([Int], ArrayValues)]
arrays) ([[Int]] -> InterpM rep [PrimValue])
-> [[Int]] -> InterpM rep [PrimValue]
forall a b. (a -> b) -> a -> b
$ [Int] -> [[Int]]
allCoordinates [Int]
result_shape
          pure <$> newArrayValue result_shape (arrayValuesType first_values) result_values
  where
    lookupArray :: VName -> InterpM rep ([Int], ArrayValues)
lookupArray VName
name = do
      array <- Env -> VName -> InterpM rep Val
forall rep. Env -> VName -> InterpM rep Val
evalVar Env
env VName
name
      case array of
        ArrayValue [Int]
shape ArrayValues
values -> ([Int], ArrayValues) -> InterpM rep ([Int], ArrayValues)
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Int]
shape, ArrayValues
values)
        Val
_ -> [Char] -> InterpM rep ([Int], ArrayValues)
forall a. HasCallStack => [Char] -> a
error [Char]
"cannot concatenate a non-array value"

    actualSize :: [([Int], b)] -> Int
actualSize =
      [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> ([([Int], b)] -> [Int]) -> [([Int], b)] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (([Int], b) -> Int) -> [([Int], b)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (([Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
concat_dim) ([Int] -> Int) -> (([Int], b) -> [Int]) -> ([Int], b) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Int], b) -> [Int]
forall a b. (a, b) -> a
fst)

    valueAt :: [([Int], ArrayValues)] -> [Int] -> InterpM rep PrimValue
valueAt [([Int], ArrayValues)]
arrays [Int]
coordinate = do
      let concat_index :: Int
concat_index = [Int]
coordinate [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
concat_dim
          ([Int]
source_shape, ArrayValues
source_values, Int
local_index) =
            Int -> [([Int], ArrayValues)] -> ([Int], ArrayValues, Int)
forall {t} {b}. (Ord t, Num t) => t -> [([t], b)] -> ([t], b, t)
findSource Int
concat_index [([Int], ArrayValues)]
arrays
          source_coordinate :: [Int]
source_coordinate =
            Int -> Int -> [Int] -> [Int]
forall a. Int -> a -> [a] -> [a]
replaceAt Int
concat_dim Int
local_index [Int]
coordinate
          offset :: Int
offset =
            [Int] -> [Int] -> Int
linearIndex [Int]
source_shape [Int]
source_coordinate

      ArrayValues -> Int -> InterpM rep PrimValue
forall rep. ArrayValues -> Int -> InterpM rep PrimValue
readArrayValue ArrayValues
source_values Int
offset

    findSource :: t -> [([t], b)] -> ([t], b, t)
findSource t
_ [] =
      [Char] -> ([t], b, t)
forall a. HasCallStack => [Char] -> a
error [Char]
"invalid concat coordinate"
    findSource t
index (([t]
shape, b
values) : [([t], b)]
arrays)
      | t
index t -> t -> Bool
forall a. Ord a => a -> a -> Bool
< t
size =
          ([t]
shape, b
values, t
index)
      | Bool
otherwise =
          t -> [([t], b)] -> ([t], b, t)
findSource (t
index t -> t -> t
forall a. Num a => a -> a -> a
- t
size) [([t], b)]
arrays
      where
        size :: t
size = [t]
shape [t] -> Int -> t
forall a. HasCallStack => [a] -> Int -> a
!! Int
concat_dim
evalBasicOp Env
env (Update Safety
safety VName
array_name Slice SubExp
slice_exp SubExp
value_exp) = do
  array <- Env -> VName -> InterpM rep Val
forall rep. Env -> VName -> InterpM rep Val
evalVar Env
env VName
array_name
  slice <- evalSlice env slice_exp
  replacement <- evalSubExp env value_exp

  case array of
    ArrayValue [Int]
shape ArrayValues
values -> do
      in_bounds <-
        case Safety
safety of
          Safety
Safe -> Bool -> InterpM rep Bool
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> InterpM rep Bool) -> Bool -> InterpM rep Bool
forall a b. (a -> b) -> a -> b
$ [Int] -> Slice Int -> Bool
sliceInBounds [Int]
shape Slice Int
slice
          Safety
Unsafe -> Bool
True Bool -> InterpM rep () -> InterpM rep Bool
forall a b. a -> InterpM rep b -> InterpM rep a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ [Int] -> Slice Int -> InterpM rep ()
forall rep. [Int] -> Slice Int -> InterpM rep ()
checkSlice [Int]
shape Slice Int
slice
      when in_bounds $ do
        replacement_values <- valElements replacement
        zipWithM_
          (writeArrayValue values)
          (map (linearIndex shape) $ sliceCoordinates slice)
          replacement_values
      pure [array]
    Val
_ ->
      [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"cannot update a non-array value"
evalBasicOp Env
env (FlatIndex VName
array_name FlatSlice SubExp
flat_slice) = do
  array <- Env -> VName -> InterpM rep Val
forall rep. Env -> VName -> InterpM rep Val
evalVar Env
env VName
array_name
  (result_shape, offsets) <- evalFlatSlice env flat_slice

  case array of
    ArrayValue [Int]
_ ArrayValues
values -> do
      selected_values <- (Int -> InterpM rep PrimValue) -> [Int] -> InterpM rep [PrimValue]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (ArrayValues -> Int -> InterpM rep PrimValue
forall rep. ArrayValues -> Int -> InterpM rep PrimValue
readArrayValue ArrayValues
values) [Int]
offsets
      case (result_shape, selected_values) of
        ([], [PrimValue
primitive_value]) -> [Val] -> InterpM rep [Val]
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> Val
PrimVal PrimValue
primitive_value]
        ([], [PrimValue]
_) -> [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"invalid scalar flat index"
        ([Int], [PrimValue])
_ -> Val -> [Val]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Val -> [Val]) -> InterpM rep Val -> InterpM rep [Val]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int] -> PrimType -> [PrimValue] -> InterpM rep Val
forall rep. [Int] -> PrimType -> [PrimValue] -> InterpM rep Val
newArrayValue [Int]
result_shape (ArrayValues -> PrimType
arrayValuesType ArrayValues
values) [PrimValue]
selected_values
    Val
_ ->
      [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"cannot flat-index a non-array value"
evalBasicOp Env
env (FlatUpdate VName
source_name FlatSlice SubExp
flat_slice VName
replacement_name) = do
  source <- Env -> VName -> InterpM rep Val
forall rep. Env -> VName -> InterpM rep Val
evalVar Env
env VName
source_name
  replacement <- evalVar env replacement_name
  (_, offsets) <- evalFlatSlice env flat_slice
  case source of
    ArrayValue [Int]
_ ArrayValues
source_values
      | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (ArrayValues -> Int -> Bool
validOffset ArrayValues
source_values) [Int]
offsets ->
          Text -> InterpM rep [Val]
forall rep a. Text -> InterpM rep a
interpError Text
"flat update out of bounds"
      | Bool
otherwise -> do
          replacement_values <- Val -> InterpM rep [PrimValue]
forall rep. Val -> InterpM rep [PrimValue]
valElements Val
replacement
          zipWithM_ (writeArrayValue source_values) offsets replacement_values
          pure [source]
    Val
_ ->
      [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"cannot flat-update a non-array value"
  where
    validOffset :: ArrayValues -> Int -> Bool
validOffset ArrayValues
values Int
offset =
      Int
offset Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
offset Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< ArrayValues -> Int
arrayValuesLength ArrayValues
values
evalBasicOp Env
env (Scratch PrimType
element_type [SubExp]
dimension_exps) = do
  dimensions <- (SubExp -> InterpM rep Int) -> [SubExp] -> InterpM rep [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Env -> SubExp -> InterpM rep Int
forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env) [SubExp]
dimension_exps
  if any (< 0) dimensions
    then interpError "scratch dimensions cannot be negative"
    else
      let element_count = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
dimensions
          blank_value = PrimType -> PrimValue
P.blankPrimValue PrimType
element_type
       in pure <$> newArrayValue dimensions element_type (replicate element_count blank_value)
evalBasicOp Env
env (UserParam Name
_ SubExp
default_sub_exp) =
  Val -> [Val]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Val -> [Val]) -> InterpM rep Val -> InterpM rep [Val]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> SubExp -> InterpM rep Val
forall rep. Env -> SubExp -> InterpM rep Val
evalSubExp Env
env SubExp
default_sub_exp
evalBasicOp Env
env (UpdateAcc Safety
safety VName
accumulator_name [SubExp]
index_exps [SubExp]
value_exps) =
  Env -> Safety -> VName -> [SubExp] -> [SubExp] -> InterpM rep [Val]
forall rep.
Env -> Safety -> VName -> [SubExp] -> [SubExp] -> InterpM rep [Val]
evalUpdateAcc Env
env Safety
safety VName
accumulator_name [SubExp]
index_exps [SubExp]
value_exps

evalUpdateAcc ::
  Env ->
  Safety ->
  VName ->
  [SubExp] ->
  [SubExp] ->
  InterpM rep [Val]
evalUpdateAcc :: forall rep.
Env -> Safety -> VName -> [SubExp] -> [SubExp] -> InterpM rep [Val]
evalUpdateAcc Env
env Safety
safety VName
accumulator_name [SubExp]
index_exps [SubExp]
value_exps = do
  accumulator <- Env -> VName -> InterpM rep Val
forall rep. Env -> VName -> InterpM rep Val
evalVar Env
env VName
accumulator_name
  indices <- mapM (evalInt env) index_exps
  values <- mapM (evalSubExp env) value_exps

  case accumulator of
    AccValue Accumulator
acc -> do
      Accumulator -> [Int] -> [Val] -> InterpM rep ()
forall {rep}. Accumulator -> [Int] -> [Val] -> InterpM rep ()
updateAccumulator Accumulator
acc [Int]
indices [Val]
values
      [Val] -> InterpM rep [Val]
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Val
accumulator]
    Val
_ ->
      [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"UpdateAcc argument is not an accumulator"
  where
    updateAccumulator :: Accumulator -> [Int] -> [Val] -> InterpM rep ()
updateAccumulator Accumulator
acc [Int]
indices [Val]
new_values
      | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [Int] -> [Int] -> Bool
forall {a}. (Ord a, Num a) => [a] -> [a] -> Bool
indicesInBounds (Accumulator -> [Int]
accShape Accumulator
acc) [Int]
indices =
          case Safety
safety of
            Safety
Safe -> () -> InterpM rep ()
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            Safety
Unsafe -> Text -> InterpM rep ()
forall rep a. Text -> InterpM rep a
interpError Text
"unsafe accumulator update out of bounds"
      | Bool
otherwise = do
          operator <- (InterpEnv rep -> Maybe (FunEnv rep, Env, Lambda rep))
-> InterpM rep (Maybe (FunEnv rep, Env, Lambda rep))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((InterpEnv rep -> Maybe (FunEnv rep, Env, Lambda rep))
 -> InterpM rep (Maybe (FunEnv rep, Env, Lambda rep)))
-> (InterpEnv rep -> Maybe (FunEnv rep, Env, Lambda rep))
-> InterpM rep (Maybe (FunEnv rep, Env, Lambda rep))
forall a b. (a -> b) -> a -> b
$ VName
-> Map VName (FunEnv rep, Env, Lambda rep)
-> Maybe (FunEnv rep, Env, Lambda rep)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (Accumulator -> VName
accCert Accumulator
acc) (Map VName (FunEnv rep, Env, Lambda rep)
 -> Maybe (FunEnv rep, Env, Lambda rep))
-> (InterpEnv rep -> Map VName (FunEnv rep, Env, Lambda rep))
-> InterpEnv rep
-> Maybe (FunEnv rep, Env, Lambda rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InterpEnv rep -> Map VName (FunEnv rep, Env, Lambda rep)
forall rep.
InterpEnv rep -> Map VName (FunEnv rep, Env, Lambda rep)
interpAccOps
          replacement_values <-
            case operator of
              Maybe (FunEnv rep, Env, Lambda rep)
Nothing -> [Val] -> InterpM rep [Val]
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Val]
new_values
              Just (FunEnv rep
operator_funs, Env
operator_env, Lambda rep
operator_lambda) -> do
                old_values <-
                  (Val -> InterpM rep Val) -> [Val] -> InterpM rep [Val]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Int] -> Val -> InterpM rep Val
forall rep. [Int] -> Val -> InterpM rep Val
readAccumulatorElement [Int]
indices) ([Val] -> InterpM rep [Val]) -> [Val] -> InterpM rep [Val]
forall a b. (a -> b) -> a -> b
$ Accumulator -> [Val]
accArrays Accumulator
acc
                evalLambda operator_funs operator_env operator_lambda $
                  map int64Val indices <> old_values <> new_values

          zipWithM_
            (writeAccumulatorElementInPlace indices)
            (accArrays acc)
            replacement_values

    indicesInBounds :: [a] -> [a] -> Bool
indicesInBounds [a]
shape [a]
indices =
      [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ([Bool] -> Bool) -> [Bool] -> Bool
forall a b. (a -> b) -> a -> b
$ (a -> a -> Bool) -> [a] -> [a] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\a
size a
index -> a
index a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
0 Bool -> Bool -> Bool
&& a
index a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
size) [a]
shape [a]
indices

evalErrorMsg :: Env -> ErrorMsg SubExp -> InterpM rep T.Text
evalErrorMsg :: forall rep. Env -> ErrorMsg SubExp -> InterpM rep Text
evalErrorMsg Env
env (ErrorMsg [ErrorMsgPart SubExp]
parts) = do
  evaluatedParts <- (ErrorMsgPart SubExp -> InterpM rep Text)
-> [ErrorMsgPart SubExp] -> InterpM rep [Text]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ErrorMsgPart SubExp -> InterpM rep Text
forall {rep}. ErrorMsgPart SubExp -> InterpM rep Text
evalPart [ErrorMsgPart SubExp]
parts
  pure $ T.concat ("Error " : evaluatedParts)
  where
    evalPart :: ErrorMsgPart SubExp -> InterpM rep Text
evalPart (ErrorString Text
text) = Text -> InterpM rep Text
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
text
    evalPart (ErrorVal PrimType
_ SubExp
sub_exp) =
      PrimValue -> Text
renderErrorValue (PrimValue -> Text) -> InterpM rep PrimValue -> InterpM rep Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Env -> SubExp -> InterpM rep Val
forall rep. Env -> SubExp -> InterpM rep Val
evalSubExp Env
env SubExp
sub_exp InterpM rep Val
-> (Val -> InterpM rep PrimValue) -> InterpM rep PrimValue
forall a b. InterpM rep a -> (a -> InterpM rep b) -> InterpM rep b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Val -> InterpM rep PrimValue
forall rep. Val -> InterpM rep PrimValue
expectPrimVal)

renderErrorValue :: PrimValue -> T.Text
renderErrorValue :: PrimValue -> Text
renderErrorValue (IntValue IntValue
val) =
  Integer -> Text
forall a. Show a => a -> Text
showText (IntValue -> Integer
forall int. Integral int => IntValue -> int
P.valueIntegral IntValue
val :: Integer)
renderErrorValue (FloatValue (Float16Value Half
val)) =
  Half -> Text
forall a. Show a => a -> Text
showText Half
val
renderErrorValue (FloatValue (Float32Value Float
val)) =
  Float -> Text
forall a. Show a => a -> Text
showText Float
val
renderErrorValue (FloatValue (Float64Value Double
val)) =
  Double -> Text
forall a. Show a => a -> Text
showText Double
val
renderErrorValue (BoolValue Bool
True) = Text
"true"
renderErrorValue (BoolValue Bool
False) = Text
"false"
renderErrorValue PrimValue
UnitValue = Text
"()"

evalFlatSlice :: Env -> FlatSlice SubExp -> InterpM rep ([Int], [Int])
evalFlatSlice :: forall rep. Env -> FlatSlice SubExp -> InterpM rep ([Int], [Int])
evalFlatSlice Env
env (FlatSlice SubExp
offset_exp [FlatDimIndex SubExp]
dimensions) = do
  offset <- Env -> SubExp -> InterpM rep Int
forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env SubExp
offset_exp
  evaluated_dimensions <- mapM evalDimension dimensions

  let result_shape = ((Int, Int) -> Int) -> [(Int, Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Int) -> Int
forall a b. (a, b) -> a
fst [(Int, Int)]
evaluated_dimensions
      strides = ((Int, Int) -> Int) -> [(Int, Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Int) -> Int
forall a b. (a, b) -> b
snd [(Int, Int)]
evaluated_dimensions
      offsets = ([Int] -> Int) -> [[Int]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ((Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+) (Int -> Int) -> ([Int] -> Int) -> [Int] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> ([Int] -> [Int]) -> [Int] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) [Int]
strides) ([[Int]] -> [Int]) -> [[Int]] -> [Int]
forall a b. (a -> b) -> a -> b
$ [Int] -> [[Int]]
allCoordinates [Int]
result_shape
  if any (< 0) result_shape
    then interpError "flat slice dimensions cannot be negative"
    else pure (result_shape, offsets)
  where
    evalDimension :: FlatDimIndex SubExp -> InterpM rep (Int, Int)
evalDimension (FlatDimIndex SubExp
size_exp SubExp
stride_exp) = do
      size <- Env -> SubExp -> InterpM rep Int
forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env SubExp
size_exp
      stride <- evalInt env stride_exp
      pure (size, stride)

replaceAt :: Int -> a -> [a] -> [a]
replaceAt :: forall a. Int -> a -> [a] -> [a]
replaceAt Int
index a
val [a]
xs =
  Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
index [a]
xs [a] -> [a] -> [a]
forall a. Semigroup a => a -> a -> a
<> [a
val] [a] -> [a] -> [a]
forall a. Semigroup a => a -> a -> a
<> Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop (Int
index Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [a]
xs

-- The elements of a value in row-major order.
valElements :: Val -> InterpM rep [PrimValue]
valElements :: forall rep. Val -> InterpM rep [PrimValue]
valElements (PrimVal PrimValue
primitive_value) = [PrimValue] -> InterpM rep [PrimValue]
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue
primitive_value]
valElements (ArrayValue [Int]
_ ArrayValues
values) = ArrayValues -> InterpM rep [PrimValue]
forall rep. ArrayValues -> InterpM rep [PrimValue]
arrayValues ArrayValues
values
valElements AccValue {} = [Char] -> InterpM rep [PrimValue]
forall a. HasCallStack => [Char] -> a
error [Char]
"accumulators have no elements"

-- All coordinates of an array with the given shape, in row-major order.
allCoordinates :: [Int] -> [[Int]]
allCoordinates :: [Int] -> [[Int]]
allCoordinates = (Int -> [Int]) -> [Int] -> [[Int]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((Int -> [Int]) -> [Int] -> [[Int]])
-> (Int -> [Int]) -> [Int] -> [[Int]]
forall a b. (a -> b) -> a -> b
$ \Int
size -> [Int
0 .. Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]

evalSlice :: Env -> Slice SubExp -> InterpM rep (Slice Int)
evalSlice :: forall rep. Env -> Slice SubExp -> InterpM rep (Slice Int)
evalSlice Env
env = (SubExp -> InterpM rep Int)
-> Slice SubExp -> InterpM rep (Slice Int)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Slice a -> f (Slice b)
traverse ((SubExp -> InterpM rep Int)
 -> Slice SubExp -> InterpM rep (Slice Int))
-> (SubExp -> InterpM rep Int)
-> Slice SubExp
-> InterpM rep (Slice Int)
forall a b. (a -> b) -> a -> b
$ Env -> SubExp -> InterpM rep Int
forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env

sliceInBounds :: [Int] -> Slice Int -> Bool
sliceInBounds :: [Int] -> Slice Int -> Bool
sliceInBounds [Int]
shape (Slice [DimIndex Int]
dimensions) =
  [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ([Bool] -> Bool) -> [Bool] -> Bool
forall a b. (a -> b) -> a -> b
$ (Int -> DimIndex Int -> Bool) -> [Int] -> [DimIndex Int] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> DimIndex Int -> Bool
forall {a}. (Ord a, Num a) => a -> DimIndex a -> Bool
dimensionInBounds [Int]
shape [DimIndex Int]
dimensions
  where
    dimensionInBounds :: a -> DimIndex a -> Bool
dimensionInBounds a
size (DimFix a
index) =
      a -> a -> Bool
forall {a}. (Ord a, Num a) => a -> a -> Bool
validIndex a
size a
index
    dimensionInBounds a
size (DimSlice a
start a
count a
stride) =
      a
count a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0
        Bool -> Bool -> Bool
|| (a
count a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
0 Bool -> Bool -> Bool
&& a -> a -> Bool
forall {a}. (Ord a, Num a) => a -> a -> Bool
validIndex a
size a
start Bool -> Bool -> Bool
&& a -> a -> Bool
forall {a}. (Ord a, Num a) => a -> a -> Bool
validIndex a
size (a
start a -> a -> a
forall a. Num a => a -> a -> a
+ (a
count a -> a -> a
forall a. Num a => a -> a -> a
- a
1) a -> a -> a
forall a. Num a => a -> a -> a
* a
stride))

    validIndex :: a -> a -> Bool
validIndex a
size a
index = a
index a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
0 Bool -> Bool -> Bool
&& a
index a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
size

checkSlice :: [Int] -> Slice Int -> InterpM rep ()
checkSlice :: forall rep. [Int] -> Slice Int -> InterpM rep ()
checkSlice [Int]
shape Slice Int
slice
  | (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) ([Int] -> Bool) -> [Int] -> Bool
forall a b. (a -> b) -> a -> b
$ Slice Int -> [Int]
forall d. Slice d -> [d]
sliceDims Slice Int
slice =
      Text -> InterpM rep ()
forall rep a. Text -> InterpM rep a
interpError Text
"slice length cannot be negative"
  | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [Int] -> Slice Int -> Bool
sliceInBounds [Int]
shape Slice Int
slice =
      Text -> InterpM rep ()
forall rep a. Text -> InterpM rep a
interpError Text
"array index out of bounds"
  | Bool
otherwise =
      () -> InterpM rep ()
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- The coordinates in the source array selected by a slice, in row-major order
-- of the result.
sliceCoordinates :: Slice Int -> [[Int]]
sliceCoordinates :: Slice Int -> [[Int]]
sliceCoordinates Slice Int
slice =
  ([Int] -> [Int]) -> [[Int]] -> [[Int]]
forall a b. (a -> b) -> [a] -> [b]
map (Slice Int -> [Int] -> [Int]
forall d. Num d => Slice d -> [d] -> [d]
fixSlice Slice Int
slice) ([[Int]] -> [[Int]]) -> [[Int]] -> [[Int]]
forall a b. (a -> b) -> a -> b
$ [Int] -> [[Int]]
allCoordinates ([Int] -> [[Int]]) -> [Int] -> [[Int]]
forall a b. (a -> b) -> a -> b
$ Slice Int -> [Int]
forall d. Slice d -> [d]
sliceDims Slice Int
slice

indexArray :: [Int] -> ArrayValues -> Slice Int -> InterpM rep Val
indexArray :: forall rep. [Int] -> ArrayValues -> Slice Int -> InterpM rep Val
indexArray [Int]
shape ArrayValues
values Slice Int
slice = do
  [Int] -> Slice Int -> InterpM rep ()
forall rep. [Int] -> Slice Int -> InterpM rep ()
checkSlice [Int]
shape Slice Int
slice
  selected_values <-
    ([Int] -> InterpM rep PrimValue)
-> [[Int]] -> InterpM rep [PrimValue]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (ArrayValues -> Int -> InterpM rep PrimValue
forall rep. ArrayValues -> Int -> InterpM rep PrimValue
readArrayValue ArrayValues
values (Int -> InterpM rep PrimValue)
-> ([Int] -> Int) -> [Int] -> InterpM rep PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [Int] -> Int
linearIndex [Int]
shape) ([[Int]] -> InterpM rep [PrimValue])
-> [[Int]] -> InterpM rep [PrimValue]
forall a b. (a -> b) -> a -> b
$ Slice Int -> [[Int]]
sliceCoordinates Slice Int
slice
  case (sliceDims slice, selected_values) of
    ([], [PrimValue
primitive_value]) -> Val -> InterpM rep Val
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Val -> InterpM rep Val) -> Val -> InterpM rep Val
forall a b. (a -> b) -> a -> b
$ PrimValue -> Val
PrimVal PrimValue
primitive_value
    ([], [PrimValue]
_) -> [Char] -> InterpM rep Val
forall a. HasCallStack => [Char] -> a
error [Char]
"invalid scalar index"
    ([Int]
result_shape, [PrimValue]
_) ->
      [Int] -> PrimType -> [PrimValue] -> InterpM rep Val
forall rep. [Int] -> PrimType -> [PrimValue] -> InterpM rep Val
newArrayValue [Int]
result_shape (ArrayValues -> PrimType
arrayValuesType ArrayValues
values) [PrimValue]
selected_values

linearIndex :: [Int] -> [Int] -> Int
linearIndex :: [Int] -> [Int] -> Int
linearIndex [Int]
shape [Int]
indices =
  (Int -> (Int, Int) -> Int) -> Int -> [(Int, Int)] -> Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\Int
acc (Int
dim_size, Int
index) -> Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
dim_size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
index) Int
0 ([(Int, Int)] -> Int) -> [(Int, Int)] -> Int
forall a b. (a -> b) -> a -> b
$ [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
shape [Int]
indices

evalSOAC :: FunEnv rep -> Env -> SOAC rep -> InterpM rep [Val]
evalSOAC :: forall rep. FunEnv rep -> Env -> SOAC rep -> InterpM rep [Val]
evalSOAC FunEnv rep
funs Env
env (Screma SubExp
width_exp [VName]
input_names ScremaForm rep
form) =
  FunEnv rep
-> Env -> SubExp -> [VName] -> ScremaForm rep -> InterpM rep [Val]
forall rep.
FunEnv rep
-> Env -> SubExp -> [VName] -> ScremaForm rep -> InterpM rep [Val]
evalScrema FunEnv rep
funs Env
env SubExp
width_exp [VName]
input_names ScremaForm rep
form
evalSOAC FunEnv rep
funs Env
env (Stream SubExp
width_exp [VName]
input_names [SubExp]
initial_accumulators Lambda rep
lambda) =
  FunEnv rep
-> Env
-> SubExp
-> [VName]
-> [SubExp]
-> Lambda rep
-> InterpM rep [Val]
forall rep.
FunEnv rep
-> Env
-> SubExp
-> [VName]
-> [SubExp]
-> Lambda rep
-> InterpM rep [Val]
evalStream FunEnv rep
funs Env
env SubExp
width_exp [VName]
input_names [SubExp]
initial_accumulators Lambda rep
lambda
evalSOAC FunEnv rep
funs Env
env (Hist SubExp
width_exp [VName]
input_names [HistOp rep]
hist_ops Lambda rep
lambda) = FunEnv rep
-> Env
-> SubExp
-> [VName]
-> [HistOp rep]
-> Lambda rep
-> InterpM rep [Val]
forall rep.
FunEnv rep
-> Env
-> SubExp
-> [VName]
-> [HistOp rep]
-> Lambda rep
-> InterpM rep [Val]
evalHist FunEnv rep
funs Env
env SubExp
width_exp [VName]
input_names [HistOp rep]
hist_ops Lambda rep
lambda
evalSOAC FunEnv rep
funs Env
env (FlatMap SubExp
width_exp [VName]
input_names ExtLambda rep
lambda) = FunEnv rep
-> Env -> SubExp -> [VName] -> ExtLambda rep -> InterpM rep [Val]
forall rep.
FunEnv rep
-> Env -> SubExp -> [VName] -> ExtLambda rep -> InterpM rep [Val]
evalFlatMap FunEnv rep
funs Env
env SubExp
width_exp [VName]
input_names ExtLambda rep
lambda
evalSOAC FunEnv rep
_ Env
_ JVP {} = [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"evalSOAC: JVP is not implemented"
evalSOAC FunEnv rep
_ Env
_ VJP {} = [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"evalSOAC: VJP is not implemented"
evalSOAC FunEnv rep
_ Env
_ WithVJP {} = [Char] -> InterpM rep [Val]
forall a. HasCallStack => [Char] -> a
error [Char]
"evalSOAC: WithVJP is not implemented"

evalFlatMap ::
  FunEnv rep ->
  Env ->
  SubExp ->
  [VName] ->
  ExtLambda rep ->
  InterpM rep [Val]
evalFlatMap :: forall rep.
FunEnv rep
-> Env -> SubExp -> [VName] -> ExtLambda rep -> InterpM rep [Val]
evalFlatMap FunEnv rep
funs Env
env SubExp
width_exp [VName]
input_names ExtLambda rep
lambda = do
  width <- Env -> SubExp -> InterpM rep Int
forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env SubExp
width_exp
  inputs <- mapM (evalVar env) input_names
  rows <- mapM (runIteration inputs) [0 .. width - 1]
  let sizes = ((Int, [Val]) -> Int) -> [(Int, [Val])] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, [Val]) -> Int
forall a b. (a, b) -> a
fst [(Int, [Val])]
rows
      value_rows = ((Int, [Val]) -> [Val]) -> [(Int, [Val])] -> [[Val]]
forall a b. (a -> b) -> [a] -> [b]
map (Int, [Val]) -> [Val]
forall a b. (a, b) -> b
snd [(Int, [Val])]
rows
      offsets = [Int] -> [Int]
forall a. HasCallStack => [a] -> [a]
init ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
0 [Int]
sizes
      total_size = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int]
sizes
      return_types = Int -> [ExtType] -> [ExtType]
forall a. Int -> [a] -> [a]
drop Int
1 ([ExtType] -> [ExtType]) -> [ExtType] -> [ExtType]
forall a b. (a -> b) -> a -> b
$ ExtLambda rep -> [ExtType]
forall rep t. GLambda rep t -> [t]
lambdaReturnType ExtLambda rep
lambda
      columns
        | [[Val]] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[Val]]
value_rows = Int -> [Val] -> [[Val]]
forall a. Int -> a -> [a]
replicate ([ExtType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ExtType]
return_types) []
        | Bool
otherwise = [[Val]] -> [[Val]]
forall a. [[a]] -> [[a]]
L.transpose [[Val]]
value_rows
  values <-
    zipWithM
      (collectFlatMapOutput env total_size)
      return_types
      columns
  let flags = (Int -> [Bool]) -> [Int] -> [Bool]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Int -> [Bool]
segmentFlags [Int]
sizes
  sizes_array <- newArrayValue [width] int64_type $ map int64Prim sizes
  flags_array <- newArrayValue [total_size] Bool $ map BoolValue flags
  offsets_array <- newArrayValue [width] int64_type $ map int64Prim offsets

  pure $
    [int64Val total_size, sizes_array, flags_array, offsets_array]
      <> values
  where
    runIteration :: [Val] -> Int -> InterpM rep (Int, [Val])
runIteration [Val]
inputs Int
index = do
      input_rows <- (Val -> InterpM rep Val) -> [Val] -> InterpM rep [Val]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Int -> Val -> InterpM rep Val
forall rep. Int -> Val -> InterpM rep Val
rowAt Int
index) [Val]
inputs
      results <- evalLambda funs env lambda input_rows

      case results of
        Val
size_value : [Val]
values -> do
          size <- Val -> InterpM rep PrimValue
forall rep. Val -> InterpM rep PrimValue
expectPrimVal Val
size_value InterpM rep PrimValue
-> (PrimValue -> InterpM rep Int) -> InterpM rep Int
forall a b. InterpM rep a -> (a -> InterpM rep b) -> InterpM rep b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= PrimValue -> InterpM rep Int
forall rep. PrimValue -> InterpM rep Int
expectInt
          if size < 0
            then interpError "FlatMap segment size cannot be negative"
            else pure (size, values)
        [] ->
          [Char] -> InterpM rep (Int, [Val])
forall a. HasCallStack => [Char] -> a
error [Char]
"FlatMap lambda returned no segment size"

    segmentFlags :: Int -> [Bool]
segmentFlags Int
size
      | Int
size Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 = Bool
True Bool -> [Bool] -> [Bool]
forall a. a -> [a] -> [a]
: Int -> Bool -> [Bool]
forall a. Int -> a -> [a]
replicate (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Bool
False
      | Bool
otherwise = []

    int64_type :: PrimType
int64_type = IntType -> PrimType
IntType IntType
Int64
    int64Prim :: Int -> PrimValue
int64Prim = IntValue -> PrimValue
IntValue (IntValue -> PrimValue) -> (Int -> IntValue) -> Int -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> IntValue
Int64Value (Int64 -> IntValue) -> (Int -> Int64) -> Int -> IntValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral

evalHist ::
  FunEnv rep ->
  Env ->
  SubExp ->
  [VName] ->
  [HistOp rep] ->
  Lambda rep ->
  InterpM rep [Val]
evalHist :: forall rep.
FunEnv rep
-> Env
-> SubExp
-> [VName]
-> [HistOp rep]
-> Lambda rep
-> InterpM rep [Val]
evalHist FunEnv rep
funs Env
env SubExp
width_exp [VName]
input_names [HistOp rep]
hist_ops Lambda rep
bucket_lambda = do
  width <- Env -> SubExp -> InterpM rep Int
forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env SubExp
width_exp
  inputs <- mapM (evalVar env) input_names
  initial_histograms <- mapM (mapM (evalVar env) . histDest) hist_ops
  final_histograms <-
    foldM
      (runIteration inputs)
      initial_histograms
      [0 .. width - 1]

  pure $ concat final_histograms
  where
    index_counts :: [Int]
index_counts = (HistOp rep -> Int) -> [HistOp rep] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank (Shape -> Int) -> (HistOp rep -> Shape) -> HistOp rep -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HistOp rep -> Shape
forall rep. HistOp rep -> Shape
histShape) [HistOp rep]
hist_ops
    value_counts :: [Int]
value_counts = (HistOp rep -> Int) -> [HistOp rep] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ([VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([VName] -> Int) -> (HistOp rep -> [VName]) -> HistOp rep -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HistOp rep -> [VName]
forall rep. HistOp rep -> [VName]
histDest) [HistOp rep]
hist_ops

    runIteration :: [Val] -> [[Val]] -> Int -> InterpM rep [[Val]]
runIteration [Val]
inputs [[Val]]
histograms Int
iteration = do
      input_rows <- (Val -> InterpM rep Val) -> [Val] -> InterpM rep [Val]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Int -> Val -> InterpM rep Val
forall rep. Int -> Val -> InterpM rep Val
rowAt Int
iteration) [Val]
inputs
      bucket_results <- evalLambda funs env bucket_lambda input_rows

      let (index_groups, remaining) = splitGroups index_counts bucket_results
          (value_groups, _) = splitGroups value_counts remaining

      index_groups' <-
        mapM
          (mapM (expectPrimVal >=> expectInt))
          index_groups

      forM (zip3 hist_ops index_groups' (zip value_groups histograms)) $
        \(HistOp rep
hist_operation, [Int]
indices, ([Val], [Val])
value_and_histograms) ->
          HistOp rep -> [Int] -> ([Val], [Val]) -> InterpM rep [Val]
updateHistogram HistOp rep
hist_operation [Int]
indices ([Val], [Val])
value_and_histograms

    updateHistogram :: HistOp rep -> [Int] -> ([Val], [Val]) -> InterpM rep [Val]
updateHistogram HistOp rep
hist_operation [Int]
indices ([Val]
values, [Val]
histograms)
      | Bool -> Bool
not ([Int] -> [Val] -> Bool
inBounds [Int]
indices [Val]
histograms) =
          [Val] -> InterpM rep [Val]
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Val]
histograms
      | Bool
otherwise = do
          old_bins <- (Val -> InterpM rep Val) -> [Val] -> InterpM rep [Val]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Int] -> Val -> InterpM rep Val
forall rep. [Int] -> Val -> InterpM rep Val
readHistogramBin [Int]
indices) [Val]
histograms
          new_bins <-
            evalLambda funs env (histOp hist_operation) (old_bins <> values)
          zipWithM (writeHistogramBin indices) histograms new_bins

    inBounds :: [Int] -> [Val] -> Bool
inBounds [Int]
indices [Val]
histograms =
      case [Val]
histograms of
        [] -> Bool
False
        ArrayValue [Int]
shape ArrayValues
_ : [Val]
_ ->
          [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
indices Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
shape
            Bool -> Bool -> Bool
&& [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Int -> Int -> Bool) -> [Int] -> [Int] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Bool
forall {a}. (Ord a, Num a) => a -> a -> Bool
validIndex [Int]
indices [Int]
shape)
        [Val]
_ -> Bool
False

    validIndex :: a -> a -> Bool
validIndex a
index a
dimension =
      a
index a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
0 Bool -> Bool -> Bool
&& a
index a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
dimension

evalStream :: FunEnv rep -> Env -> SubExp -> [VName] -> [SubExp] -> Lambda rep -> InterpM rep [Val]
evalStream :: forall rep.
FunEnv rep
-> Env
-> SubExp
-> [VName]
-> [SubExp]
-> Lambda rep
-> InterpM rep [Val]
evalStream FunEnv rep
funs Env
env SubExp
width_exp [VName]
input_names [SubExp]
initial_accumulators Lambda rep
lambda = do
  width <- Env -> SubExp -> InterpM rep Val
forall rep. Env -> SubExp -> InterpM rep Val
evalSubExp Env
env SubExp
width_exp
  inputs <- mapM (evalVar env) input_names
  accumulators <- mapM (evalSubExp env) initial_accumulators
  evalLambda funs env lambda $ width : accumulators <> inputs

evalScrema :: FunEnv rep -> Env -> SubExp -> [VName] -> ScremaForm rep -> InterpM rep [Val]
evalScrema :: forall rep.
FunEnv rep
-> Env -> SubExp -> [VName] -> ScremaForm rep -> InterpM rep [Val]
evalScrema FunEnv rep
funs Env
env SubExp
width_exp [VName]
input_names (ScremaForm Lambda rep
pre_lambda [Scan rep]
scans [Reduce rep]
reductions Lambda rep
post_lambda) = do
  width <- Env -> SubExp -> InterpM rep Int
forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env SubExp
width_exp
  inputs <- mapM (evalVar env) input_names
  initial_scan_states <- mapM (mapM (evalSubExp env) . scanNeutral) scans
  initial_reduction_states <- mapM (mapM (evalSubExp env) . redNeutral) reductions
  (_, final_reduction_states, reversed_output_rows) <-
    foldM
      (runIteration inputs)
      (initial_scan_states, initial_reduction_states, [])
      [0 .. width - 1]
  collected_outputs <-
    collectScremaOutputs
      env
      width
      (lambdaReturnType post_lambda)
      (reverse reversed_output_rows)

  pure $ concat final_reduction_states <> collected_outputs
  where
    scan_sizes :: [Int]
scan_sizes =
      (Scan rep -> Int) -> [Scan rep] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([SubExp] -> Int) -> (Scan rep -> [SubExp]) -> Scan rep -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Scan rep -> [SubExp]
forall rep. Scan rep -> [SubExp]
scanNeutral) [Scan rep]
scans
    reduction_sizes :: [Int]
reduction_sizes =
      (Reduce rep -> Int) -> [Reduce rep] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([SubExp] -> Int) -> (Reduce rep -> [SubExp]) -> Reduce rep -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Reduce rep -> [SubExp]
forall rep. Reduce rep -> [SubExp]
redNeutral) [Reduce rep]
reductions

    runIteration :: [Val]
-> ([[Val]], [[Val]], [[Val]])
-> Int
-> InterpM rep ([[Val]], [[Val]], [[Val]])
runIteration
      [Val]
inputs
      ([[Val]]
scan_states, [[Val]]
reduction_states, [[Val]]
output_rows)
      Int
index = do
        input_rows <- (Val -> InterpM rep Val) -> [Val] -> InterpM rep [Val]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Int -> Val -> InterpM rep Val
forall rep. Int -> Val -> InterpM rep Val
rowAt Int
index) [Val]
inputs
        pre_results <- evalLambda funs env pre_lambda input_rows
        let (scan_contributions, after_scans) = splitGroups scan_sizes pre_results
            (reduction_contributions, map_values) = splitGroups reduction_sizes after_scans
        next_scan_states <- updateScanStates funs env scans scan_states scan_contributions
        next_reduction_states <- updateReductionStates funs env reductions reduction_states reduction_contributions
        post_results <- evalLambda funs env post_lambda (concat next_scan_states <> map_values)
        pure (next_scan_states, next_reduction_states, post_results : output_rows)

int64Val :: Int -> Val
int64Val :: Int -> Val
int64Val = PrimValue -> Val
PrimVal (PrimValue -> Val) -> (Int -> PrimValue) -> Int -> Val
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntValue -> PrimValue
IntValue (IntValue -> PrimValue) -> (Int -> IntValue) -> Int -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> IntValue
Int64Value (Int64 -> IntValue) -> (Int -> Int64) -> Int -> IntValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral

evalSegSpace :: FunEnv rep -> M.Map VName Val -> Seg.SegSpace -> Seg.KernelBody rep -> InterpM rep ([Int], [([Int], [KernelResultValue])])
evalSegSpace :: forall rep.
FunEnv rep
-> Env
-> SegSpace
-> KernelBody rep
-> InterpM rep ([Int], [([Int], [KernelResultValue])])
evalSegSpace FunEnv rep
funs Env
env space :: SegSpace
space@(Seg.SegSpace VName
_ [(VName, SubExp)]
dimensions) KernelBody rep
body = do
  sizes <- ((VName, SubExp) -> InterpM rep Int)
-> [(VName, SubExp)] -> InterpM rep [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Env -> SubExp -> InterpM rep Int
forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env (SubExp -> InterpM rep Int)
-> ((VName, SubExp) -> SubExp)
-> (VName, SubExp)
-> InterpM rep Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, SubExp) -> SubExp
forall a b. (a, b) -> b
snd) [(VName, SubExp)]
dimensions
  rows <- mapM (runWorker sizes) $ allCoordinates sizes
  pure (sizes, rows)
  where
    runWorker :: [Int] -> [Int] -> InterpM rep ([Int], [KernelResultValue])
runWorker [Int]
sizes [Int]
coordinate = do
      let worker_env :: Env
worker_env = Env -> SegSpace -> [Int] -> [Int] -> Env
segSpaceEnv Env
env SegSpace
space [Int]
sizes [Int]
coordinate
      values <- FunEnv rep
-> Env -> KernelBody rep -> InterpM rep [KernelResultValue]
forall rep.
FunEnv rep
-> Env -> KernelBody rep -> InterpM rep [KernelResultValue]
evalKernelBody FunEnv rep
funs Env
worker_env KernelBody rep
body
      pure (coordinate, values)

segSpaceEnv :: Env -> Seg.SegSpace -> [Int] -> [Int] -> Env
segSpaceEnv :: Env -> SegSpace -> [Int] -> [Int] -> Env
segSpaceEnv Env
env (Seg.SegSpace VName
flat [(VName, SubExp)]
dimensions) [Int]
shape [Int]
coordinate =
  Env -> Env -> Env
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union Env
bindings Env
env
  where
    bindings :: Env
bindings =
      [(VName, Val)] -> Env
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Val)] -> Env) -> [(VName, Val)] -> Env
forall a b. (a -> b) -> a -> b
$
        (VName
flat, Int -> Val
int64Val (Int -> Val) -> Int -> Val
forall a b. (a -> b) -> a -> b
$ [Int] -> [Int] -> Int
linearIndex [Int]
shape [Int]
coordinate)
          (VName, Val) -> [(VName, Val)] -> [(VName, Val)]
forall a. a -> [a] -> [a]
: [VName] -> [Val] -> [(VName, Val)]
forall a b. [a] -> [b] -> [(a, b)]
zip (((VName, SubExp) -> VName) -> [(VName, SubExp)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, SubExp) -> VName
forall a b. (a, b) -> a
fst [(VName, SubExp)]
dimensions) ((Int -> Val) -> [Int] -> [Val]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Val
int64Val [Int]
coordinate)

evalShape :: Env -> Shape -> InterpM rep [Int]
evalShape :: forall rep. Env -> Shape -> InterpM rep [Int]
evalShape Env
env = (SubExp -> InterpM rep Int) -> [SubExp] -> InterpM rep [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Env -> SubExp -> InterpM rep Int
forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env) ([SubExp] -> InterpM rep [Int])
-> (Shape -> [SubExp]) -> Shape -> InterpM rep [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims

expectKernelValue :: KernelResultValue -> InterpM rep Val
expectKernelValue :: forall rep. KernelResultValue -> InterpM rep Val
expectKernelValue (KernelValue Val
result_value) = Val -> InterpM rep Val
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Val
result_value
expectKernelValue KernelTile {} =
  [Char] -> InterpM rep Val
forall a. HasCallStack => [Char] -> a
error [Char]
"TileReturns cannot be used as a segmented contribution"
expectKernelValue KernelRegTile {} =
  [Char] -> InterpM rep Val
forall a. HasCallStack => [Char] -> a
error [Char]
"RegTileReturns cannot be used as a segmented contribution"

segmentRows :: Int -> Int -> [a] -> [[a]]
segmentRows :: forall a. Int -> Int -> [a] -> [[a]]
segmentRows Int
segment_count Int
segment_width =
  Int -> [a] -> [[a]]
forall {t} {a}. (Ord t, Num t) => t -> [a] -> [[a]]
go Int
segment_count
  where
    go :: t -> [a] -> [[a]]
go t
remaining [a]
rows
      | t
remaining t -> t -> Bool
forall a. Ord a => a -> a -> Bool
<= t
0 = []
      | Bool
otherwise =
          let ([a]
segment, [a]
rest) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
segment_width [a]
rows
           in [a]
segment [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: t -> [a] -> [[a]]
go (t
remaining t -> t -> t
forall a. Num a => a -> a -> a
- t
1) [a]
rest

splitSegContributions :: [Seg.SegBinOp rep] -> [Val] -> [[Val]]
splitSegContributions :: forall rep. [SegBinOp rep] -> [Val] -> [[Val]]
splitSegContributions [SegBinOp rep]
operators =
  ([[Val]], [Val]) -> [[Val]]
forall a b. (a, b) -> a
fst (([[Val]], [Val]) -> [[Val]])
-> ([Val] -> ([[Val]], [Val])) -> [Val] -> [[Val]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [Val] -> ([[Val]], [Val])
forall a. [Int] -> [a] -> ([[a]], [a])
splitGroups ((SegBinOp rep -> Int) -> [SegBinOp rep] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([SubExp] -> Int)
-> (SegBinOp rep -> [SubExp]) -> SegBinOp rep -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SegBinOp rep -> [SubExp]
forall rep. SegBinOp rep -> [SubExp]
Seg.segBinOpNeutral) [SegBinOp rep]
operators)

initialSegBinOp ::
  Env ->
  Seg.SegBinOp rep ->
  InterpM rep [Val]
initialSegBinOp :: forall rep. Env -> SegBinOp rep -> InterpM rep [Val]
initialSegBinOp Env
env SegBinOp rep
operator = do
  neutral <- (SubExp -> InterpM rep Val) -> [SubExp] -> InterpM rep [Val]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Env -> SubExp -> InterpM rep Val
forall rep. Env -> SubExp -> InterpM rep Val
evalSubExp Env
env) ([SubExp] -> InterpM rep [Val]) -> [SubExp] -> InterpM rep [Val]
forall a b. (a -> b) -> a -> b
$ SegBinOp rep -> [SubExp]
forall rep. SegBinOp rep -> [SubExp]
Seg.segBinOpNeutral SegBinOp rep
operator
  vector_shape <- evalShape env $ Seg.segBinOpShape operator

  if null vector_shape
    then pure neutral
    else
      collectOutputs
        env
        vector_shape
        (lambdaReturnType $ Seg.segBinOpLambda operator)
        (replicate (product vector_shape) neutral)

applySegBinOp ::
  FunEnv rep ->
  Env ->
  Seg.SegBinOp rep ->
  [Val] ->
  [Val] ->
  InterpM rep [Val]
applySegBinOp :: forall rep.
FunEnv rep
-> Env -> SegBinOp rep -> [Val] -> [Val] -> InterpM rep [Val]
applySegBinOp FunEnv rep
funs Env
env SegBinOp rep
operator [Val]
state [Val]
contribution = do
  vector_shape <- Env -> Shape -> InterpM rep [Int]
forall rep. Env -> Shape -> InterpM rep [Int]
evalShape Env
env (Shape -> InterpM rep [Int]) -> Shape -> InterpM rep [Int]
forall a b. (a -> b) -> a -> b
$ SegBinOp rep -> Shape
forall rep. SegBinOp rep -> Shape
Seg.segBinOpShape SegBinOp rep
operator
  let operator_lambda = SegBinOp rep -> Lambda rep
forall rep. SegBinOp rep -> Lambda rep
Seg.segBinOpLambda SegBinOp rep
operator

  if null vector_shape
    then evalLambda funs env operator_lambda $ state <> contribution
    else do
      result_rows <-
        forM (allCoordinates vector_shape) $ \[Int]
coordinate -> do
          state_elements <-
            (Val -> InterpM rep Val) -> [Val] -> InterpM rep [Val]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Int] -> Val -> InterpM rep Val
forall rep. [Int] -> Val -> InterpM rep Val
readAccumulatorElement [Int]
coordinate) [Val]
state
          contribution_elements <-
            mapM (readAccumulatorElement coordinate) contribution
          evalLambda
            funs
            env
            operator_lambda
            (state_elements <> contribution_elements)

      collectOutputs
        env
        vector_shape
        (lambdaReturnType operator_lambda)
        result_rows

initialSegStates ::
  Env ->
  [Seg.SegBinOp rep] ->
  InterpM rep [[Val]]
initialSegStates :: forall rep. Env -> [SegBinOp rep] -> InterpM rep [[Val]]
initialSegStates Env
env =
  (SegBinOp rep -> InterpM rep [Val])
-> [SegBinOp rep] -> InterpM rep [[Val]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((SegBinOp rep -> InterpM rep [Val])
 -> [SegBinOp rep] -> InterpM rep [[Val]])
-> (SegBinOp rep -> InterpM rep [Val])
-> [SegBinOp rep]
-> InterpM rep [[Val]]
forall a b. (a -> b) -> a -> b
$ Env -> SegBinOp rep -> InterpM rep [Val]
forall rep. Env -> SegBinOp rep -> InterpM rep [Val]
initialSegBinOp Env
env

updateSegStates ::
  FunEnv rep ->
  Env ->
  [Seg.SegBinOp rep] ->
  [[Val]] ->
  [[Val]] ->
  InterpM rep [[Val]]
updateSegStates :: forall rep.
FunEnv rep
-> Env
-> [SegBinOp rep]
-> [[Val]]
-> [[Val]]
-> InterpM rep [[Val]]
updateSegStates FunEnv rep
funs Env
env [SegBinOp rep]
operators [[Val]]
states [[Val]]
contributions =
  [(SegBinOp rep, [Val], [Val])]
-> ((SegBinOp rep, [Val], [Val]) -> InterpM rep [Val])
-> InterpM rep [[Val]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([SegBinOp rep]
-> [[Val]] -> [[Val]] -> [(SegBinOp rep, [Val], [Val])]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [SegBinOp rep]
operators [[Val]]
states [[Val]]
contributions) (((SegBinOp rep, [Val], [Val]) -> InterpM rep [Val])
 -> InterpM rep [[Val]])
-> ((SegBinOp rep, [Val], [Val]) -> InterpM rep [Val])
-> InterpM rep [[Val]]
forall a b. (a -> b) -> a -> b
$ \(SegBinOp rep
operator, [Val]
state, [Val]
contribution) ->
    FunEnv rep
-> Env -> SegBinOp rep -> [Val] -> [Val] -> InterpM rep [Val]
forall rep.
FunEnv rep
-> Env -> SegBinOp rep -> [Val] -> [Val] -> InterpM rep [Val]
applySegBinOp FunEnv rep
funs Env
env SegBinOp rep
operator [Val]
state [Val]
contribution

evalSegOp ::
  FunEnv rep -> Env -> Seg.SegOp level rep -> InterpM rep [Val]
evalSegOp :: forall rep level.
FunEnv rep -> Env -> SegOp level rep -> InterpM rep [Val]
evalSegOp FunEnv rep
funs Env
env (Seg.SegMap level
_ SegSpace
space [TypeBase Shape NoMode]
types KernelBody rep
body) =
  FunEnv rep
-> Env
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> InterpM rep [Val]
forall rep.
FunEnv rep
-> Env
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> InterpM rep [Val]
evalSegMap FunEnv rep
funs Env
env SegSpace
space [TypeBase Shape NoMode]
types KernelBody rep
body
evalSegOp FunEnv rep
funs Env
env (Seg.SegRed level
_ SegSpace
space [TypeBase Shape NoMode]
types KernelBody rep
body [SegBinOp rep]
operators) =
  FunEnv rep
-> Env
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> [SegBinOp rep]
-> InterpM rep [Val]
forall rep.
FunEnv rep
-> Env
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> [SegBinOp rep]
-> InterpM rep [Val]
evalSegRed FunEnv rep
funs Env
env SegSpace
space [TypeBase Shape NoMode]
types KernelBody rep
body [SegBinOp rep]
operators
evalSegOp FunEnv rep
funs Env
env (Seg.SegScan level
_ SegSpace
space [TypeBase Shape NoMode]
_ KernelBody rep
body [SegBinOp rep]
operators SegPostOp rep
post_operator) =
  FunEnv rep
-> Env
-> SegSpace
-> KernelBody rep
-> [SegBinOp rep]
-> SegPostOp rep
-> InterpM rep [Val]
forall rep.
FunEnv rep
-> Env
-> SegSpace
-> KernelBody rep
-> [SegBinOp rep]
-> SegPostOp rep
-> InterpM rep [Val]
evalSegScan FunEnv rep
funs Env
env SegSpace
space KernelBody rep
body [SegBinOp rep]
operators SegPostOp rep
post_operator
evalSegOp FunEnv rep
funs Env
env (Seg.SegHist level
_ SegSpace
space [TypeBase Shape NoMode]
_ KernelBody rep
body [HistOp rep]
operators) =
  FunEnv rep
-> Env
-> SegSpace
-> KernelBody rep
-> [HistOp rep]
-> InterpM rep [Val]
forall rep.
FunEnv rep
-> Env
-> SegSpace
-> KernelBody rep
-> [HistOp rep]
-> InterpM rep [Val]
evalSegHist FunEnv rep
funs Env
env SegSpace
space KernelBody rep
body [HistOp rep]
operators

evalSegMap ::
  FunEnv rep ->
  Env ->
  Seg.SegSpace ->
  [Type] ->
  Seg.KernelBody rep ->
  InterpM rep [Val]
evalSegMap :: forall rep.
FunEnv rep
-> Env
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> InterpM rep [Val]
evalSegMap FunEnv rep
funs Env
env SegSpace
space [TypeBase Shape NoMode]
types KernelBody rep
body = do
  (shape, indexed_rows) <- FunEnv rep
-> Env
-> SegSpace
-> KernelBody rep
-> InterpM rep ([Int], [([Int], [KernelResultValue])])
forall rep.
FunEnv rep
-> Env
-> SegSpace
-> KernelBody rep
-> InterpM rep ([Int], [([Int], [KernelResultValue])])
evalSegSpace FunEnv rep
funs Env
env SegSpace
space KernelBody rep
body
  collectSegOutputs env shape types (bodyResult body) indexed_rows

evalSegRed ::
  FunEnv rep ->
  Env ->
  Seg.SegSpace ->
  [Type] ->
  Seg.KernelBody rep ->
  [Seg.SegBinOp rep] ->
  InterpM rep [Val]
evalSegRed :: forall rep.
FunEnv rep
-> Env
-> SegSpace
-> [TypeBase Shape NoMode]
-> KernelBody rep
-> [SegBinOp rep]
-> InterpM rep [Val]
evalSegRed FunEnv rep
funs Env
env SegSpace
space [TypeBase Shape NoMode]
types KernelBody rep
body [SegBinOp rep]
operators = do
  (shape, indexed_rows) <- FunEnv rep
-> Env
-> SegSpace
-> KernelBody rep
-> InterpM rep ([Int], [([Int], [KernelResultValue])])
forall rep.
FunEnv rep
-> Env
-> SegSpace
-> KernelBody rep
-> InterpM rep ([Int], [([Int], [KernelResultValue])])
evalSegSpace FunEnv rep
funs Env
env SegSpace
space KernelBody rep
body

  let segment_shape = [Int] -> [Int]
forall a. HasCallStack => [a] -> [a]
init [Int]
shape
      segment_width = [Int] -> Int
forall a. HasCallStack => [a] -> a
last [Int]
shape
      segment_count = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
segment_shape
      reduction_count = [SegBinOp rep] -> Int
forall rep. [SegBinOp rep] -> Int
Seg.segBinOpResults [SegBinOp rep]
operators
      reduction_types = Int -> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a. Int -> [a] -> [a]
take Int
reduction_count [TypeBase Shape NoMode]
types
      map_types = Int -> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a. Int -> [a] -> [a]
drop Int
reduction_count [TypeBase Shape NoMode]
types
      reduction_result_rows =
        (([Int], [KernelResultValue]) -> [KernelResultValue])
-> [([Int], [KernelResultValue])] -> [[KernelResultValue]]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> [KernelResultValue] -> [KernelResultValue]
forall a. Int -> [a] -> [a]
take Int
reduction_count ([KernelResultValue] -> [KernelResultValue])
-> (([Int], [KernelResultValue]) -> [KernelResultValue])
-> ([Int], [KernelResultValue])
-> [KernelResultValue]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Int], [KernelResultValue]) -> [KernelResultValue]
forall a b. (a, b) -> b
snd) [([Int], [KernelResultValue])]
indexed_rows

  reduction_rows <-
    mapM (mapM expectKernelValue) reduction_result_rows

  let segments =
        Int -> Int -> [[Val]] -> [[[Val]]]
forall a. Int -> Int -> [a] -> [[a]]
segmentRows
          Int
segment_count
          Int
segment_width
          [[Val]]
reduction_rows

  reduced_rows <- mapM reduceSegment segments

  reduction_outputs <-
    collectOutputs
      env
      segment_shape
      reduction_types
      reduced_rows

  map_outputs <-
    collectSegOutputs
      env
      shape
      map_types
      (drop reduction_count $ bodyResult body)
      (map (fmap $ drop reduction_count) indexed_rows)

  pure $ reduction_outputs <> map_outputs
  where
    reduceSegment :: t [Val] -> InterpM rep [Val]
reduceSegment t [Val]
rows = do
      initial_states <- Env -> [SegBinOp rep] -> InterpM rep [[Val]]
forall rep. Env -> [SegBinOp rep] -> InterpM rep [[Val]]
initialSegStates Env
env [SegBinOp rep]
operators
      final_states <- foldM reduceRow initial_states rows
      pure $ concat final_states

    reduceRow :: [[Val]] -> [Val] -> InterpM rep [[Val]]
reduceRow [[Val]]
states [Val]
values =
      FunEnv rep
-> Env
-> [SegBinOp rep]
-> [[Val]]
-> [[Val]]
-> InterpM rep [[Val]]
forall rep.
FunEnv rep
-> Env
-> [SegBinOp rep]
-> [[Val]]
-> [[Val]]
-> InterpM rep [[Val]]
updateSegStates FunEnv rep
funs Env
env [SegBinOp rep]
operators [[Val]]
states ([[Val]] -> InterpM rep [[Val]]) -> [[Val]] -> InterpM rep [[Val]]
forall a b. (a -> b) -> a -> b
$
        [SegBinOp rep] -> [Val] -> [[Val]]
forall rep. [SegBinOp rep] -> [Val] -> [[Val]]
splitSegContributions [SegBinOp rep]
operators [Val]
values

evalSegScan ::
  FunEnv rep ->
  Env ->
  Seg.SegSpace ->
  Seg.KernelBody rep ->
  [Seg.SegBinOp rep] ->
  Seg.SegPostOp rep ->
  InterpM rep [Val]
evalSegScan :: forall rep.
FunEnv rep
-> Env
-> SegSpace
-> KernelBody rep
-> [SegBinOp rep]
-> SegPostOp rep
-> InterpM rep [Val]
evalSegScan FunEnv rep
funs Env
env SegSpace
space KernelBody rep
body [SegBinOp rep]
operators SegPostOp rep
post_operator = do
  (shape, indexed_rows) <- FunEnv rep
-> Env
-> SegSpace
-> KernelBody rep
-> InterpM rep ([Int], [([Int], [KernelResultValue])])
forall rep.
FunEnv rep
-> Env
-> SegSpace
-> KernelBody rep
-> InterpM rep ([Int], [([Int], [KernelResultValue])])
evalSegSpace FunEnv rep
funs Env
env SegSpace
space KernelBody rep
body
  worker_rows <- mapM evaluateWorker indexed_rows

  let segment_shape = [Int] -> [Int]
forall a. HasCallStack => [a] -> [a]
init [Int]
shape
      segment_width = [Int] -> Int
forall a. HasCallStack => [a] -> a
last [Int]
shape
      segment_count = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
segment_shape
      segments =
        Int -> Int -> [([Int], [Val])] -> [[([Int], [Val])]]
forall a. Int -> Int -> [a] -> [[a]]
segmentRows
          Int
segment_count
          Int
segment_width
          [([Int], [Val])]
worker_rows

  output_rows <- concat <$> mapM (scanSegment shape) segments

  collectOutputs
    env
    shape
    (lambdaReturnType post_lambda)
    output_rows
  where
    evaluateWorker :: (a, t KernelResultValue) -> InterpM rep (a, t Val)
evaluateWorker (a
coordinate, t KernelResultValue
values) = do
      values' <- (KernelResultValue -> InterpM rep Val)
-> t KernelResultValue -> InterpM rep (t Val)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> t a -> m (t b)
mapM KernelResultValue -> InterpM rep Val
forall rep. KernelResultValue -> InterpM rep Val
expectKernelValue t KernelResultValue
values
      pure (coordinate, values')

    contribution_count :: Int
contribution_count = [SegBinOp rep] -> Int
forall rep. [SegBinOp rep] -> Int
Seg.segBinOpResults [SegBinOp rep]
operators
    post_lambda :: GLambda rep (TypeBase Shape NoMode)
post_lambda =
      SegPostOp rep -> GLambda rep (TypeBase Shape NoMode)
forall rep. SegPostOp rep -> Lambda rep
Seg.segPostOpLambda SegPostOp rep
post_operator
    scanSegment :: [Int] -> t ([Int], [Val]) -> InterpM rep [[Val]]
scanSegment [Int]
shape t ([Int], [Val])
rows = do
      initial_states <- Env -> [SegBinOp rep] -> InterpM rep [[Val]]
forall rep. Env -> [SegBinOp rep] -> InterpM rep [[Val]]
initialSegStates Env
env [SegBinOp rep]
operators
      (_, reversed_outputs) <-
        foldM
          (scanRow shape)
          (initial_states, [])
          rows
      pure $ reverse reversed_outputs

    scanRow :: [Int]
-> ([[Val]], [[Val]])
-> ([Int], [Val])
-> InterpM rep ([[Val]], [[Val]])
scanRow [Int]
shape ([[Val]]
states, [[Val]]
outputs) ([Int]
coordinate, [Val]
values) = do
      let ([Val]
contribution_values, [Val]
map_values) =
            Int -> [Val] -> ([Val], [Val])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
contribution_count [Val]
values

      next_states <-
        FunEnv rep
-> Env
-> [SegBinOp rep]
-> [[Val]]
-> [[Val]]
-> InterpM rep [[Val]]
forall rep.
FunEnv rep
-> Env
-> [SegBinOp rep]
-> [[Val]]
-> [[Val]]
-> InterpM rep [[Val]]
updateSegStates
          FunEnv rep
funs
          Env
env
          [SegBinOp rep]
operators
          [[Val]]
states
          ([SegBinOp rep] -> [Val] -> [[Val]]
forall rep. [SegBinOp rep] -> [Val] -> [[Val]]
splitSegContributions [SegBinOp rep]
operators [Val]
contribution_values)

      post_results <-
        evalLambda
          funs
          (segSpaceEnv env space shape coordinate)
          post_lambda
          (concat next_states <> map_values)

      pure (next_states, post_results : outputs)

evalSegHist ::
  FunEnv rep ->
  Env ->
  Seg.SegSpace ->
  Seg.KernelBody rep ->
  [Seg.HistOp rep] ->
  InterpM rep [Val]
evalSegHist :: forall rep.
FunEnv rep
-> Env
-> SegSpace
-> KernelBody rep
-> [HistOp rep]
-> InterpM rep [Val]
evalSegHist FunEnv rep
funs Env
env SegSpace
space KernelBody rep
body [HistOp rep]
operators = do
  (_, indexed_rows) <- FunEnv rep
-> Env
-> SegSpace
-> KernelBody rep
-> InterpM rep ([Int], [([Int], [KernelResultValue])])
forall rep.
FunEnv rep
-> Env
-> SegSpace
-> KernelBody rep
-> InterpM rep ([Int], [([Int], [KernelResultValue])])
evalSegSpace FunEnv rep
funs Env
env SegSpace
space KernelBody rep
body
  initial_histograms <- mapM (mapM (evalVar env) . Seg.histDest) operators
  final_histograms <- foldM updateFromWorker initial_histograms indexed_rows
  pure $ concat final_histograms
  where
    index_counts :: [Int]
index_counts = (HistOp rep -> Int) -> [HistOp rep] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank (Shape -> Int) -> (HistOp rep -> Shape) -> HistOp rep -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HistOp rep -> Shape
forall rep. HistOp rep -> Shape
Seg.histShape) [HistOp rep]
operators
    value_counts :: [Int]
value_counts = (HistOp rep -> Int) -> [HistOp rep] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ([VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([VName] -> Int) -> (HistOp rep -> [VName]) -> HistOp rep -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HistOp rep -> [VName]
forall rep. HistOp rep -> [VName]
Seg.histDest) [HistOp rep]
operators

    updateFromWorker :: [[Val]] -> ([Int], [KernelResultValue]) -> InterpM rep [[Val]]
updateFromWorker [[Val]]
histograms ([Int]
coordinate, [KernelResultValue]
worker_values) = do
      let segment_indices :: [Int]
segment_indices = [Int] -> [Int]
forall a. HasCallStack => [a] -> [a]
init [Int]
coordinate

      worker_values' <- (KernelResultValue -> InterpM rep Val)
-> [KernelResultValue] -> InterpM rep [Val]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM KernelResultValue -> InterpM rep Val
forall rep. KernelResultValue -> InterpM rep Val
expectKernelValue [KernelResultValue]
worker_values

      let (index_groups, remaining) = splitGroups index_counts worker_values'
          (value_groups, _) = splitGroups value_counts remaining

      evaluated_indices <-
        mapM (mapM (expectPrimVal >=> expectInt)) index_groups

      forM (zip3 operators evaluated_indices (zip value_groups histograms)) $
        \(HistOp rep
operator, [Int]
bucket_indices, ([Val], [Val])
values_and_histograms) ->
          [Int] -> HistOp rep -> [Int] -> ([Val], [Val]) -> InterpM rep [Val]
updateHistogram
            [Int]
segment_indices
            HistOp rep
operator
            [Int]
bucket_indices
            ([Val], [Val])
values_and_histograms

    updateHistogram :: [Int] -> HistOp rep -> [Int] -> ([Val], [Val]) -> InterpM rep [Val]
updateHistogram
      [Int]
segment_indices
      HistOp rep
operator
      [Int]
bucket_indices
      ([Val]
new_values, [Val]
histograms)
        | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [Int] -> [Val] -> Bool
forall {t :: * -> *}. Foldable t => [Int] -> t Val -> Bool
indicesInBounds [Int]
full_indices [Val]
histograms =
            [Val] -> InterpM rep [Val]
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Val]
histograms
        | Bool
otherwise = do
            vector_shape <- Env -> Shape -> InterpM rep [Int]
forall rep. Env -> Shape -> InterpM rep [Int]
evalShape Env
env (Shape -> InterpM rep [Int]) -> Shape -> InterpM rep [Int]
forall a b. (a -> b) -> a -> b
$ HistOp rep -> Shape
forall rep. HistOp rep -> Shape
Seg.histOpShape HistOp rep
operator
            old_values <- mapM (readHistogramBin full_indices) histograms

            replacement_values <-
              if null vector_shape
                then
                  evalLambda
                    funs
                    env
                    (Seg.histOp operator)
                    (old_values <> new_values)
                else do
                  result_rows <-
                    forM (allCoordinates vector_shape) $ \[Int]
coordinate -> do
                      old_elements <-
                        (Val -> InterpM rep Val) -> [Val] -> InterpM rep [Val]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Int] -> Val -> InterpM rep Val
forall rep. [Int] -> Val -> InterpM rep Val
readAccumulatorElement [Int]
coordinate) [Val]
old_values
                      new_elements <-
                        mapM (readAccumulatorElement coordinate) new_values
                      evalLambda
                        funs
                        env
                        (Seg.histOp operator)
                        (old_elements <> new_elements)

                  collectOutputs
                    env
                    vector_shape
                    (lambdaReturnType $ Seg.histOp operator)
                    result_rows

            zipWithM
              (writeHistogramBin full_indices)
              histograms
              replacement_values
        where
          full_indices :: [Int]
full_indices =
            [Int]
segment_indices [Int] -> [Int] -> [Int]
forall a. Semigroup a => a -> a -> a
<> [Int]
bucket_indices

    indicesInBounds :: [Int] -> t Val -> Bool
indicesInBounds [Int]
indices =
      (Val -> Bool) -> t Val -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((Val -> Bool) -> t Val -> Bool) -> (Val -> Bool) -> t Val -> Bool
forall a b. (a -> b) -> a -> b
$ [Int] -> Val -> Bool
histogramIndicesInBounds [Int]
indices

    histogramIndicesInBounds :: [Int] -> Val -> Bool
histogramIndicesInBounds
      [Int]
indices
      (ArrayValue [Int]
histogram_shape ArrayValues
_) =
        [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
indices Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
histogram_shape
          Bool -> Bool -> Bool
&& [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and
            ((Int -> Int -> Bool) -> [Int] -> [Int] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Bool
forall {a}. (Ord a, Num a) => a -> a -> Bool
validIndex [Int]
indices [Int]
histogram_shape)
    histogramIndicesInBounds [Int]
_ Val
_ =
      Bool
False

    validIndex :: a -> a -> Bool
validIndex a
index a
dimension =
      a
index a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
0 Bool -> Bool -> Bool
&& a
index a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
dimension

splitGroups :: [Int] -> [a] -> ([[a]], [a])
splitGroups :: forall a. [Int] -> [a] -> ([[a]], [a])
splitGroups [] [a]
values = ([], [a]
values)
splitGroups (Int
size : [Int]
sizes) [a]
values =
  let ([a]
group, [a]
rest) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
size [a]
values
      ([[a]]
groups, [a]
remaining) = [Int] -> [a] -> ([[a]], [a])
forall a. [Int] -> [a] -> ([[a]], [a])
splitGroups [Int]
sizes [a]
rest
   in ([a]
group [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: [[a]]
groups, [a]
remaining)

evalLambda ::
  FunEnv rep ->
  Env ->
  GLambda rep return_type ->
  [Val] ->
  InterpM rep [Val]
evalLambda :: forall rep return_type.
FunEnv rep
-> Env -> GLambda rep return_type -> [Val] -> InterpM rep [Val]
evalLambda FunEnv rep
funs Env
env (Lambda [LParam rep]
params [return_type]
_ Body rep
body) [Val]
args =
  FunEnv rep -> Env -> Body rep -> InterpM rep [Val]
forall rep. FunEnv rep -> Env -> Body rep -> InterpM rep [Val]
evalBody FunEnv rep
funs (Env -> Env -> Env
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union ([(VName, Val)] -> Env
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Val)] -> Env) -> [(VName, Val)] -> Env
forall a b. (a -> b) -> a -> b
$ [VName] -> [Val] -> [(VName, Val)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((LParam rep -> VName) -> [LParam rep] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map LParam rep -> VName
forall dec. Param dec -> VName
paramName [LParam rep]
params) [Val]
args) Env
env) Body rep
body

updateScanStates ::
  FunEnv rep ->
  Env ->
  [Scan rep] ->
  [[Val]] ->
  [[Val]] ->
  InterpM rep [[Val]]
updateScanStates :: forall rep.
FunEnv rep
-> Env -> [Scan rep] -> [[Val]] -> [[Val]] -> InterpM rep [[Val]]
updateScanStates FunEnv rep
funs Env
env [Scan rep]
scans [[Val]]
states [[Val]]
contributions =
  (Scan rep -> ([Val], [Val]) -> InterpM rep [Val])
-> [Scan rep] -> [([Val], [Val])] -> InterpM rep [[Val]]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM Scan rep -> ([Val], [Val]) -> InterpM rep [Val]
updateOne [Scan rep]
scans ([[Val]] -> [[Val]] -> [([Val], [Val])]
forall a b. [a] -> [b] -> [(a, b)]
zip [[Val]]
states [[Val]]
contributions)
  where
    updateOne :: Scan rep -> ([Val], [Val]) -> InterpM rep [Val]
updateOne Scan rep
scan ([Val]
state, [Val]
contribution) =
      FunEnv rep
-> Env
-> GLambda rep (TypeBase Shape NoMode)
-> [Val]
-> InterpM rep [Val]
forall rep return_type.
FunEnv rep
-> Env -> GLambda rep return_type -> [Val] -> InterpM rep [Val]
evalLambda FunEnv rep
funs Env
env (Scan rep -> GLambda rep (TypeBase Shape NoMode)
forall rep. Scan rep -> Lambda rep
scanLambda Scan rep
scan) ([Val]
state [Val] -> [Val] -> [Val]
forall a. Semigroup a => a -> a -> a
<> [Val]
contribution)

updateReductionStates ::
  FunEnv rep ->
  Env ->
  [Reduce rep] ->
  [[Val]] ->
  [[Val]] ->
  InterpM rep [[Val]]
updateReductionStates :: forall rep.
FunEnv rep
-> Env -> [Reduce rep] -> [[Val]] -> [[Val]] -> InterpM rep [[Val]]
updateReductionStates FunEnv rep
funs Env
env [Reduce rep]
reductions [[Val]]
states [[Val]]
contributions =
  (Reduce rep -> ([Val], [Val]) -> InterpM rep [Val])
-> [Reduce rep] -> [([Val], [Val])] -> InterpM rep [[Val]]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM Reduce rep -> ([Val], [Val]) -> InterpM rep [Val]
updateOne [Reduce rep]
reductions ([[Val]] -> [[Val]] -> [([Val], [Val])]
forall a b. [a] -> [b] -> [(a, b)]
zip [[Val]]
states [[Val]]
contributions)
  where
    updateOne :: Reduce rep -> ([Val], [Val]) -> InterpM rep [Val]
updateOne Reduce rep
reduction ([Val]
state, [Val]
contribution) =
      FunEnv rep
-> Env
-> GLambda rep (TypeBase Shape NoMode)
-> [Val]
-> InterpM rep [Val]
forall rep return_type.
FunEnv rep
-> Env -> GLambda rep return_type -> [Val] -> InterpM rep [Val]
evalLambda FunEnv rep
funs Env
env (Reduce rep -> GLambda rep (TypeBase Shape NoMode)
forall rep. Reduce rep -> Lambda rep
redLambda Reduce rep
reduction) ([Val]
state [Val] -> [Val] -> [Val]
forall a. Semigroup a => a -> a -> a
<> [Val]
contribution)

rowAt :: Int -> Val -> InterpM rep Val
rowAt :: forall rep. Int -> Val -> InterpM rep Val
rowAt Int
index (ArrayValue (Int
_ : [Int]
row_shape) ArrayValues
values)
  | [Int] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
row_shape =
      PrimValue -> Val
PrimVal (PrimValue -> Val) -> InterpM rep PrimValue -> InterpM rep Val
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ArrayValues -> Int -> InterpM rep PrimValue
forall rep. ArrayValues -> Int -> InterpM rep PrimValue
readArrayValue ArrayValues
values Int
index
  | Bool
otherwise = do
      let row_size :: Int
row_size = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
row_shape
          offset :: Int
offset = Int
index Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
row_size
      row_values <-
        (Int -> InterpM rep PrimValue) -> [Int] -> InterpM rep [PrimValue]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (ArrayValues -> Int -> InterpM rep PrimValue
forall rep. ArrayValues -> Int -> InterpM rep PrimValue
readArrayValue ArrayValues
values) [Int
offset .. Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
row_size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
      newArrayValue row_shape (arrayValuesType values) row_values
rowAt Int
_ accumulator :: Val
accumulator@AccValue {} =
  Val -> InterpM rep Val
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Val
accumulator
rowAt Int
_ Val
_ =
  [Char] -> InterpM rep Val
forall a. HasCallStack => [Char] -> a
error [Char]
"cannot extract a row from this value"

readAccumulatorElement :: [Int] -> Val -> InterpM rep Val
readAccumulatorElement :: forall rep. [Int] -> Val -> InterpM rep Val
readAccumulatorElement [Int]
indices (ArrayValue [Int]
shape ArrayValues
values) = do
  let index_rank :: Int
index_rank = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
indices
      element_shape :: [Int]
element_shape = Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop Int
index_rank [Int]
shape
      element_size :: Int
element_size = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
element_shape
      offset :: Int
offset = [Int] -> [Int] -> Int
linearIndex (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
index_rank [Int]
shape) [Int]
indices Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
element_size
  case [Int]
element_shape of
    [] -> PrimValue -> Val
PrimVal (PrimValue -> Val) -> InterpM rep PrimValue -> InterpM rep Val
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ArrayValues -> Int -> InterpM rep PrimValue
forall rep. ArrayValues -> Int -> InterpM rep PrimValue
readArrayValue ArrayValues
values Int
offset
    [Int]
_ -> do
      element_values <- (Int -> InterpM rep PrimValue) -> [Int] -> InterpM rep [PrimValue]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (ArrayValues -> Int -> InterpM rep PrimValue
forall rep. ArrayValues -> Int -> InterpM rep PrimValue
readArrayValue ArrayValues
values) [Int
offset .. Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
element_size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
      newArrayValue element_shape (arrayValuesType values) element_values
readAccumulatorElement [Int]
_ Val
_ =
  [Char] -> InterpM rep Val
forall a. HasCallStack => [Char] -> a
error [Char]
"accumulator backing value must be an array"

writeAccumulatorElementInPlace :: [Int] -> Val -> Val -> InterpM rep ()
writeAccumulatorElementInPlace :: forall rep. [Int] -> Val -> Val -> InterpM rep ()
writeAccumulatorElementInPlace [Int]
indices (ArrayValue [Int]
shape ArrayValues
values) Val
replacement = do
  let index_rank :: Int
index_rank = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
indices
      element_size :: Int
element_size = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop Int
index_rank [Int]
shape
      offset :: Int
offset = [Int] -> [Int] -> Int
linearIndex (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
index_rank [Int]
shape) [Int]
indices Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
element_size
  replacement_values <- Val -> InterpM rep [PrimValue]
forall rep. Val -> InterpM rep [PrimValue]
valElements Val
replacement
  zipWithM_ (writeArrayValue values) [offset ..] replacement_values
writeAccumulatorElementInPlace [Int]
_ Val
_ Val
_ =
  [Char] -> InterpM rep ()
forall a. HasCallStack => [Char] -> a
error [Char]
"accumulator backing value must be an array"

readHistogramBin :: [Int] -> Val -> InterpM rep Val
readHistogramBin :: forall rep. [Int] -> Val -> InterpM rep Val
readHistogramBin [Int]
indices (ArrayValue [Int]
shape ArrayValues
values) = do
  let rank :: Int
rank = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
indices
      bin_shape :: [Int]
bin_shape = Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop Int
rank [Int]
shape
      bin_size :: Int
bin_size = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
bin_shape
      offset :: Int
offset = [Int] -> [Int] -> Int
linearIndex (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
rank [Int]
shape) [Int]
indices Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
bin_size
  case [Int]
bin_shape of
    [] -> PrimValue -> Val
PrimVal (PrimValue -> Val) -> InterpM rep PrimValue -> InterpM rep Val
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ArrayValues -> Int -> InterpM rep PrimValue
forall rep. ArrayValues -> Int -> InterpM rep PrimValue
readArrayValue ArrayValues
values Int
offset
    [Int]
_ -> do
      bin_values <- (Int -> InterpM rep PrimValue) -> [Int] -> InterpM rep [PrimValue]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (ArrayValues -> Int -> InterpM rep PrimValue
forall rep. ArrayValues -> Int -> InterpM rep PrimValue
readArrayValue ArrayValues
values) [Int
offset .. Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bin_size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
      newArrayValue bin_shape (arrayValuesType values) bin_values
readHistogramBin [Int]
_ Val
_ =
  [Char] -> InterpM rep Val
forall a. HasCallStack => [Char] -> a
error [Char]
"Hist destination must be an array"

writeHistogramBin :: [Int] -> Val -> Val -> InterpM rep Val
writeHistogramBin :: forall rep. [Int] -> Val -> Val -> InterpM rep Val
writeHistogramBin [Int]
indices histogram :: Val
histogram@(ArrayValue [Int]
shape ArrayValues
values) Val
replacement = do
  let rank :: Int
rank = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
indices
      bin_size :: Int
bin_size = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop Int
rank [Int]
shape
      offset :: Int
offset = [Int] -> [Int] -> Int
linearIndex (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
rank [Int]
shape) [Int]
indices Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
bin_size
  replacement_values <- Val -> InterpM rep [PrimValue]
forall rep. Val -> InterpM rep [PrimValue]
valElements Val
replacement
  zipWithM_ (writeArrayValue values) [offset ..] replacement_values
  pure histogram
writeHistogramBin [Int]
_ Val
_ Val
_ =
  [Char] -> InterpM rep Val
forall a. HasCallStack => [Char] -> a
error [Char]
"Hist destination must be an array"

collectFlatMapOutput ::
  Env ->
  Int ->
  ExtType ->
  [Val] ->
  InterpM rep Val
collectFlatMapOutput :: forall rep. Env -> Int -> ExtType -> [Val] -> InterpM rep Val
collectFlatMapOutput Env
env Int
total_size ExtType
result_type [Val]
rows
  | ExtType -> Bool
flatMapNonuniform ExtType
result_type =
      case [Val]
rows of
        ArrayValue (Int
_ : [Int]
row_shape) ArrayValues
values : [Val]
_ -> do
          element_values <- [[PrimValue]] -> [PrimValue]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[PrimValue]] -> [PrimValue])
-> InterpM rep [[PrimValue]] -> InterpM rep [PrimValue]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Val -> InterpM rep [PrimValue])
-> [Val] -> InterpM rep [[PrimValue]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Val -> InterpM rep [PrimValue]
forall rep. Val -> InterpM rep [PrimValue]
valElements [Val]
rows
          newArrayValue (total_size : row_shape) (arrayValuesType values) element_values
        [Val]
_ -> do
          (element_type, row_shape) <- Bool -> InterpM rep (PrimType, [Int])
forall {rep}. Bool -> InterpM rep (PrimType, [Int])
emptyArrayType Bool
True
          newArrayValue (total_size : row_shape) element_type []
  | Bool
otherwise =
      case ExtType
result_type of
        Prim PrimType
element_type ->
          [Int] -> PrimType -> [PrimValue] -> InterpM rep Val
forall rep. [Int] -> PrimType -> [PrimValue] -> InterpM rep Val
newArrayValue [[Val] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Val]
rows] PrimType
element_type ([PrimValue] -> InterpM rep Val)
-> InterpM rep [PrimValue] -> InterpM rep Val
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Val -> InterpM rep PrimValue) -> [Val] -> InterpM rep [PrimValue]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Val -> InterpM rep PrimValue
forall rep. Val -> InterpM rep PrimValue
expectPrimVal [Val]
rows
        Array PrimType
element_type ExtShape
_ NoMode
_ ->
          case [Val]
rows of
            ArrayValue [Int]
row_shape ArrayValues
_ : [Val]
_ -> do
              element_values <- [[PrimValue]] -> [PrimValue]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[PrimValue]] -> [PrimValue])
-> InterpM rep [[PrimValue]] -> InterpM rep [PrimValue]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Val -> InterpM rep [PrimValue])
-> [Val] -> InterpM rep [[PrimValue]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Val -> InterpM rep [PrimValue]
forall rep. Val -> InterpM rep [PrimValue]
valElements [Val]
rows
              newArrayValue (length rows : row_shape) element_type element_values
            [Val]
_ -> do
              (_, row_shape) <- Bool -> InterpM rep (PrimType, [Int])
forall {rep}. Bool -> InterpM rep (PrimType, [Int])
emptyArrayType Bool
False
              newArrayValue (0 : row_shape) element_type []
        Acc {} ->
          [Char] -> InterpM rep Val
forall a. HasCallStack => [Char] -> a
error [Char]
"FlatMap accumulator outputs are unsupported"
        Mem {} ->
          [Char] -> InterpM rep Val
forall a. HasCallStack => [Char] -> a
error [Char]
"FlatMap memory outputs are unsupported"
  where
    emptyArrayType :: Bool -> InterpM rep (PrimType, [Int])
emptyArrayType Bool
drop_existential =
      case ExtType
result_type of
        Array PrimType
element_type (Shape [ExtSize]
dimensions) NoMode
_ -> do
          let dimensions' :: [ExtSize]
dimensions'
                | Bool
drop_existential = Int -> [ExtSize] -> [ExtSize]
forall a. Int -> [a] -> [a]
drop Int
1 [ExtSize]
dimensions
                | Bool
otherwise = [ExtSize]
dimensions
          shape <- (ExtSize -> InterpM rep Int) -> [ExtSize] -> InterpM rep [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ExtSize -> InterpM rep Int
forall {rep}. ExtSize -> InterpM rep Int
evalFreeDimension [ExtSize]
dimensions'
          pure (element_type, shape)
        ExtType
_ ->
          [Char] -> InterpM rep (PrimType, [Int])
forall a. HasCallStack => [Char] -> a
error [Char]
"nonuniform FlatMap result must be an array"

    evalFreeDimension :: ExtSize -> InterpM rep Int
evalFreeDimension (Free SubExp
dimension) =
      Env -> SubExp -> InterpM rep Int
forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env SubExp
dimension
    evalFreeDimension (Ext Int
_) =
      [Char] -> InterpM rep Int
forall a. HasCallStack => [Char] -> a
error [Char]
"unexpected existential FlatMap result dimension"

collectOutputs ::
  Env ->
  [Int] ->
  [Type] ->
  [[Val]] ->
  InterpM rep [Val]
collectOutputs :: forall rep.
Env
-> [Int] -> [TypeBase Shape NoMode] -> [[Val]] -> InterpM rep [Val]
collectOutputs Env
env [Int]
outer_shape [TypeBase Shape NoMode]
return_types [[Val]]
iteration_results =
  (TypeBase Shape NoMode -> [Val] -> InterpM rep Val)
-> [TypeBase Shape NoMode] -> [[Val]] -> InterpM rep [Val]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM TypeBase Shape NoMode -> [Val] -> InterpM rep Val
forall {o} {rep}. TypeBase Shape o -> [Val] -> InterpM rep Val
collectOne [TypeBase Shape NoMode]
return_types [[Val]]
columns
  where
    columns :: [[Val]]
columns
      | [[Val]] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[Val]]
iteration_results =
          Int -> [Val] -> [[Val]]
forall a. Int -> a -> [a]
replicate ([TypeBase Shape NoMode] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase Shape NoMode]
return_types) []
      | Bool
otherwise =
          [[Val]] -> [[Val]]
forall a. [[a]] -> [[a]]
L.transpose [[Val]]
iteration_results

    collectOne :: TypeBase Shape o -> [Val] -> InterpM rep Val
collectOne (Prim PrimType
element_type) [Val]
rows = do
      values <- (Val -> InterpM rep PrimValue) -> [Val] -> InterpM rep [PrimValue]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Val -> InterpM rep PrimValue
forall rep. Val -> InterpM rep PrimValue
expectPrimVal [Val]
rows
      case (outer_shape, values) of
        ([], [PrimValue
val]) -> Val -> InterpM rep Val
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Val -> InterpM rep Val) -> Val -> InterpM rep Val
forall a b. (a -> b) -> a -> b
$ PrimValue -> Val
PrimVal PrimValue
val
        ([], [PrimValue]
_) -> [Char] -> InterpM rep Val
forall a. HasCallStack => [Char] -> a
error [Char]
"invalid scalar output count"
        ([Int], [PrimValue])
_ -> [Int] -> PrimType -> [PrimValue] -> InterpM rep Val
forall rep. [Int] -> PrimType -> [PrimValue] -> InterpM rep Val
newArrayValue [Int]
outer_shape PrimType
element_type [PrimValue]
values
    collectOne (Array PrimType
element_type Shape
annotated_shape o
_) [] = do
      row_shape <- Env -> Shape -> InterpM rep [Int]
forall rep. Env -> Shape -> InterpM rep [Int]
evalShape Env
env Shape
annotated_shape
      newArrayValue (outer_shape <> row_shape) element_type []
    collectOne (Array PrimType
element_type Shape
_ o
_) rows :: [Val]
rows@(ArrayValue [Int]
row_shape ArrayValues
_ : [Val]
_) = do
      element_values <- [[PrimValue]] -> [PrimValue]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[PrimValue]] -> [PrimValue])
-> InterpM rep [[PrimValue]] -> InterpM rep [PrimValue]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Val -> InterpM rep [PrimValue])
-> [Val] -> InterpM rep [[PrimValue]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Val -> InterpM rep [PrimValue]
forall rep. Val -> InterpM rep [PrimValue]
valElements [Val]
rows
      newArrayValue (outer_shape <> row_shape) element_type element_values
    collectOne Array {} [Val]
_ =
      [Char] -> InterpM rep Val
forall a. HasCallStack => [Char] -> a
error [Char]
"expected array output"
    collectOne (Acc VName
certificate Shape
_ [TypeBase Shape NoMode]
_) [] =
      Val -> InterpM rep Val
forall {f :: * -> *}. Applicative f => Val -> f Val
expectAccValue (Val -> InterpM rep Val) -> InterpM rep Val -> InterpM rep Val
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Env -> VName -> InterpM rep Val
forall rep. Env -> VName -> InterpM rep Val
evalVar Env
env VName
certificate
    collectOne Acc {} (Val
row : [Val]
_) =
      Val -> InterpM rep Val
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Val
row
    collectOne Mem {} [Val]
_ =
      [Char] -> InterpM rep Val
forall a. HasCallStack => [Char] -> a
error [Char]
"memory outputs are unsupported"

    expectAccValue :: Val -> f Val
expectAccValue accumulator :: Val
accumulator@AccValue {} = Val -> f Val
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Val
accumulator
    expectAccValue Val
_ = [Char] -> f Val
forall a. HasCallStack => [Char] -> a
error [Char]
"accumulator output has no backing storage"

collectScremaOutputs :: Env -> Int -> [Type] -> [[Val]] -> InterpM rep [Val]
collectScremaOutputs :: forall rep.
Env
-> Int -> [TypeBase Shape NoMode] -> [[Val]] -> InterpM rep [Val]
collectScremaOutputs Env
env Int
width = Env
-> [Int] -> [TypeBase Shape NoMode] -> [[Val]] -> InterpM rep [Val]
forall rep.
Env
-> [Int] -> [TypeBase Shape NoMode] -> [[Val]] -> InterpM rep [Val]
collectOutputs Env
env [Int
width]

collectSegOutputs ::
  Env ->
  [Int] ->
  [Type] ->
  [Seg.KernelResult] ->
  [([Int], [KernelResultValue])] ->
  InterpM rep [Val]
collectSegOutputs :: forall rep.
Env
-> [Int]
-> [TypeBase Shape NoMode]
-> [KernelResult]
-> [([Int], [KernelResultValue])]
-> InterpM rep [Val]
collectSegOutputs Env
env [Int]
worker_shape [TypeBase Shape NoMode]
return_types [KernelResult]
result_specs [([Int], [KernelResultValue])]
worker_rows
  | [([Int], [KernelResultValue])] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [([Int], [KernelResultValue])]
worker_rows =
      (TypeBase Shape NoMode -> KernelResult -> InterpM rep Val)
-> [TypeBase Shape NoMode] -> [KernelResult] -> InterpM rep [Val]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM TypeBase Shape NoMode -> KernelResult -> InterpM rep Val
forall {rep}.
TypeBase Shape NoMode -> KernelResult -> InterpM rep Val
collectEmpty [TypeBase Shape NoMode]
return_types [KernelResult]
result_specs
  | Bool
otherwise =
      (TypeBase Shape NoMode
 -> [([Int], KernelResultValue)] -> InterpM rep Val)
-> [TypeBase Shape NoMode]
-> [[([Int], KernelResultValue)]]
-> InterpM rep [Val]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM TypeBase Shape NoMode
-> [([Int], KernelResultValue)] -> InterpM rep Val
forall {rep}.
TypeBase Shape NoMode
-> [([Int], KernelResultValue)] -> InterpM rep Val
collectOne [TypeBase Shape NoMode]
return_types [[([Int], KernelResultValue)]]
columns
  where
    columns :: [[([Int], KernelResultValue)]]
columns = do
      result_index <- [Int
0 .. [TypeBase Shape NoMode] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase Shape NoMode]
return_types Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
      pure $ map (fmap (!! result_index)) worker_rows

    collectEmpty :: TypeBase Shape NoMode -> KernelResult -> InterpM rep Val
collectEmpty TypeBase Shape NoMode
result_type KernelResult
result_spec = do
      output_shape <- KernelResult -> InterpM rep [Int]
forall {rep}. KernelResult -> InterpM rep [Int]
resultShape KernelResult
result_spec
      collectOutputs env output_shape [result_type] [] >>= expectOne

    collectOne :: TypeBase Shape NoMode
-> [([Int], KernelResultValue)] -> InterpM rep Val
collectOne TypeBase Shape NoMode
result_type [([Int], KernelResultValue)]
rows =
      case [([Int], KernelResultValue)]
rows of
        [] -> [Char] -> InterpM rep Val
forall a. HasCallStack => [Char] -> a
error [Char]
"internal empty segmented output"
        ([Int]
_, KernelValue {}) : [([Int], KernelResultValue)]
_ -> do
          values <- (([Int], KernelResultValue) -> InterpM rep Val)
-> [([Int], KernelResultValue)] -> InterpM rep [Val]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (KernelResultValue -> InterpM rep Val
forall {f :: * -> *}. Applicative f => KernelResultValue -> f Val
expectOrdinary (KernelResultValue -> InterpM rep Val)
-> (([Int], KernelResultValue) -> KernelResultValue)
-> ([Int], KernelResultValue)
-> InterpM rep Val
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Int], KernelResultValue) -> KernelResultValue
forall a b. (a, b) -> b
snd) [([Int], KernelResultValue)]
rows
          collectOutputs env worker_shape [result_type] (map pure values)
            >>= expectOne
        ([Int]
_, KernelTile [(Int, Int)]
dimensions Val
_) : [([Int], KernelResultValue)]
_ -> do
          let output_shape :: [Int]
output_shape = ((Int, Int) -> Int) -> [(Int, Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Int) -> Int
forall a b. (a, b) -> a
fst [(Int, Int)]
dimensions
              tile_sizes :: [Int]
tile_sizes = ((Int, Int) -> Int) -> [(Int, Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Int) -> Int
forall a b. (a, b) -> b
snd [(Int, Int)]
dimensions
          TypeBase Shape NoMode
-> [Int]
-> [([Int], KernelResultValue)]
-> ([Int] -> ([Int], [Int]))
-> InterpM rep Val
forall {a} {rep}.
Eq a =>
TypeBase Shape NoMode
-> [Int]
-> [(a, KernelResultValue)]
-> ([Int] -> (a, [Int]))
-> InterpM rep Val
collectTiled
            TypeBase Shape NoMode
result_type
            [Int]
output_shape
            [([Int], KernelResultValue)]
rows
            ([Int] -> [Int] -> ([Int], [Int])
forall {c}. Integral c => [c] -> [c] -> ([c], [c])
tileSourceIndices [Int]
tile_sizes)
        ([Int]
_, KernelRegTile [(Int, Int, Int)]
dimensions Val
_) : [([Int], KernelResultValue)]
_ -> do
          let output_shape :: [Int]
output_shape = ((Int, Int, Int) -> Int) -> [(Int, Int, Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Int, Int) -> Int
forall {a} {b} {c}. (a, b, c) -> a
firstOf3 [(Int, Int, Int)]
dimensions
              block_sizes :: [Int]
block_sizes = ((Int, Int, Int) -> Int) -> [(Int, Int, Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Int, Int) -> Int
forall {a} {b} {c}. (a, b, c) -> b
secondOf3 [(Int, Int, Int)]
dimensions
              reg_sizes :: [Int]
reg_sizes = ((Int, Int, Int) -> Int) -> [(Int, Int, Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Int, Int) -> Int
forall {a} {b} {c}. (a, b, c) -> c
thirdOf3 [(Int, Int, Int)]
dimensions
          TypeBase Shape NoMode
-> [Int]
-> [([Int], KernelResultValue)]
-> ([Int] -> ([Int], [Int]))
-> InterpM rep Val
forall {a} {rep}.
Eq a =>
TypeBase Shape NoMode
-> [Int]
-> [(a, KernelResultValue)]
-> ([Int] -> (a, [Int]))
-> InterpM rep Val
collectTiled
            TypeBase Shape NoMode
result_type
            [Int]
output_shape
            [([Int], KernelResultValue)]
rows
            ([Int] -> [Int] -> [Int] -> ([Int], [Int])
forall {c}. Integral c => [c] -> [c] -> [c] -> ([c], [c])
regTileSourceIndices [Int]
block_sizes [Int]
reg_sizes)

    collectTiled :: TypeBase Shape NoMode
-> [Int]
-> [(a, KernelResultValue)]
-> ([Int] -> (a, [Int]))
-> InterpM rep Val
collectTiled TypeBase Shape NoMode
result_type [Int]
output_shape [(a, KernelResultValue)]
rows [Int] -> (a, [Int])
source_indices = do
      values <-
        ([Int] -> InterpM rep Val) -> [[Int]] -> InterpM rep [Val]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([(a, KernelResultValue)]
-> ([Int] -> (a, [Int])) -> [Int] -> InterpM rep Val
forall {a} {p} {rep}.
Eq a =>
[(a, KernelResultValue)]
-> (p -> (a, [Int])) -> p -> InterpM rep Val
tiledValueAt [(a, KernelResultValue)]
rows [Int] -> (a, [Int])
source_indices) ([[Int]] -> InterpM rep [Val]) -> [[Int]] -> InterpM rep [Val]
forall a b. (a -> b) -> a -> b
$ [Int] -> [[Int]]
allCoordinates [Int]
output_shape
      collectOutputs env output_shape [result_type] (map pure values)
        >>= expectOne

    tiledValueAt :: [(a, KernelResultValue)]
-> (p -> (a, [Int])) -> p -> InterpM rep Val
tiledValueAt [(a, KernelResultValue)]
rows p -> (a, [Int])
source_indices p
output_coordinate = do
      let (a
worker_coordinate, [Int]
source_coordinate) =
            p -> (a, [Int])
source_indices p
output_coordinate
          tile :: Val
tile = case a -> [(a, KernelResultValue)] -> Maybe KernelResultValue
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup a
worker_coordinate [(a, KernelResultValue)]
rows of
            Just (KernelTile [(Int, Int)]
_ Val
tile_value) -> Val
tile_value
            Just (KernelRegTile [(Int, Int, Int)]
_ Val
tile_value) -> Val
tile_value
            Maybe KernelResultValue
_ -> [Char] -> Val
forall a. HasCallStack => [Char] -> a
error [Char]
"inconsistent segmented result layout"
      [Int] -> Val -> InterpM rep Val
forall rep. [Int] -> Val -> InterpM rep Val
readAccumulatorElement [Int]
source_coordinate Val
tile

    tileSourceIndices :: [c] -> [c] -> ([c], [c])
tileSourceIndices [c]
tile_sizes [c]
output_coordinate =
      ( (c -> c -> c) -> [c] -> [c] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith c -> c -> c
forall a. Integral a => a -> a -> a
div [c]
output_coordinate [c]
tile_sizes,
        (c -> c -> c) -> [c] -> [c] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith c -> c -> c
forall a. Integral a => a -> a -> a
mod [c]
output_coordinate [c]
tile_sizes
      )

    regTileSourceIndices :: [c] -> [c] -> [c] -> ([c], [c])
regTileSourceIndices [c]
block_sizes [c]
reg_sizes [c]
output_coordinate =
      ([c]
workers, [c]
block_indices [c] -> [c] -> [c]
forall a. Semigroup a => a -> a -> a
<> [c]
reg_indices)
      where
        tile_sizes :: [c]
tile_sizes = (c -> c -> c) -> [c] -> [c] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith c -> c -> c
forall a. Num a => a -> a -> a
(*) [c]
block_sizes [c]
reg_sizes
        workers :: [c]
workers = (c -> c -> c) -> [c] -> [c] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith c -> c -> c
forall a. Integral a => a -> a -> a
div [c]
output_coordinate [c]
tile_sizes
        within_tiles :: [c]
within_tiles = (c -> c -> c) -> [c] -> [c] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith c -> c -> c
forall a. Integral a => a -> a -> a
mod [c]
output_coordinate [c]
tile_sizes
        block_indices :: [c]
block_indices = (c -> c -> c) -> [c] -> [c] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith c -> c -> c
forall a. Integral a => a -> a -> a
div [c]
within_tiles [c]
reg_sizes
        reg_indices :: [c]
reg_indices = (c -> c -> c) -> [c] -> [c] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith c -> c -> c
forall a. Integral a => a -> a -> a
mod [c]
within_tiles [c]
reg_sizes

    resultShape :: KernelResult -> InterpM rep [Int]
resultShape Seg.Returns {} = [Int] -> InterpM rep [Int]
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Int]
worker_shape
    resultShape (Seg.TileReturns Certs
_ [(SubExp, SubExp)]
dimensions VName
_) =
      ((SubExp, SubExp) -> InterpM rep Int)
-> [(SubExp, SubExp)] -> InterpM rep [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Env -> SubExp -> InterpM rep Int
forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env (SubExp -> InterpM rep Int)
-> ((SubExp, SubExp) -> SubExp)
-> (SubExp, SubExp)
-> InterpM rep Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SubExp, SubExp) -> SubExp
forall a b. (a, b) -> a
fst) [(SubExp, SubExp)]
dimensions
    resultShape (Seg.RegTileReturns Certs
_ [(SubExp, SubExp, SubExp)]
dimensions VName
_) =
      ((SubExp, SubExp, SubExp) -> InterpM rep Int)
-> [(SubExp, SubExp, SubExp)] -> InterpM rep [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Env -> SubExp -> InterpM rep Int
forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env (SubExp -> InterpM rep Int)
-> ((SubExp, SubExp, SubExp) -> SubExp)
-> (SubExp, SubExp, SubExp)
-> InterpM rep Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SubExp, SubExp, SubExp) -> SubExp
forall {a} {b} {c}. (a, b, c) -> a
firstOf3) [(SubExp, SubExp, SubExp)]
dimensions

    expectOrdinary :: KernelResultValue -> f Val
expectOrdinary (KernelValue Val
result_value) = Val -> f Val
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Val
result_value
    expectOrdinary KernelResultValue
_ = [Char] -> f Val
forall a. HasCallStack => [Char] -> a
error [Char]
"inconsistent segmented result layout"

    expectOne :: [a] -> f a
expectOne [a
result_value] = a -> f a
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
result_value
    expectOne [a]
_ = [Char] -> f a
forall a. HasCallStack => [Char] -> a
error [Char]
"invalid segmented output count"

    firstOf3 :: (a, b, c) -> a
firstOf3 (a
first, b
_, c
_) = a
first
    secondOf3 :: (a, b, c) -> b
secondOf3 (a
_, b
second, c
_) = b
second
    thirdOf3 :: (a, b, c) -> c
thirdOf3 (a
_, b
_, c
third) = c
third

-- | Run a program in the IR specified by rep.
runProgram ::
  (Typed (FParamInfo rep)) =>
  OpEvaluator rep ->
  Prog rep ->
  Name ->
  [V.Value] ->
  IO (Either T.Text [V.Value])
runProgram :: forall rep.
Typed (FParamInfo rep) =>
OpEvaluator rep
-> Prog rep -> Name -> [Value] -> IO (Either Text [Value])
runProgram OpEvaluator rep
eval_op Prog rep
prog Name
entry [Value]
inputs = ExceptT Text IO [Value] -> IO (Either Text [Value])
forall e (m :: * -> *) a. ExceptT e m a -> m (Either e a)
runExceptT (ExceptT Text IO [Value] -> IO (Either Text [Value]))
-> ExceptT Text IO [Value] -> IO (Either Text [Value])
forall a b. (a -> b) -> a -> b
$ (ReaderT (InterpEnv rep) (ExceptT Text IO) [Value]
 -> InterpEnv rep -> ExceptT Text IO [Value])
-> InterpEnv rep
-> ReaderT (InterpEnv rep) (ExceptT Text IO) [Value]
-> ExceptT Text IO [Value]
forall a b c. (a -> b -> c) -> b -> a -> c
flip ReaderT (InterpEnv rep) (ExceptT Text IO) [Value]
-> InterpEnv rep -> ExceptT Text IO [Value]
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT InterpEnv rep
initial_interp_env (ReaderT (InterpEnv rep) (ExceptT Text IO) [Value]
 -> ExceptT Text IO [Value])
-> (InterpM rep [Value]
    -> ReaderT (InterpEnv rep) (ExceptT Text IO) [Value])
-> InterpM rep [Value]
-> ExceptT Text IO [Value]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InterpM rep [Value]
-> ReaderT (InterpEnv rep) (ExceptT Text IO) [Value]
forall rep a.
InterpM rep a -> ReaderT (InterpEnv rep) (ExceptT Text IO) a
unInterpM (InterpM rep [Value] -> ExceptT Text IO [Value])
-> InterpM rep [Value] -> ExceptT Text IO [Value]
forall a b. (a -> b) -> a -> b
$ do
  let funs :: Map Name (FunDef rep)
funs = [(Name, FunDef rep)] -> Map Name (FunDef rep)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(FunDef rep -> Name
forall rep. FunDef rep -> Name
funDefName FunDef rep
fun, FunDef rep
fun) | FunDef rep
fun <- Prog rep -> [FunDef rep]
forall rep. Prog rep -> [FunDef rep]
progFuns Prog rep
prog]
  consts_env <- Map Name (FunDef rep) -> Env -> [Stm rep] -> InterpM rep Env
forall {rep}. FunEnv rep -> Env -> [Stm rep] -> InterpM rep Env
foldConsts Map Name (FunDef rep)
funs Env
forall a. Monoid a => a
mempty (Stms rep -> [Stm rep]
forall rep. Stms rep -> [Stm rep]
stmsToList (Prog rep -> Stms rep
forall rep. Prog rep -> Stms rep
progConsts Prog rep
prog)) -- top-level consts
  fun <- findEntry prog entry
  converted_inputs <- mapM fromValue inputs
  let params = FunDef rep -> [FParam rep]
forall rep. FunDef rep -> [FParam rep]
funDefParams FunDef rep
fun
      value_count = [([Val], Val)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [([Val], Val)]
converted_inputs
      context_count = [FParam rep] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [FParam rep]
params Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
value_count
  if context_count < 0
    then interpError "entry point argument count mismatch"
    else do
      let (context_params, value_params) = splitAt context_count params
      shape_bindings <-
        foldM bindInputShape mempty $ zip value_params converted_inputs
      context_args <- mapM (lookupContextArg shape_bindings) context_params
      let value_args = (([Val], Val) -> Val) -> [([Val], Val)] -> [Val]
forall a b. (a -> b) -> [a] -> [b]
map ([Val], Val) -> Val
forall a b. (a, b) -> b
snd [([Val], Val)]
converted_inputs
          arg_vals = [Val]
context_args [Val] -> [Val] -> [Val]
forall a. Semigroup a => a -> a -> a
<> [Val]
value_args
          env =
            Env -> Env -> Env
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union
              ([(VName, Val)] -> Env
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Val)] -> Env) -> [(VName, Val)] -> Env
forall a b. (a -> b) -> a -> b
$ [VName] -> [Val] -> [(VName, Val)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((FParam rep -> VName) -> [FParam rep] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map FParam rep -> VName
forall dec. Param dec -> VName
paramName [FParam rep]
params) [Val]
arg_vals)
              Env
consts_env
      results <- evalBody funs env (funDefBody fun)
      result_signedness <- entryResultSignedness prog fun
      zipWithM
        toValue
        result_signedness
        (drop (length results - length result_signedness) results)
  where
    initial_interp_env :: InterpEnv rep
initial_interp_env =
      InterpEnv
        { interpOpEvaluator :: OpEvaluator rep
interpOpEvaluator = OpEvaluator rep
eval_op,
          interpAccOps :: Map VName (Map Name (FunDef rep), Env, Lambda rep)
interpAccOps = Map VName (Map Name (FunDef rep), Env, Lambda rep)
forall a. Monoid a => a
mempty
        }

    foldConsts :: FunEnv rep -> Env -> [Stm rep] -> InterpM rep Env
foldConsts FunEnv rep
_ Env
e [] = Env -> InterpM rep Env
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Env
e
    foldConsts FunEnv rep
funs Env
e (Stm rep
s : [Stm rep]
ss) = FunEnv rep -> Env -> Stm rep -> InterpM rep Env
forall rep. FunEnv rep -> Env -> Stm rep -> InterpM rep Env
evalStm FunEnv rep
funs Env
e Stm rep
s InterpM rep Env -> (Env -> InterpM rep Env) -> InterpM rep Env
forall a b. InterpM rep a -> (a -> InterpM rep b) -> InterpM rep b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Env
e' -> FunEnv rep -> Env -> [Stm rep] -> InterpM rep Env
foldConsts FunEnv rep
funs Env
e' [Stm rep]
ss

    bindInputShape :: Env -> (Param dec, ([Val], b)) -> InterpM rep Env
bindInputShape Env
bindings (Param dec
param, ([Val]
actual_dims, b
_))
      | [SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
expected_dims Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Val] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Val]
actual_dims =
          Text -> InterpM rep Env
forall rep a. Text -> InterpM rep a
interpError Text
"entry point input rank mismatch"
      | Bool
otherwise =
          (Env -> (SubExp, Val) -> InterpM rep Env)
-> Env -> [(SubExp, Val)] -> InterpM rep Env
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM Env -> (SubExp, Val) -> InterpM rep Env
forall {rep}. Env -> (SubExp, Val) -> InterpM rep Env
bindDimension Env
bindings ([(SubExp, Val)] -> InterpM rep Env)
-> [(SubExp, Val)] -> InterpM rep Env
forall a b. (a -> b) -> a -> b
$ [SubExp] -> [Val] -> [(SubExp, Val)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
expected_dims [Val]
actual_dims
      where
        expected_dims :: [SubExp]
expected_dims = TypeBase Shape NoMode -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims (TypeBase Shape NoMode -> [SubExp])
-> TypeBase Shape NoMode -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Param dec -> TypeBase Shape NoMode
forall dec. Typed dec => Param dec -> TypeBase Shape NoMode
paramType Param dec
param

    bindDimension :: Env -> (SubExp, Val) -> InterpM rep Env
bindDimension Env
bindings (Var VName
name, Val
actual_dim) =
      case VName -> Env -> Maybe Val
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name Env
bindings of
        Maybe Val
Nothing -> Env -> InterpM rep Env
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Env -> InterpM rep Env) -> Env -> InterpM rep Env
forall a b. (a -> b) -> a -> b
$ VName -> Val -> Env -> Env
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
name Val
actual_dim Env
bindings
        Just Val
expected_dim
          | Val -> Val -> Bool
sameDimension Val
expected_dim Val
actual_dim -> Env -> InterpM rep Env
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Env
bindings
          | Bool
otherwise -> Text -> InterpM rep Env
forall rep a. Text -> InterpM rep a
interpError Text
"entry point input shape mismatch"
    bindDimension Env
bindings (Constant PrimValue
expected, PrimVal PrimValue
actual)
      | PrimValue
expected PrimValue -> PrimValue -> Bool
forall a. Eq a => a -> a -> Bool
== PrimValue
actual = Env -> InterpM rep Env
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Env
bindings
      | Bool
otherwise = Text -> InterpM rep Env
forall rep a. Text -> InterpM rep a
interpError Text
"entry point input shape mismatch"
    bindDimension Env
_ (SubExp, Val)
_ =
      [Char] -> InterpM rep Env
forall a. HasCallStack => [Char] -> a
error [Char]
"invalid entry point input dimension"

    lookupContextArg :: Map VName a -> Param dec -> f a
lookupContextArg Map VName a
bindings Param dec
param =
      f a -> (a -> f a) -> Maybe a -> f a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
        ([Char] -> f a
forall a. HasCallStack => [Char] -> a
error [Char]
"unbound entry point shape parameter")
        a -> f a
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        (VName -> Map VName a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (Param dec -> VName
forall dec. Param dec -> VName
paramName Param dec
param) Map VName a
bindings)

    sameDimension :: Val -> Val -> Bool
sameDimension (PrimVal PrimValue
x) (PrimVal PrimValue
y) = PrimValue
x PrimValue -> PrimValue -> Bool
forall a. Eq a => a -> a -> Bool
== PrimValue
y
    sameDimension Val
_ Val
_ = Bool
False

entryResultSignedness :: Prog rep -> FunDef rep -> InterpM rep [Signedness]
entryResultSignedness :: forall rep. Prog rep -> FunDef rep -> InterpM rep [Signedness]
entryResultSignedness Prog rep
prog FunDef rep
fun =
  case FunDef rep -> Maybe EntryPoint
forall rep. FunDef rep -> Maybe EntryPoint
funDefEntryPoint FunDef rep
fun of
    Just (Name
_, [EntryParam]
_, EntryResult
result, Maybe Text
_) ->
      OpaqueTypes -> EntryPointType -> InterpM rep [Signedness]
forall {f :: * -> *}.
Monad f =>
OpaqueTypes -> EntryPointType -> f [Signedness]
entryPointTypeSignedness (Prog rep -> OpaqueTypes
forall rep. Prog rep -> OpaqueTypes
progTypes Prog rep
prog) (EntryPointType -> InterpM rep [Signedness])
-> EntryPointType -> InterpM rep [Signedness]
forall a b. (a -> b) -> a -> b
$ EntryResult -> EntryPointType
entryResultType EntryResult
result
    Maybe EntryPoint
Nothing ->
      [Char] -> InterpM rep [Signedness]
forall a. HasCallStack => [Char] -> a
error [Char]
"function is not an entry point"
  where
    entryPointTypeSignedness :: OpaqueTypes -> EntryPointType -> f [Signedness]
entryPointTypeSignedness OpaqueTypes
_ (TypeTransparent ValueType
value_type) =
      [Signedness] -> f [Signedness]
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [ValueType -> Signedness
valueTypeSignedness ValueType
value_type]
    entryPointTypeSignedness types :: OpaqueTypes
types@(OpaqueTypes [(Name, (OpaqueType, Maybe Text))]
opaque_types) (TypeOpaque Name
name) =
      case Name
-> [(Name, (OpaqueType, Maybe Text))]
-> Maybe (OpaqueType, Maybe Text)
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
name [(Name, (OpaqueType, Maybe Text))]
opaque_types of
        Maybe (OpaqueType, Maybe Text)
Nothing -> [Char] -> f [Signedness]
forall a. HasCallStack => [Char] -> a
error ([Char] -> f [Signedness]) -> [Char] -> f [Signedness]
forall a b. (a -> b) -> a -> b
$ [Char]
"unknown opaque type: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Name -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Name
name
        Just (OpaqueType
opaque_type, Maybe Text
_) -> OpaqueTypes -> OpaqueType -> f [Signedness]
opaqueTypeSignedness OpaqueTypes
types OpaqueType
opaque_type

    opaqueTypeSignedness :: OpaqueTypes -> OpaqueType -> f [Signedness]
opaqueTypeSignedness OpaqueTypes
_ (OpaqueArray Int
_ Name
_ [ValueType]
value_types) =
      [Signedness] -> f [Signedness]
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Signedness] -> f [Signedness]) -> [Signedness] -> f [Signedness]
forall a b. (a -> b) -> a -> b
$ (ValueType -> Signedness) -> [ValueType] -> [Signedness]
forall a b. (a -> b) -> [a] -> [b]
map ValueType -> Signedness
valueTypeSignedness [ValueType]
value_types
    opaqueTypeSignedness OpaqueTypes
types (OpaqueRecordArray Int
_ Name
_ [(Name, EntryPointType)]
fields) =
      [[Signedness]] -> [Signedness]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Signedness]] -> [Signedness])
-> f [[Signedness]] -> f [Signedness]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Name, EntryPointType) -> f [Signedness])
-> [(Name, EntryPointType)] -> f [[Signedness]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (OpaqueTypes -> EntryPointType -> f [Signedness]
entryPointTypeSignedness OpaqueTypes
types (EntryPointType -> f [Signedness])
-> ((Name, EntryPointType) -> EntryPointType)
-> (Name, EntryPointType)
-> f [Signedness]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, EntryPointType) -> EntryPointType
forall a b. (a, b) -> b
snd) [(Name, EntryPointType)]
fields
    opaqueTypeSignedness OpaqueTypes
_ (OpaqueRecord []) =
      [Signedness] -> f [Signedness]
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Signedness
Signed]
    opaqueTypeSignedness OpaqueTypes
types (OpaqueRecord [(Name, EntryPointType)]
fields) =
      [[Signedness]] -> [Signedness]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Signedness]] -> [Signedness])
-> f [[Signedness]] -> f [Signedness]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Name, EntryPointType) -> f [Signedness])
-> [(Name, EntryPointType)] -> f [[Signedness]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (OpaqueTypes -> EntryPointType -> f [Signedness]
entryPointTypeSignedness OpaqueTypes
types (EntryPointType -> f [Signedness])
-> ((Name, EntryPointType) -> EntryPointType)
-> (Name, EntryPointType)
-> f [Signedness]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, EntryPointType) -> EntryPointType
forall a b. (a, b) -> b
snd) [(Name, EntryPointType)]
fields
    opaqueTypeSignedness OpaqueTypes
_ (OpaqueSum [ValueType]
value_types [(Name, [(EntryPointType, [Int])])]
_) =
      [Signedness] -> f [Signedness]
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Signedness] -> f [Signedness]) -> [Signedness] -> f [Signedness]
forall a b. (a -> b) -> a -> b
$ (ValueType -> Signedness) -> [ValueType] -> [Signedness]
forall a b. (a -> b) -> [a] -> [b]
map ValueType -> Signedness
valueTypeSignedness [ValueType]
value_types

    valueTypeSignedness :: ValueType -> Signedness
valueTypeSignedness (ValueType Signedness
signedness Rank
_ PrimType
_) = Signedness
signedness

findEntry :: Prog rep -> Name -> InterpM rep (FunDef rep)
findEntry :: forall rep. Prog rep -> Name -> InterpM rep (FunDef rep)
findEntry Prog rep
prog Name
name =
  InterpM rep (FunDef rep)
-> (FunDef rep -> InterpM rep (FunDef rep))
-> Maybe (FunDef rep)
-> InterpM rep (FunDef rep)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
    (Text -> InterpM rep (FunDef rep)
forall rep a. Text -> InterpM rep a
interpError (Text -> InterpM rep (FunDef rep))
-> Text -> InterpM rep (FunDef rep)
forall a b. (a -> b) -> a -> b
$ Text
"entry point not found: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Name -> Text
forall a. Pretty a => a -> Text
prettyText Name
name)
    FunDef rep -> InterpM rep (FunDef rep)
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    (Maybe (FunDef rep) -> InterpM rep (FunDef rep))
-> Maybe (FunDef rep) -> InterpM rep (FunDef rep)
forall a b. (a -> b) -> a -> b
$ Name -> [(Name, FunDef rep)] -> Maybe (FunDef rep)
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup
      Name
name
      [ (Name
entry_name, FunDef rep
fun)
      | FunDef rep
fun <- Prog rep -> [FunDef rep]
forall rep. Prog rep -> [FunDef rep]
progFuns Prog rep
prog,
        Just (Name
entry_name, [EntryParam]
_, EntryResult
_, Maybe Text
_) <- [FunDef rep -> Maybe EntryPoint
forall rep. FunDef rep -> Maybe EntryPoint
funDefEntryPoint FunDef rep
fun]
      ]

fromValue :: V.Value -> InterpM rep ([Val], Val)
fromValue :: forall rep. Value -> InterpM rep ([Val], Val)
fromValue (V.I8Value Vector Int
shape Vector Int8
values) =
  Vector Int
-> PrimType
-> (Int8 -> PrimValue)
-> Vector Int8
-> InterpM rep ([Val], Val)
forall a rep.
Storable a =>
Vector Int
-> PrimType
-> (a -> PrimValue)
-> Vector a
-> InterpM rep ([Val], Val)
fromPrimitiveVector Vector Int
shape (IntType -> PrimType
IntType IntType
Int8) (IntValue -> PrimValue
IntValue (IntValue -> PrimValue) -> (Int8 -> IntValue) -> Int8 -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int8 -> IntValue
Int8Value) Vector Int8
values
fromValue (V.I16Value Vector Int
shape Vector Int16
values) =
  Vector Int
-> PrimType
-> (Int16 -> PrimValue)
-> Vector Int16
-> InterpM rep ([Val], Val)
forall a rep.
Storable a =>
Vector Int
-> PrimType
-> (a -> PrimValue)
-> Vector a
-> InterpM rep ([Val], Val)
fromPrimitiveVector Vector Int
shape (IntType -> PrimType
IntType IntType
Int16) (IntValue -> PrimValue
IntValue (IntValue -> PrimValue)
-> (Int16 -> IntValue) -> Int16 -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int16 -> IntValue
Int16Value) Vector Int16
values
fromValue (V.I32Value Vector Int
shape Vector Int32
values) =
  Vector Int
-> PrimType
-> (Int32 -> PrimValue)
-> Vector Int32
-> InterpM rep ([Val], Val)
forall a rep.
Storable a =>
Vector Int
-> PrimType
-> (a -> PrimValue)
-> Vector a
-> InterpM rep ([Val], Val)
fromPrimitiveVector Vector Int
shape (IntType -> PrimType
IntType IntType
Int32) (IntValue -> PrimValue
IntValue (IntValue -> PrimValue)
-> (Int32 -> IntValue) -> Int32 -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int32 -> IntValue
Int32Value) Vector Int32
values
fromValue (V.I64Value Vector Int
shape Vector Int64
values) =
  Vector Int
-> PrimType
-> (Int64 -> PrimValue)
-> Vector Int64
-> InterpM rep ([Val], Val)
forall a rep.
Storable a =>
Vector Int
-> PrimType
-> (a -> PrimValue)
-> Vector a
-> InterpM rep ([Val], Val)
fromPrimitiveVector Vector Int
shape (IntType -> PrimType
IntType IntType
Int64) (IntValue -> PrimValue
IntValue (IntValue -> PrimValue)
-> (Int64 -> IntValue) -> Int64 -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> IntValue
Int64Value) Vector Int64
values
fromValue (V.U8Value Vector Int
shape Vector Word8
values) =
  Vector Int
-> PrimType
-> (Word8 -> PrimValue)
-> Vector Word8
-> InterpM rep ([Val], Val)
forall a rep.
Storable a =>
Vector Int
-> PrimType
-> (a -> PrimValue)
-> Vector a
-> InterpM rep ([Val], Val)
fromPrimitiveVector Vector Int
shape (IntType -> PrimType
IntType IntType
Int8) (IntValue -> PrimValue
IntValue (IntValue -> PrimValue)
-> (Word8 -> IntValue) -> Word8 -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int8 -> IntValue
Int8Value (Int8 -> IntValue) -> (Word8 -> Int8) -> Word8 -> IntValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word8 -> Int8
forall a b. (Integral a, Num b) => a -> b
fromIntegral) Vector Word8
values
fromValue (V.U16Value Vector Int
shape Vector Word16
values) =
  Vector Int
-> PrimType
-> (Word16 -> PrimValue)
-> Vector Word16
-> InterpM rep ([Val], Val)
forall a rep.
Storable a =>
Vector Int
-> PrimType
-> (a -> PrimValue)
-> Vector a
-> InterpM rep ([Val], Val)
fromPrimitiveVector Vector Int
shape (IntType -> PrimType
IntType IntType
Int16) (IntValue -> PrimValue
IntValue (IntValue -> PrimValue)
-> (Word16 -> IntValue) -> Word16 -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int16 -> IntValue
Int16Value (Int16 -> IntValue) -> (Word16 -> Int16) -> Word16 -> IntValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word16 -> Int16
forall a b. (Integral a, Num b) => a -> b
fromIntegral) Vector Word16
values
fromValue (V.U32Value Vector Int
shape Vector Word32
values) =
  Vector Int
-> PrimType
-> (Word32 -> PrimValue)
-> Vector Word32
-> InterpM rep ([Val], Val)
forall a rep.
Storable a =>
Vector Int
-> PrimType
-> (a -> PrimValue)
-> Vector a
-> InterpM rep ([Val], Val)
fromPrimitiveVector Vector Int
shape (IntType -> PrimType
IntType IntType
Int32) (IntValue -> PrimValue
IntValue (IntValue -> PrimValue)
-> (Word32 -> IntValue) -> Word32 -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int32 -> IntValue
Int32Value (Int32 -> IntValue) -> (Word32 -> Int32) -> Word32 -> IntValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word32 -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral) Vector Word32
values
fromValue (V.U64Value Vector Int
shape Vector Word64
values) =
  Vector Int
-> PrimType
-> (Word64 -> PrimValue)
-> Vector Word64
-> InterpM rep ([Val], Val)
forall a rep.
Storable a =>
Vector Int
-> PrimType
-> (a -> PrimValue)
-> Vector a
-> InterpM rep ([Val], Val)
fromPrimitiveVector Vector Int
shape (IntType -> PrimType
IntType IntType
Int64) (IntValue -> PrimValue
IntValue (IntValue -> PrimValue)
-> (Word64 -> IntValue) -> Word64 -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> IntValue
Int64Value (Int64 -> IntValue) -> (Word64 -> Int64) -> Word64 -> IntValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word64 -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral) Vector Word64
values
fromValue (V.F16Value Vector Int
shape Vector Half
values) =
  Vector Int
-> PrimType
-> (Half -> PrimValue)
-> Vector Half
-> InterpM rep ([Val], Val)
forall a rep.
Storable a =>
Vector Int
-> PrimType
-> (a -> PrimValue)
-> Vector a
-> InterpM rep ([Val], Val)
fromPrimitiveVector Vector Int
shape (FloatType -> PrimType
FloatType FloatType
Float16) (FloatValue -> PrimValue
FloatValue (FloatValue -> PrimValue)
-> (Half -> FloatValue) -> Half -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Half -> FloatValue
Float16Value) Vector Half
values
fromValue (V.F32Value Vector Int
shape Vector Float
values) =
  Vector Int
-> PrimType
-> (Float -> PrimValue)
-> Vector Float
-> InterpM rep ([Val], Val)
forall a rep.
Storable a =>
Vector Int
-> PrimType
-> (a -> PrimValue)
-> Vector a
-> InterpM rep ([Val], Val)
fromPrimitiveVector Vector Int
shape (FloatType -> PrimType
FloatType FloatType
Float32) (FloatValue -> PrimValue
FloatValue (FloatValue -> PrimValue)
-> (Float -> FloatValue) -> Float -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Float -> FloatValue
Float32Value) Vector Float
values
fromValue (V.F64Value Vector Int
shape Vector Double
values) =
  Vector Int
-> PrimType
-> (Double -> PrimValue)
-> Vector Double
-> InterpM rep ([Val], Val)
forall a rep.
Storable a =>
Vector Int
-> PrimType
-> (a -> PrimValue)
-> Vector a
-> InterpM rep ([Val], Val)
fromPrimitiveVector Vector Int
shape (FloatType -> PrimType
FloatType FloatType
Float64) (FloatValue -> PrimValue
FloatValue (FloatValue -> PrimValue)
-> (Double -> FloatValue) -> Double -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> FloatValue
Float64Value) Vector Double
values
fromValue (V.BoolValue Vector Int
shape Vector Bool
values) =
  Vector Int
-> PrimType
-> (Bool -> PrimValue)
-> Vector Bool
-> InterpM rep ([Val], Val)
forall a rep.
Storable a =>
Vector Int
-> PrimType
-> (a -> PrimValue)
-> Vector a
-> InterpM rep ([Val], Val)
fromPrimitiveVector Vector Int
shape PrimType
Bool Bool -> PrimValue
BoolValue Vector Bool
values

fromPrimitiveVector ::
  (Storable a) =>
  SVec.Vector Int ->
  PrimType ->
  (a -> PrimValue) ->
  SVec.Vector a ->
  InterpM rep ([Val], Val)
fromPrimitiveVector :: forall a rep.
Storable a =>
Vector Int
-> PrimType
-> (a -> PrimValue)
-> Vector a
-> InterpM rep ([Val], Val)
fromPrimitiveVector Vector Int
shape PrimType
element_type a -> PrimValue
wrap Vector a
values
  | (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) [Int]
dimensions =
      Text -> InterpM rep ([Val], Val)
forall rep a. Text -> InterpM rep a
interpError Text
"input array dimensions cannot be negative"
  | [Int] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
dimensions =
      case [PrimValue]
primitive_values of
        [PrimValue
primitive_value] -> ([Val], Val) -> InterpM rep ([Val], Val)
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], PrimValue -> Val
PrimVal PrimValue
primitive_value)
        [PrimValue]
_ -> Text -> InterpM rep ([Val], Val)
forall rep a. Text -> InterpM rep a
interpError Text
"invalid scalar input storage"
  | [PrimValue] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [PrimValue]
primitive_values Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
dimensions =
      Text -> InterpM rep ([Val], Val)
forall rep a. Text -> InterpM rep a
interpError Text
"invalid input array storage"
  | Bool
otherwise = do
      array <- [Int] -> PrimType -> [PrimValue] -> InterpM rep Val
forall rep. [Int] -> PrimType -> [PrimValue] -> InterpM rep Val
newArrayValue [Int]
dimensions PrimType
element_type [PrimValue]
primitive_values
      pure
        ( map
            (PrimVal . IntValue . Int64Value . fromIntegral)
            dimensions,
          array
        )
  where
    dimensions :: [Int]
dimensions = Vector Int -> [Int]
forall a. Storable a => Vector a -> [a]
SVec.toList Vector Int
shape
    primitive_values :: [PrimValue]
primitive_values = (a -> PrimValue) -> [a] -> [PrimValue]
forall a b. (a -> b) -> [a] -> [b]
map a -> PrimValue
wrap ([a] -> [PrimValue]) -> [a] -> [PrimValue]
forall a b. (a -> b) -> a -> b
$ Vector a -> [a]
forall a. Storable a => Vector a -> [a]
SVec.toList Vector a
values

toValue :: Signedness -> Val -> InterpM rep V.Value
toValue :: forall rep. Signedness -> Val -> InterpM rep Value
toValue Signedness
signedness (PrimVal PrimValue
primitive_value) =
  Signedness -> [Int] -> PrimType -> [PrimValue] -> InterpM rep Value
forall rep.
Signedness -> [Int] -> PrimType -> [PrimValue] -> InterpM rep Value
toPrimitiveValue Signedness
signedness [] (PrimValue -> PrimType
P.primValueType PrimValue
primitive_value) [PrimValue
primitive_value]
toValue Signedness
signedness (ArrayValue [Int]
shape ArrayValues
values) = do
  primitive_values <- ArrayValues -> InterpM rep [PrimValue]
forall rep. ArrayValues -> InterpM rep [PrimValue]
arrayValues ArrayValues
values
  toPrimitiveValue signedness shape (arrayValuesType values) primitive_values
toValue Signedness
_ AccValue {} =
  [Char] -> InterpM rep Value
forall a. HasCallStack => [Char] -> a
error [Char]
"accumulators cannot be represented as external values"

toPrimitiveValue :: Signedness -> [Int] -> PrimType -> [PrimValue] -> InterpM rep V.Value
toPrimitiveValue :: forall rep.
Signedness -> [Int] -> PrimType -> [PrimValue] -> InterpM rep Value
toPrimitiveValue Signedness
Signed [Int]
shape (IntType IntType
Int8) [PrimValue]
values =
  Value -> InterpM rep Value
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> InterpM rep Value) -> Value -> InterpM rep Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int8 -> Value
V.I8Value ([Int] -> Vector Int
shapeVector [Int]
shape) (Vector Int8 -> Value) -> Vector Int8 -> Value
forall a b. (a -> b) -> a -> b
$ [Int8] -> Vector Int8
forall a. Storable a => [a] -> Vector a
SVec.fromList ([Int8] -> Vector Int8) -> [Int8] -> Vector Int8
forall a b. (a -> b) -> a -> b
$ (PrimValue -> Int8) -> [PrimValue] -> [Int8]
forall a b. (a -> b) -> [a] -> [b]
map PrimValue -> Int8
expectInt8 [PrimValue]
values
  where
    expectInt8 :: PrimValue -> Int8
expectInt8 (IntValue (Int8Value Int8
element)) = Int8
element
    expectInt8 PrimValue
_ = [Char] -> Int8
forall a. HasCallStack => [Char] -> a
error [Char]
"expected an i8 value"
toPrimitiveValue Signedness
Signed [Int]
shape (IntType IntType
Int16) [PrimValue]
values =
  Value -> InterpM rep Value
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> InterpM rep Value) -> Value -> InterpM rep Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int16 -> Value
V.I16Value ([Int] -> Vector Int
shapeVector [Int]
shape) (Vector Int16 -> Value) -> Vector Int16 -> Value
forall a b. (a -> b) -> a -> b
$ [Int16] -> Vector Int16
forall a. Storable a => [a] -> Vector a
SVec.fromList ([Int16] -> Vector Int16) -> [Int16] -> Vector Int16
forall a b. (a -> b) -> a -> b
$ (PrimValue -> Int16) -> [PrimValue] -> [Int16]
forall a b. (a -> b) -> [a] -> [b]
map PrimValue -> Int16
expectInt16 [PrimValue]
values
  where
    expectInt16 :: PrimValue -> Int16
expectInt16 (IntValue (Int16Value Int16
element)) = Int16
element
    expectInt16 PrimValue
_ = [Char] -> Int16
forall a. HasCallStack => [Char] -> a
error [Char]
"expected an i16 value"
toPrimitiveValue Signedness
Signed [Int]
shape (IntType IntType
Int32) [PrimValue]
values =
  Value -> InterpM rep Value
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> InterpM rep Value) -> Value -> InterpM rep Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int32 -> Value
V.I32Value ([Int] -> Vector Int
shapeVector [Int]
shape) (Vector Int32 -> Value) -> Vector Int32 -> Value
forall a b. (a -> b) -> a -> b
$ [Int32] -> Vector Int32
forall a. Storable a => [a] -> Vector a
SVec.fromList ([Int32] -> Vector Int32) -> [Int32] -> Vector Int32
forall a b. (a -> b) -> a -> b
$ (PrimValue -> Int32) -> [PrimValue] -> [Int32]
forall a b. (a -> b) -> [a] -> [b]
map PrimValue -> Int32
expectInt32 [PrimValue]
values
  where
    expectInt32 :: PrimValue -> Int32
expectInt32 (IntValue (Int32Value Int32
element)) = Int32
element
    expectInt32 PrimValue
_ = [Char] -> Int32
forall a. HasCallStack => [Char] -> a
error [Char]
"expected an i32 value"
toPrimitiveValue Signedness
Signed [Int]
shape (IntType IntType
Int64) [PrimValue]
values =
  Value -> InterpM rep Value
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> InterpM rep Value) -> Value -> InterpM rep Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int64 -> Value
V.I64Value ([Int] -> Vector Int
shapeVector [Int]
shape) (Vector Int64 -> Value) -> Vector Int64 -> Value
forall a b. (a -> b) -> a -> b
$ [Int64] -> Vector Int64
forall a. Storable a => [a] -> Vector a
SVec.fromList ([Int64] -> Vector Int64) -> [Int64] -> Vector Int64
forall a b. (a -> b) -> a -> b
$ (PrimValue -> Int64) -> [PrimValue] -> [Int64]
forall a b. (a -> b) -> [a] -> [b]
map PrimValue -> Int64
expectInt64 [PrimValue]
values
  where
    expectInt64 :: PrimValue -> Int64
expectInt64 (IntValue (Int64Value Int64
element)) = Int64
element
    expectInt64 PrimValue
_ = [Char] -> Int64
forall a. HasCallStack => [Char] -> a
error [Char]
"expected an i64 value"
toPrimitiveValue Signedness
Unsigned [Int]
shape (IntType IntType
Int8) [PrimValue]
values =
  Value -> InterpM rep Value
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> InterpM rep Value) -> Value -> InterpM rep Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Word8 -> Value
V.U8Value ([Int] -> Vector Int
shapeVector [Int]
shape) (Vector Word8 -> Value) -> Vector Word8 -> Value
forall a b. (a -> b) -> a -> b
$ [Word8] -> Vector Word8
forall a. Storable a => [a] -> Vector a
SVec.fromList ([Word8] -> Vector Word8) -> [Word8] -> Vector Word8
forall a b. (a -> b) -> a -> b
$ (PrimValue -> Word8) -> [PrimValue] -> [Word8]
forall a b. (a -> b) -> [a] -> [b]
map PrimValue -> Word8
forall {b}. Num b => PrimValue -> b
expectUInt8 [PrimValue]
values
  where
    expectUInt8 :: PrimValue -> b
expectUInt8 (IntValue (Int8Value Int8
element)) = Int8 -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int8
element
    expectUInt8 PrimValue
_ = [Char] -> b
forall a. HasCallStack => [Char] -> a
error [Char]
"expected a u8 value"
toPrimitiveValue Signedness
Unsigned [Int]
shape (IntType IntType
Int16) [PrimValue]
values =
  Value -> InterpM rep Value
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> InterpM rep Value) -> Value -> InterpM rep Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Word16 -> Value
V.U16Value ([Int] -> Vector Int
shapeVector [Int]
shape) (Vector Word16 -> Value) -> Vector Word16 -> Value
forall a b. (a -> b) -> a -> b
$ [Word16] -> Vector Word16
forall a. Storable a => [a] -> Vector a
SVec.fromList ([Word16] -> Vector Word16) -> [Word16] -> Vector Word16
forall a b. (a -> b) -> a -> b
$ (PrimValue -> Word16) -> [PrimValue] -> [Word16]
forall a b. (a -> b) -> [a] -> [b]
map PrimValue -> Word16
forall {b}. Num b => PrimValue -> b
expectUInt16 [PrimValue]
values
  where
    expectUInt16 :: PrimValue -> b
expectUInt16 (IntValue (Int16Value Int16
element)) = Int16 -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int16
element
    expectUInt16 PrimValue
_ = [Char] -> b
forall a. HasCallStack => [Char] -> a
error [Char]
"expected a u16 value"
toPrimitiveValue Signedness
Unsigned [Int]
shape (IntType IntType
Int32) [PrimValue]
values =
  Value -> InterpM rep Value
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> InterpM rep Value) -> Value -> InterpM rep Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Word32 -> Value
V.U32Value ([Int] -> Vector Int
shapeVector [Int]
shape) (Vector Word32 -> Value) -> Vector Word32 -> Value
forall a b. (a -> b) -> a -> b
$ [Word32] -> Vector Word32
forall a. Storable a => [a] -> Vector a
SVec.fromList ([Word32] -> Vector Word32) -> [Word32] -> Vector Word32
forall a b. (a -> b) -> a -> b
$ (PrimValue -> Word32) -> [PrimValue] -> [Word32]
forall a b. (a -> b) -> [a] -> [b]
map PrimValue -> Word32
forall {b}. Num b => PrimValue -> b
expectUInt32 [PrimValue]
values
  where
    expectUInt32 :: PrimValue -> b
expectUInt32 (IntValue (Int32Value Int32
element)) = Int32 -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int32
element
    expectUInt32 PrimValue
_ = [Char] -> b
forall a. HasCallStack => [Char] -> a
error [Char]
"expected a u32 value"
toPrimitiveValue Signedness
Unsigned [Int]
shape (IntType IntType
Int64) [PrimValue]
values =
  Value -> InterpM rep Value
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> InterpM rep Value) -> Value -> InterpM rep Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Word64 -> Value
V.U64Value ([Int] -> Vector Int
shapeVector [Int]
shape) (Vector Word64 -> Value) -> Vector Word64 -> Value
forall a b. (a -> b) -> a -> b
$ [Word64] -> Vector Word64
forall a. Storable a => [a] -> Vector a
SVec.fromList ([Word64] -> Vector Word64) -> [Word64] -> Vector Word64
forall a b. (a -> b) -> a -> b
$ (PrimValue -> Word64) -> [PrimValue] -> [Word64]
forall a b. (a -> b) -> [a] -> [b]
map PrimValue -> Word64
forall {b}. Num b => PrimValue -> b
expectUInt64 [PrimValue]
values
  where
    expectUInt64 :: PrimValue -> b
expectUInt64 (IntValue (Int64Value Int64
element)) = Int64 -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
element
    expectUInt64 PrimValue
_ = [Char] -> b
forall a. HasCallStack => [Char] -> a
error [Char]
"expected a u64 value"
toPrimitiveValue Signedness
_ [Int]
shape (FloatType FloatType
Float16) [PrimValue]
values =
  Value -> InterpM rep Value
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> InterpM rep Value) -> Value -> InterpM rep Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Half -> Value
V.F16Value ([Int] -> Vector Int
shapeVector [Int]
shape) (Vector Half -> Value) -> Vector Half -> Value
forall a b. (a -> b) -> a -> b
$ [Half] -> Vector Half
forall a. Storable a => [a] -> Vector a
SVec.fromList ([Half] -> Vector Half) -> [Half] -> Vector Half
forall a b. (a -> b) -> a -> b
$ (PrimValue -> Half) -> [PrimValue] -> [Half]
forall a b. (a -> b) -> [a] -> [b]
map PrimValue -> Half
expectFloat16 [PrimValue]
values
  where
    expectFloat16 :: PrimValue -> Half
expectFloat16 (FloatValue (Float16Value Half
element)) = Half
element
    expectFloat16 PrimValue
_ = [Char] -> Half
forall a. HasCallStack => [Char] -> a
error [Char]
"expected an f16 value"
toPrimitiveValue Signedness
_ [Int]
shape (FloatType FloatType
Float32) [PrimValue]
values =
  Value -> InterpM rep Value
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> InterpM rep Value) -> Value -> InterpM rep Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Float -> Value
V.F32Value ([Int] -> Vector Int
shapeVector [Int]
shape) (Vector Float -> Value) -> Vector Float -> Value
forall a b. (a -> b) -> a -> b
$ [Float] -> Vector Float
forall a. Storable a => [a] -> Vector a
SVec.fromList ([Float] -> Vector Float) -> [Float] -> Vector Float
forall a b. (a -> b) -> a -> b
$ (PrimValue -> Float) -> [PrimValue] -> [Float]
forall a b. (a -> b) -> [a] -> [b]
map PrimValue -> Float
expectFloat32 [PrimValue]
values
  where
    expectFloat32 :: PrimValue -> Float
expectFloat32 (FloatValue (Float32Value Float
element)) = Float
element
    expectFloat32 PrimValue
_ = [Char] -> Float
forall a. HasCallStack => [Char] -> a
error [Char]
"expected an f32 value"
toPrimitiveValue Signedness
_ [Int]
shape (FloatType FloatType
Float64) [PrimValue]
values =
  Value -> InterpM rep Value
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> InterpM rep Value) -> Value -> InterpM rep Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Double -> Value
V.F64Value ([Int] -> Vector Int
shapeVector [Int]
shape) (Vector Double -> Value) -> Vector Double -> Value
forall a b. (a -> b) -> a -> b
$ [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
SVec.fromList ([Double] -> Vector Double) -> [Double] -> Vector Double
forall a b. (a -> b) -> a -> b
$ (PrimValue -> Double) -> [PrimValue] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map PrimValue -> Double
expectFloat64 [PrimValue]
values
  where
    expectFloat64 :: PrimValue -> Double
expectFloat64 (FloatValue (Float64Value Double
element)) = Double
element
    expectFloat64 PrimValue
_ = [Char] -> Double
forall a. HasCallStack => [Char] -> a
error [Char]
"expected an f64 value"
toPrimitiveValue Signedness
_ [Int]
shape PrimType
Bool [PrimValue]
values =
  Value -> InterpM rep Value
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> InterpM rep Value) -> Value -> InterpM rep Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Bool -> Value
V.BoolValue ([Int] -> Vector Int
shapeVector [Int]
shape) (Vector Bool -> Value) -> Vector Bool -> Value
forall a b. (a -> b) -> a -> b
$ [Bool] -> Vector Bool
forall a. Storable a => [a] -> Vector a
SVec.fromList ([Bool] -> Vector Bool) -> [Bool] -> Vector Bool
forall a b. (a -> b) -> a -> b
$ (PrimValue -> Bool) -> [PrimValue] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map PrimValue -> Bool
expectBool [PrimValue]
values
  where
    expectBool :: PrimValue -> Bool
expectBool (BoolValue Bool
element) = Bool
element
    expectBool PrimValue
_ = [Char] -> Bool
forall a. HasCallStack => [Char] -> a
error [Char]
"expected a bool value"
toPrimitiveValue Signedness
_ [Int]
_ PrimType
Unit [PrimValue]
_ =
  Text -> InterpM rep Value
forall rep a. Text -> InterpM rep a
interpError Text
"unit values cannot be represented as external values"

shapeVector :: [Int] -> SVec.Vector Int
shapeVector :: [Int] -> Vector Int
shapeVector = [Int] -> Vector Int
forall a. Storable a => [a] -> Vector a
SVec.fromList

-- | Run a program in the SOAC IR
runSOACS :: Prog SOACS -> Name -> [V.Value] -> IO (Either T.Text [V.Value])
runSOACS :: Prog SOACS -> Name -> [Value] -> IO (Either Text [Value])
runSOACS = OpEvaluator SOACS
-> Prog SOACS -> Name -> [Value] -> IO (Either Text [Value])
forall rep.
Typed (FParamInfo rep) =>
OpEvaluator rep
-> Prog rep -> Name -> [Value] -> IO (Either Text [Value])
runProgram OpEvaluator SOACS
FunEnv SOACS -> Env -> SOAC SOACS -> InterpM SOACS [Val]
forall rep. FunEnv rep -> Env -> SOAC rep -> InterpM rep [Val]
evalSOAC

evalGPUOp :: OpEvaluator GPU
evalGPUOp :: OpEvaluator GPU
evalGPUOp FunEnv GPU
funs Env
env (SegOp SegOp SegLevel GPU
op) = FunEnv GPU -> Env -> SegOp SegLevel GPU -> InterpM GPU [Val]
forall rep level.
FunEnv rep -> Env -> SegOp level rep -> InterpM rep [Val]
evalSegOp FunEnv GPU
funs Env
env SegOp SegLevel GPU
op
evalGPUOp FunEnv GPU
_ Env
env (SizeOp SizeOp
op) = Env -> SizeOp -> InterpM GPU [Val]
forall rep. Env -> SizeOp -> InterpM rep [Val]
evalSizeOp Env
env SizeOp
op
evalGPUOp FunEnv GPU
funs Env
env (OtherOp SOAC GPU
op) = FunEnv GPU -> Env -> SOAC GPU -> InterpM GPU [Val]
forall rep. FunEnv rep -> Env -> SOAC rep -> InterpM rep [Val]
evalSOAC FunEnv GPU
funs Env
env SOAC GPU
op
evalGPUOp FunEnv GPU
funs Env
env (GPUBody [TypeBase Shape NoMode]
types Body GPU
body) = do
  values <- FunEnv GPU -> Env -> Body GPU -> InterpM GPU [Val]
forall rep. FunEnv rep -> Env -> Body rep -> InterpM rep [Val]
evalBody FunEnv GPU
funs Env
env Body GPU
body
  collectOutputs env [1] types [values]

sizeValue :: SizeClass -> Int
sizeValue :: SizeClass -> Int
sizeValue (SizeThreshold KernelPath
_ (Just Int64
n)) = Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
n
sizeValue SizeClass
SizeThreadBlock = Int
256
sizeValue SizeClass
SizeGrid = Int
65535
sizeValue SizeClass
SizeTile = Int
32
sizeValue SizeClass
SizeRegTile = Int
4
sizeValue SizeClass
SizeSharedMemory = Int
48 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1024
sizeValue SizeClass
SizeRegisters = Int
65536
sizeValue SizeClass
SizeCache = Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1024 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1024
sizeValue SizeClass
_ = Int
32768

evalSizeOp :: Env -> SizeOp -> InterpM rep [Val]
evalSizeOp :: forall rep. Env -> SizeOp -> InterpM rep [Val]
evalSizeOp Env
_ (GetSize Name
_ SizeClass
cls) =
  [Val] -> InterpM rep [Val]
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Int -> Val
int64Val (Int -> Val) -> Int -> Val
forall a b. (a -> b) -> a -> b
$ SizeClass -> Int
sizeValue SizeClass
cls]
evalSizeOp Env
_ (GetSizeMax SizeClass
cls) =
  [Val] -> InterpM rep [Val]
forall a. a -> InterpM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Int -> Val
int64Val (Int -> Val) -> Int -> Val
forall a b. (a -> b) -> a -> b
$ SizeClass -> Int
sizeValue SizeClass
cls]
evalSizeOp Env
env (CmpSizeLe Name
_ SizeClass
cls SubExp
x) = do
  limit <- Env -> SubExp -> InterpM rep Int
forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env SubExp
x
  pure [PrimVal $ BoolValue $ sizeValue cls <= limit]
evalSizeOp Env
env (CalcNumBlocks SubExp
width_exp Name
_ SubExp
block_size_exp) = do
  width <- Env -> SubExp -> InterpM rep Int
forall rep. Env -> SubExp -> InterpM rep Int
evalInt Env
env SubExp
width_exp
  block_size <- evalInt env block_size_exp
  if block_size <= 0
    then error "thread-block size must be positive"
    else
      pure
        [ int64Val $
            max 1 $
              min (sizeValue SizeGrid) $
                (width + block_size - 1) `div` block_size
        ]

-- | Run a program in the GPU IR.
runGPU :: Prog GPU -> Name -> [V.Value] -> IO (Either T.Text [V.Value])
runGPU :: Prog GPU -> Name -> [Value] -> IO (Either Text [Value])
runGPU = OpEvaluator GPU
-> Prog GPU -> Name -> [Value] -> IO (Either Text [Value])
forall rep.
Typed (FParamInfo rep) =>
OpEvaluator rep
-> Prog rep -> Name -> [Value] -> IO (Either Text [Value])
runProgram OpEvaluator GPU
evalGPUOp