Dot.hs

{-|
Graphviz DOT exporters for compiled frame graphs.

@
FG.compile graph
dotSource <- Dot.dump graph
@

Two views: 'dump' renders the dataflow — passes and resource versions —
and 'dumpSync' the compiled schedule as queue lanes ('renderSync').

Imports sharing a dotted name prefix (e.g. mip levels @img.mip0@,
@img.mip1@) render as one record node with a cell per member, keeping the
family together in an array instead of letting the layout scatter it. The
grouping covers the version-1 identities; versions minted by writes render
in the dataflow as usual.

Any other visualization can be written the same way, against the
read-only 'Snapshot' from 'FG.snapshot'.
-}
module Fragr.Snapshot.Dot
  ( dump
  , dumpWith
  , render
  , Options (..)
  , defaultOptions

    -- * The schedule view
  , dumpSync
  , renderSync
  ) where

import Control.Monad (guard)
import Data.IntMap.Strict (IntMap)
import Data.IntMap.Strict qualified as IntMap
import Data.List (nub)
import Data.Map.Strict qualified as Map
import Data.Maybe (maybeToList)
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Word (Word64)

import Fragr (Access (..), EntryInfo (..), EventId (..), FrameGraph, NodeInfo (..), PassInfo (..), PassSync (..), QueueId (..), RetireInfo (..), Snapshot (..), SyncEvent (..), Transfer (..), Wait (..), accessId, handleId, snapshot, someHandleId, transferId)

dump :: FrameGraph ctx alloc -> IO Text
dump = dumpWith defaultOptions

dumpWith :: Options -> FrameGraph ctx alloc -> IO Text
dumpWith opts g = render opts <$> snapshot g

dumpSync :: FrameGraph ctx alloc -> IO Text
dumpSync g = renderSync <$> snapshot g

data Options = Options
  { clusterImports :: Bool
  {- ^ Fence the initial import versions in a dotted @Imported@ frame,
  instead of letting them float near their consumers.
  -}
  , stratify :: Bool
  {- ^ Pin each dependency level ('Fragr.Snapshot.PassInfo') to one rank, so
  rows are levels and the row count is the critical path. Culled passes stay
  unpinned.
  -}
  , antiEdges :: Bool
  {- ^ Overlay the write-after-read edges ('Fragr.Snapshot.PassInfo') as
  dashed reader -> renamer edges. Drawn without layout constraint, so they
  annotate the dataflow rather than reshape it.
  -}
  }
  deriving stock (Eq, Show)

defaultOptions :: Options
defaultOptions =
  Options
    { clusterImports = False
    , stratify = False
    , antiEdges = True
    }

render :: Options -> Snapshot -> Text
render opts s =
  Text.unlines $
    concat
      [ ["digraph FrameGraph {"]
      , map indent header
      , map (indent . passVertex) s.passes
      , map (indent . nodeVertex) (filter standalone s.nodes)
      , map (indent . familyVertex) importFamilies
      , concatMap (map indent . writeEdges) s.passes
      , concatMap (map indent . readEdges) s.nodes
      , map indent antiEdgeLines
      , concatMap (map indent . createCluster) s.passes
      , map indent importedCluster
      , map indent levelRanks
      , ["}"]
      ]
  where
    header =
      [ "graph [rankdir=TB, splines=spline" <> (if opts.stratify then ", newrank=true" else "") <> "]"
      , "node [shape=record, style=\"rounded,filled\", fontname=\"helvetica\", fontsize=10]"
      ]

    entryById :: IntMap EntryInfo
    entryById = IntMap.fromList do
      e <- s.entries
      pure (e.entryId, e)

    nodeById :: IntMap NodeInfo
    nodeById = IntMap.fromList do
      n <- s.nodes
      pure (n.nodeId, n)

    passKey :: PassInfo -> Text
    passKey p = passIdKey p.passId

    passIdKey :: Int -> Text
    passIdKey i = "P" <> tshow i

    nodeKey :: NodeInfo -> Text
    nodeKey n = "R" <> tshow n.resourceId <> "_" <> tshow n.version

    handleKey :: Int -> Text
    handleKey i = nodeRef (nodeById IntMap.! i)

    -- Where edges attach: a family member's port, or the node's own vertex.
    nodeRef :: NodeInfo -> Text
    nodeRef n = IntMap.findWithDefault (nodeKey n) n.nodeId familyPorts

    standalone :: NodeInfo -> Bool
    standalone n = not (IntMap.member n.nodeId familyPorts)

    -- Imported v1 identities, the Imported cluster's population.
    importedV1 :: [NodeInfo]
    importedV1 = do
      n <- s.nodes
      guard (n.version == 1)
      guard (entryById IntMap.! n.resourceId).imported
      pure n

    -- Imports sharing a name prefix before the last dot, two or more
    -- strong; singleton prefixes and undotted names stay standalone.
    -- Transients never group: they live inside per-pass create clusters,
    -- and one record node cannot span two clusters.
    importFamilies :: [Family]
    importFamilies = do
      (prefix, members@(m : _ : _)) <- Map.toAscList byPrefix
      pure Family{key = "G" <> tshow m.node.nodeId, prefix, members}
      where
        byPrefix = Map.fromListWith (flip (<>)) do
          n <- importedV1
          (prefix, suffix) <- maybeToList (familyOf n.name)
          pure (prefix, [Member{node = n, suffix}])

    familyOf :: Text -> Maybe (Text, Text)
    familyOf name = do
      let (pre, suffix) = Text.breakOnEnd "." name
      guard (not (Text.null pre))
      pure (Text.dropEnd 1 pre, suffix)

    familyPorts :: IntMap Text
    familyPorts = IntMap.fromList do
      f <- importFamilies
      m <- f.members
      pure (m.node.nodeId, f.key <> ":" <> portName m)

    familyVertex f =
      f.key
        <> " [label=\"{"
        <> esc f.prefix
        <> "|{"
        <> Text.intercalate "|" (map cell f.members)
        <> "}}\", fillcolor=lightsteelblue]"
      where
        cell m = "<" <> portName m <> "> " <> esc m.suffix

    passVertex p =
      passKey p
        <> " [label=\"{"
        <> esc p.name
        <> (if p.sideEffect then " *" else "")
        <> "|refs: "
        <> tshow p.refCount
        <> ", id: "
        <> tshow p.passId
        <> "}\", fillcolor="
        <> (if p.canExecute then "orange" else "lightgray")
        <> "]"

    nodeVertex n =
      nodeKey n
        <> " [label=\"{"
        <> esc n.name
        <> (if n.version > 1 then " v" <> tshow n.version else "")
        <> (if Text.null desc then "" else "|" <> esc desc)
        <> "|id: "
        <> tshow n.resourceId
        <> ", refs: "
        <> tshow n.refCount
        <> "}\", fillcolor="
        <> fill
        <> "]"
      where
        entry = entryById IntMap.! n.resourceId
        desc = entry.description
        fill
          | not entry.imported = "skyblue"
          | n.version == 1 = "lightsteelblue"
          | otherwise = "steelblue"

    writeEdges p = fan (passKey p) (map (handleKey . accessId) p.writes) "orangered"

    -- Node id -> keys of the passes reading it, in registration order, one
    -- per pass.
    readersOfNode :: IntMap [Text]
    readersOfNode = IntMap.fromListWith (flip (<>)) do
      p <- s.passes
      i <- nub (map accessId p.reads)
      pure (i, [passKey p])

    readEdges n = fan (nodeRef n) (IntMap.findWithDefault [] n.nodeId readersOfNode) "yellowgreen"

    antiEdgeLines
      | not opts.antiEdges = []
      | otherwise = do
          r <- s.passes
          w <- r.antiAfter
          pure (passKey r <> " -> " <> passIdKey w <> " [style=dashed, color=slategray, constraint=false]")

    {- One rank per row of the layering, imports first: a level's passes sit
    on rank @2*level+1@ and the nodes they produce on the row below, so
    every edge points strictly downward and dot never has to route a flat
    one. Pinning the nodes as well as the passes is what keeps that true —
    left free, an import lands on a pass's rank and its edge is dropped.
    Culled passes and their nodes stay unpinned; @newrank@ lets a rankset
    cross the create clusters instead of emptying them.
    -}
    levelRanks
      | not opts.stratify = []
      | otherwise = do
          (_rank, keys) <- IntMap.toAscList byRank
          pure ("{ rank=same; " <> Text.unwords keys <> " }")
      where
        byRank = IntMap.fromListWith (flip (<>)) (imports <> passes <> produced)
        imports = map (\k -> (0, [k])) importedKeys
        passes = do
          p <- s.passes
          l <- maybeToList p.level
          pure (2 * l + 1, [passKey p])
        produced = do
          n <- filter standalone s.nodes
          l <- maybeToList (IntMap.lookup n.nodeId levelOfNode)
          pure (2 * l + 2, [nodeKey n])

        -- A node sits on the row below the pass that produces it: the one
        -- that writes it, or, for a transient nobody writes, the one that
        -- created it. Imports and culled passes' nodes stay unpinned.
        levelOfNode :: IntMap Int
        levelOfNode = IntMap.fromList do
          p <- s.passes
          l <- maybeToList p.level
          i <- map accessId p.writes <> map someHandleId p.creates
          pure (i, l)

    fan from targets color
      | null targets = []
      | otherwise =
          [from <> " -> { " <> Text.unwords targets <> " } [color=" <> color <> "]"]

    createCluster p
      | null p.creates = []
      | otherwise =
          [ "subgraph cluster_" <> passKey p <> " {"
          , indent "style=rounded"
          , indent (Text.unwords (passKey p : map (handleKey . someHandleId) p.creates))
          , "}"
          ]

    importedCluster
      | not opts.clusterImports = []
      | null importedKeys = []
      | otherwise =
          [ "subgraph cluster_imported {"
          , indent "style=dotted"
          , indent "label=\"Imported\""
          , indent (Text.unwords importedKeys)
          , "}"
          ]

    -- The vertices carrying the initial import versions: a family's record
    -- node stands for its members, which have no vertex of their own.
    importedKeys :: [Text]
    importedKeys = standaloneKeys <> map (.key) importFamilies
      where
        standaloneKeys = do
          n <- importedV1
          guard (standalone n)
          pure (nodeKey n)

{- |
Render the compiled schedule as queue lanes: one cluster per queue with
its executing passes in submission order, overlaid with the cross-queue
timeline waits, split-barrier event pairs, ownership transfers and retire
points. Where 'render' answers what data flows, this answers what the
'Fragr.Execute.QueueBackend' will be asked to do.
-}
renderSync :: Snapshot -> Text
renderSync s =
  Text.unlines $
    concat
      [ ["digraph FrameGraphSync {"]
      , map indent header
      , concatMap (map indent) lanes
      , map indent waitEdges
      , map indent eventEdges
      , map indent transferEdges
      , map indent retireLines
      , ["}"]
      ]
  where
    header =
      [ "graph [rankdir=LR, splines=spline]"
      , "node [shape=record, style=\"rounded,filled\", fontname=\"helvetica\", fontsize=10]"
      ]

    scheduled :: [(PassInfo, PassSync)]
    scheduled = do
      p <- s.passes
      ps <- maybeToList p.sync
      pure (p, ps)

    nodeById :: IntMap NodeInfo
    nodeById = IntMap.fromList do
      n <- s.nodes
      pure (n.nodeId, n)

    passKey :: PassInfo -> Text
    passKey p = "P" <> tshow p.passId

    -- One cluster per queue, its passes chained in submission order.
    lanes = do
      (q, lane) <- Map.toAscList byQueue
      pure $
        concat
          [
            [ "subgraph cluster_q" <> tshow q <> " {"
            , indent "style=rounded"
            , indent ("label=\"queue " <> tshow q <> "\"")
            ]
          , map (indent . laneVertex) lane
          , map indent (chain (map (passKey . fst) lane))
          , ["}"]
          ]
    byQueue = Map.fromListWith (flip (<>)) do
      pps@(_pass, ps) <- scheduled
      let QueueId q = ps.queue
      pure (q, [pps])
    laneVertex (p, ps) =
      passKey p <> " [label=\"{" <> esc p.name <> "|signal " <> tshow ps.signal <> "}\", fillcolor=orange]"
    chain keys = do
      (a, b) <- zip keys (drop 1 keys)
      pure (a <> " -> " <> b <> " [color=gray, weight=100]")

    -- (queue, signaled value) -> the signaling pass.
    signaler :: Map.Map (Int, Word64) Text
    signaler = Map.fromList do
      (p, ps) <- scheduled
      let QueueId q = ps.queue
      pure ((q, ps.signal), passKey p)

    waitEdges = do
      (p, ps) <- scheduled
      w <- ps.waits
      let QueueId wq = w.queue
      src <- maybeToList (Map.lookup (wq, w.value) signaler)
      let covers = Text.intercalate ", " (map accessText w.covers)
      pure $
        src
          <> " -> "
          <> passKey p
          <> " [label=\">="
          <> tshow w.value
          <> (if Text.null covers then "" else "\\n" <> escPlain covers)
          <> "\", color=royalblue]"

    eventEnd pick = IntMap.fromList do
      (p, ps) <- scheduled
      se <- pick ps
      let EventId e = se.event
      pure (e, passKey p)
    eventEdges = do
      (e, src) <- IntMap.toList (eventEnd (.signalEvents))
      dst <- maybeToList (IntMap.lookup e (eventEnd (.waitEvents)))
      pure (src <> " -> " <> dst <> " [label=\"e" <> tshow e <> "\", style=dashed, color=slategray]")

    -- A release edge finds its acquirers by node and source queue; a
    -- fan-out read acquires on every consumer.
    acquirers :: Map.Map (Int, Int) [(Int, Text)]
    acquirers = Map.fromListWith (flip (<>)) do
      (p, ps) <- scheduled
      t@(Transfer _h (QueueId src) _flags) <- ps.acquires
      let QueueId consumer = ps.queue
      pure ((transferId t, src), [(consumer, passKey p)])
    transferEdges = do
      (p, ps) <- scheduled
      t@(Transfer h (QueueId peer) flags) <- ps.releases
      let QueueId srcQ = ps.queue
      (consumer, dst) <- Map.findWithDefault [] (transferId t, srcQ) acquirers
      guard (consumer == peer)
      pure $
        passKey p
          <> " -> "
          <> dst
          <> " [label=\""
          <> escPlain (accessText (Access h flags))
          <> "\", color=purple]"

    -- Retire points: a note on the last executing pass touching the
    -- entry, carrying the per-queue timeline requirements.
    entryOfNode = IntMap.fromList do
      n <- s.nodes
      pure (n.nodeId, n.resourceId)
    entryName = IntMap.fromList do
      n <- s.nodes
      pure (n.resourceId, n.name)
    -- 'IntMap.fromList' keeps the last occurrence, and 'scheduled' walks
    -- in execution order: the last executing toucher per entry.
    lastToucher :: IntMap Text
    lastToucher = IntMap.fromList do
      (p, _ps) <- scheduled
      i <- map someHandleId p.creates <> map accessId (p.reads <> p.writes)
      e <- maybeToList (IntMap.lookup i entryOfNode)
      pure (e, passKey p)
    retireLines = concat do
      r <- s.retires
      src <- maybeToList (IntMap.lookup r.entryId lastToucher)
      let
        key = "T" <> tshow r.entryId
        reqs = Text.intercalate ", " do
          (QueueId q, v) <- r.requirements
          pure ("q" <> tshow q <> ">=" <> tshow v)
        rname = IntMap.findWithDefault (tshow r.entryId) r.entryId entryName
      pure
        [ key <> " [shape=note, label=\"retire " <> escPlain rname <> "\\n" <> escPlain reqs <> "\", fillcolor=khaki]"
        , src <> " -> " <> key <> " [style=dotted, color=gray50, arrowhead=none]"
        ]

    accessText :: Access -> Text
    accessText (Access h flags) = name <> maybe "" (\f -> " @" <> tshow f) flags
      where
        name = maybe (tshow (handleId h)) nodeText (IntMap.lookup (handleId h) nodeById)
        nodeText n = n.name <> (if n.version > 1 then " v" <> tshow n.version else "")

-- | One rendered import family: its vertex key, shared prefix and members.
data Family = Family
  { key :: Text
  , prefix :: Text
  , members :: [Member]
  }

-- | One family member: the imported v1 node and its name's suffix.
data Member = Member
  { node :: NodeInfo
  , suffix :: Text
  }

-- | The record port a member's edges attach to.
portName :: Member -> Text
portName m = "n" <> tshow m.node.nodeId

indent :: Text -> Text
indent = ("  " <>)

esc :: Text -> Text
esc = Text.concatMap \case
  '\\' -> "\\\\"
  '"' -> "\\\""
  '{' -> "\\{"
  '}' -> "\\}"
  '|' -> "\\|"
  '<' -> "\\<"
  '>' -> "\\>"
  c -> Text.singleton c

-- | Escape for plain (non-record) labels; record cells use 'esc'.
escPlain :: Text -> Text
escPlain = Text.concatMap \case
  '\\' -> "\\\\"
  '"' -> "\\\""
  c -> Text.singleton c

tshow :: (Show a) => a -> Text
tshow = Text.pack . show