Execute.hs

{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}

{-|
The execute phase.

The single-queue fast path and the schedule-driven multi-queue path with
its 'QueueBackend' seam.
-}
module Fragr.Execute
  ( execute
  , executeQueued
  , QueueBackend (..)
  , executingQueues
  ) where

import Control.Exception (throwIO)
import Control.Monad (void, when)
import Control.Monad.IO.Class (MonadIO (..))
import Data.Coerce (coerce)
import Data.Foldable (for_)
import Data.IORef (readIORef, writeIORef)
import Data.IntMap.Strict qualified as IntMap
import Data.IntSet qualified as IntSet
import Data.Word (Word64)
import Type.Reflection

import Fragr.Compile (canExecute)
import Fragr.Error (FragrError (..))
import Fragr.Graph (FrameGraph (..), PassNode (..), ResourceEntry (..), Resources (..), SomeEntry (..), entryAt, entryOf, requireCompiled)
import Fragr.Recycle (RecycleQueue, collect, mkRetireItem, retireItem)
import Fragr.Resource (Access (..), Resource (..))
import Fragr.Sync (Compiled (..), PassSync (..), Transfer (..))
import Fragr.Types (Handle, QueueId (..), SomeHandle (..))

{- |
Walk the passes in registration order, skipping culled ones. Per
surviving pass:

1. materialize the transients it creates;
2. run the 'preRead' / 'preWrite' hooks for accesses with flags, then
   the 'Fragr.Graph.addPreExec' flushes;
3. invoke its execution callback;
4. fire the 'Fragr.Graph.addPostExec' flushes;
5. destroy every transient whose last executing user it is.

@ctx@ is forwarded to execution callbacks and hooks, @alloc@ to
'createResource' / 'destroyResource'.
-}
{-# INLINE execute #-}
execute :: (MonadIO m) => FrameGraph ctx alloc -> ctx -> alloc -> m ()
execute g ctx alloc = liftIO do
  compiled <- requireCompiled g "execute"
  passes <- readIORef g.passesRef
  for_ passes \p ->
    when (canExecute compiled.passRefs p) do
      runPass g ctx alloc Nothing p
      for_ (IntMap.findWithDefault [] p.passId compiled.retireAfter) \e ->
        entryAt g e >>= release alloc

{- |
The seam through which 'executeQueued' hands the schedule to the
application. The library computes 'PassSync' values and calls these back;
it never interprets a queue, timeline value or event itself.
-}
data QueueBackend = QueueBackend
  { beforePass :: PassSync -> IO ()
  {- ^ before a pass runs: wait the listed timeline values and events, and
  acquire ownership of the listed resources
  -}
  , afterPass :: PassSync -> IO ()
  {- ^ after a pass runs: release the listed resources, signal the listed
  events, then signal this pass's timeline value
  -}
  , invoke :: PassSync -> IO () -> IO ()
  {- ^ invocation control over one pass's whole step ('beforePass', the
  hook/run sequence, 'afterPass'): run it now (@\_ body -> body@, a device
  queue recording commands) or stash it to run later (a host queue executed
  after the submits, against real timeline values).

  Deferring moves the pass's 'createResource' calls with it, so a deferred
  queue must produce nothing an inline-invoked pass consumes — the consumer
  runs first and finds the resource unmaterialized. Reclamation is safe by
  construction: it goes through the 'RecycleQueue', whose 'completed' values
  gate the release behind the deferred pass's own timeline signal.
  -}
  , completed :: IO [(QueueId, Word64)]
  {- ^ the currently-reached timeline value per queue, consulted to decide
  what the recycle queue may reclaim
  -}
  }

{- |
Multi-queue execution: 'execute' driving the compiled schedule through
the 'QueueBackend' instead of destroying transients inline. Per surviving
pass:

1. hand its step to 'invoke' — 'beforePass', then materialize plus the
   full hook sequence of 'runPass' (acquire / release sides included),
   then 'afterPass';
2. retire any transient whose last executing user this pass is;
3. 'collect' the recycle queue.

An import-only graph (every resource owned outside the graph) never retires
anything and may pass 'Nothing' for the recycle queue; a graph that does
need to reclaim a transient then fails upfront with 'RecycleQueueRequired'.
-}
{-# INLINE executeQueued #-}
executeQueued
  :: (MonadIO m)
  => FrameGraph ctx alloc
  -> QueueBackend
  -> Maybe RecycleQueue
  -> ctx
  -> alloc
  -> m ()
executeQueued g backend mrq ctx alloc = liftIO do
  compiled <- requireCompiled g "executeQueued"
  -- The retire schedule holds transients only, so any entry at all
  -- demands a recycle queue.
  case (mrq, concat (IntMap.elems compiled.retireAfter)) of
    (Nothing, e : _) -> do
      entry <- entryAt g e
      throwIO $ RecycleQueueRequired entry.name
    _ -> pure ()
  passes <- readIORef g.passesRef
  for_ passes \p ->
    when (canExecute compiled.passRefs p) do
      psync <- case IntMap.lookup p.passId compiled.passSync of
        Nothing -> throwIO $ InternalInvariant "no schedule for an executing pass"
        Just s -> pure s
      backend.invoke psync do
        backend.beforePass psync
        runPass g ctx alloc (Just psync) p
        backend.afterPass psync
      for_ mrq \rq -> do
        for_ (IntMap.findWithDefault [] p.passId compiled.retireAfter) \e -> do
          entry <- entryAt g e
          let reqs = IntMap.findWithDefault [] e compiled.entryRetire
          item <- mkRetireItem e reqs (release alloc entry)
          retireItem rq item
        backend.completed >>= void . collect rq
  for_ mrq \rq -> backend.completed >>= void . collect rq

{- |
The queues with at least one executing pass on the compiled schedule,
ascending. Backends size per-queue state (command buffers, timelines) from
it. Fatal before 'Fragr.Compile.compile'.
-}
{-# INLINEABLE executingQueues #-}
executingQueues :: (MonadIO m) => FrameGraph ctx alloc -> m [QueueId]
executingQueues g = liftIO do
  c <- requireCompiled g "executingQueues"
  pure $ coerce . IntSet.toAscList . IntSet.fromList . coerce $ (map (.queue) (IntMap.elems c.passSync))

{- | The per-pass body shared by 'execute' and 'executeQueued': materialize
the pass's creates, fire the acquire / read / write hooks for accesses with
flags, fire the 'Fragr.Graph.addPreExec' flushes, run the callback, fire the
release hooks, then the 'Fragr.Graph.addPostExec' flushes. The acquire and
release sides exist only under a schedule ('executeQueued').
-}
runPass :: forall ctx alloc. FrameGraph ctx alloc -> ctx -> alloc -> Maybe PassSync -> PassNode ctx alloc -> IO ()
runPass g ctx alloc msync p = do
  for_ p.creates \(SomeHandle h) -> do
    entry <- entryOf g h
    when entry.imported $
      throwIO $
        InternalInvariant "create on an imported entry"
    materialize alloc entry
  for_ (maybe [] (.acquires) msync) \(Transfer h peer flags) -> fire (\hh d fl -> preAcquire hh d fl peer) h flags
  for_ p.reads \(Access h flags) -> fire preRead h flags
  for_ p.writes \(Access h flags) -> fire preWrite h flags
  preExecs <- readIORef g.preExecRef
  for_ preExecs ($ ctx)
  p.run Resources{graph = g, pass = p, ctx}
  for_ (maybe [] (.releases) msync) \(Transfer h peer flags) -> fire (\hh d fl -> preRelease hh d fl peer) h flags
  postExecs <- readIORef g.postExecRef
  for_ postExecs ($ ctx)
  where
    fire
      :: forall r
       . (Resource r)
      => (forall s. (Resource s, Alloc s ~ alloc, Ctx s ~ ctx) => Handle s -> Desc s -> Flags s -> Ctx s -> s -> IO ())
      -> Handle r
      -> Maybe (Flags r)
      -> IO ()
    fire hook h mflags = for_ mflags run
      where
        run :: Flags r -> IO ()
        run flags = do
          entry <- entryOf g h
          case entry.payload of
            SomeEntry rep desc objRef -> case eqTypeRep rep (typeRep @r) of
              Nothing -> throwIO $ TypeMismatch "hook" (SomeTypeRep rep) (SomeTypeRep (typeRep @r))
              Just HRefl ->
                readIORef objRef >>= \case
                  Nothing -> throwIO $ InternalInvariant "hook before create"
                  Just obj -> hook h desc flags ctx obj

materialize :: alloc -> ResourceEntry ctx alloc -> IO ()
materialize alloc entry = case entry.payload of
  SomeEntry _ desc objRef -> do
    obj <- createResource desc alloc
    writeIORef objRef (Just obj)

release :: alloc -> ResourceEntry ctx alloc -> IO ()
release alloc entry = case entry.payload of
  SomeEntry _ desc objRef ->
    readIORef objRef >>= \case
      Nothing -> throwIO $ InternalInvariant "destroy before create"
      Just obj -> do
        destroyResource desc alloc obj
        writeIORef objRef Nothing