{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE TypeApplications           #-}
{-# LANGUAGE TypeFamilies               #-}
{-# OPTIONS_HADDOCK hide #-}
-- |
-- Module      : Data.Array.Accelerate.LLVM.Execute.Async
-- Copyright   : [2014..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--

module Data.Array.Accelerate.LLVM.Execute.Async
  where

import Data.Array.Accelerate.LLVM.State                             ( LLVM )
import Data.Array.Accelerate.Representation.Array
import Data.Array.Accelerate.Representation.Type

import GHC.Stack


class Monad (Par arch) => Async arch where

  -- | The monad parallel computations will be executed in. Presumably a stack
  -- with the LLVM monad at the base.
  --
  data Par arch :: * -> *

  -- | Parallel computations can communicate via futures.
  --
  type FutureR arch :: * -> *

  -- | Create a new (empty) promise, to be fulfilled at some future point.
  --
  new :: HasCallStack => Par arch (FutureR arch a)

  -- | The future is here. Multiple 'put's to the same future are not allowed
  -- and (presumably) result in a runtime error.
  --
  put :: HasCallStack => FutureR arch a -> a -> Par arch ()

  -- | Read the value stored in a future, once it is available. It is _not_
  -- required that this is a blocking operation on the host, only that it is
  -- blocking with respect to computations on the remote device.
  --
  get :: HasCallStack => FutureR arch a -> Par arch a

  -- | Fork a computation to happen in parallel. The forked computation may
  -- exchange values with other computations using Futures.
  --
  fork :: HasCallStack => Par arch () -> Par arch ()

  -- | Lift an operation from the base LLVM monad into the Par monad
  --
  liftPar :: HasCallStack => LLVM arch a -> Par arch a

  -- | Read a value stored in a future, once it is available. This is blocking
  -- with respect to both the host and remote device.
  --
  {-# INLINEABLE block #-}
  block :: HasCallStack => FutureR arch a -> Par arch a
  block = FutureR arch a -> Par arch a
forall arch a.
(Async arch, HasCallStack) =>
FutureR arch a -> Par arch a
get

  -- | Evaluate a computation in a new thread/context. This might be implemented
  -- more efficiently than the default implementation.
  --
  {-# INLINEABLE spawn #-}
  spawn :: HasCallStack => Par arch a -> Par arch a
  spawn Par arch a
m = do
    FutureR arch a
r <- Par arch (FutureR arch a)
forall arch a.
(Async arch, HasCallStack) =>
Par arch (FutureR arch a)
new
    Par arch () -> Par arch ()
forall arch.
(Async arch, HasCallStack) =>
Par arch () -> Par arch ()
fork (Par arch () -> Par arch ()) -> Par arch () -> Par arch ()
forall a b. (a -> b) -> a -> b
$ FutureR arch a -> a -> Par arch ()
forall arch a.
(Async arch, HasCallStack) =>
FutureR arch a -> a -> Par arch ()
put FutureR arch a
r (a -> Par arch ()) -> Par arch a -> Par arch ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Par arch a
m
    FutureR arch a -> Par arch a
forall arch a.
(Async arch, HasCallStack) =>
FutureR arch a -> Par arch a
get FutureR arch a
r

  -- | Create a new "future" where the value is available immediately. This
  -- might be implemented more efficiently than the default implementation.
  --
  {-# INLINEABLE newFull #-}
  newFull :: HasCallStack => a -> Par arch (FutureR arch a)
  newFull a
a = do
    FutureR arch a
r <- Par arch (FutureR arch a)
forall arch a.
(Async arch, HasCallStack) =>
Par arch (FutureR arch a)
new
    FutureR arch a -> a -> Par arch ()
forall arch a.
(Async arch, HasCallStack) =>
FutureR arch a -> a -> Par arch ()
put FutureR arch a
r a
a
    FutureR arch a -> Par arch (FutureR arch a)
forall (m :: * -> *) a. Monad m => a -> m a
return FutureR arch a
r

type family FutureArraysR arch arrs where
  FutureArraysR arch ()           = ()
  FutureArraysR arch (a, b)       = (FutureArraysR arch a, FutureArraysR arch b)
  FutureArraysR arch (Array sh e) = FutureR arch (Array sh e)

getArrays :: Async arch => ArraysR a -> FutureArraysR arch a -> Par arch a
getArrays :: ArraysR a -> FutureArraysR arch a -> Par arch a
getArrays (TupRsingle ArrayR{}) FutureArraysR arch a
a        = FutureR arch a -> Par arch a
forall arch a.
(Async arch, HasCallStack) =>
FutureR arch a -> Par arch a
get FutureArraysR arch a
FutureR arch a
a
getArrays ArraysR a
TupRunit              FutureArraysR arch a
_        = () -> Par arch ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
getArrays (TupRpair TupR ArrayR a1
r1 TupR ArrayR b
r2)      (a1, a2) = (,) (a1 -> b -> (a1, b)) -> Par arch a1 -> Par arch (b -> (a1, b))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TupR ArrayR a1 -> FutureArraysR arch a1 -> Par arch a1
forall arch a.
Async arch =>
ArraysR a -> FutureArraysR arch a -> Par arch a
getArrays TupR ArrayR a1
r1 FutureArraysR arch a1
a1 Par arch (b -> (a1, b)) -> Par arch b -> Par arch (a1, b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TupR ArrayR b -> FutureArraysR arch b -> Par arch b
forall arch a.
Async arch =>
ArraysR a -> FutureArraysR arch a -> Par arch a
getArrays TupR ArrayR b
r2 FutureArraysR arch b
a2

blockArrays :: Async arch => ArraysR a -> FutureArraysR arch a -> Par arch a
blockArrays :: ArraysR a -> FutureArraysR arch a -> Par arch a
blockArrays (TupRsingle ArrayR{}) FutureArraysR arch a
a        = FutureR arch a -> Par arch a
forall arch a.
(Async arch, HasCallStack) =>
FutureR arch a -> Par arch a
block FutureArraysR arch a
FutureR arch a
a
blockArrays ArraysR a
TupRunit              FutureArraysR arch a
_        = () -> Par arch ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
blockArrays (TupRpair TupR ArrayR a1
r1 TupR ArrayR b
r2)      (a1, a2) = (,) (a1 -> b -> (a1, b)) -> Par arch a1 -> Par arch (b -> (a1, b))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TupR ArrayR a1 -> FutureArraysR arch a1 -> Par arch a1
forall arch a.
Async arch =>
ArraysR a -> FutureArraysR arch a -> Par arch a
blockArrays TupR ArrayR a1
r1 FutureArraysR arch a1
a1 Par arch (b -> (a1, b)) -> Par arch b -> Par arch (a1, b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TupR ArrayR b -> FutureArraysR arch b -> Par arch b
forall arch a.
Async arch =>
ArraysR a -> FutureArraysR arch a -> Par arch a
blockArrays TupR ArrayR b
r2 FutureArraysR arch b
a2

-- | Create new (empty) promises for a structure of arrays, to be fulfilled
-- at some future point. Note that the promises in the structure may all be
-- fullfilled at different moments.
--
newArrays :: Async arch => ArraysR a -> Par arch (FutureArraysR arch a)
newArrays :: ArraysR a -> Par arch (FutureArraysR arch a)
newArrays ArraysR a
TupRunit               = () -> Par arch ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
newArrays (TupRsingle ArrayR{})  = Par arch (FutureArraysR arch a)
forall arch a.
(Async arch, HasCallStack) =>
Par arch (FutureR arch a)
new
newArrays (TupRpair TupR ArrayR a1
repr1 TupR ArrayR b
repr2) = (,) (FutureArraysR arch a1
 -> FutureArraysR arch b
 -> (FutureArraysR arch a1, FutureArraysR arch b))
-> Par arch (FutureArraysR arch a1)
-> Par
     arch
     (FutureArraysR arch b
      -> (FutureArraysR arch a1, FutureArraysR arch b))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TupR ArrayR a1 -> Par arch (FutureArraysR arch a1)
forall arch a.
Async arch =>
ArraysR a -> Par arch (FutureArraysR arch a)
newArrays TupR ArrayR a1
repr1 Par
  arch
  (FutureArraysR arch b
   -> (FutureArraysR arch a1, FutureArraysR arch b))
-> Par arch (FutureArraysR arch b)
-> Par arch (FutureArraysR arch a1, FutureArraysR arch b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TupR ArrayR b -> Par arch (FutureArraysR arch b)
forall arch a.
Async arch =>
ArraysR a -> Par arch (FutureArraysR arch a)
newArrays TupR ArrayR b
repr2