    {-# LANGUAGE TypeFamilies #-}
    {-# LANGUAGE UndecidableInstances #-}

    module Futhark.IR.GPU.Op
      ( -- * Size operations
        SizeOp (..),

        -- * Host operations
        HostOp (..),
        traverseHostOpStms,
        typeCheckHostOp,

        -- * SegOp refinements
        SegLevel (..),
        segVirt,
        SegVirt (..),
        SegSeqDims (..),
        KernelGrid (..),

        -- * Reexports
        module Futhark.IR.GPU.Sizes,
        module Futhark.IR.SegOp,
      )
    where

    import Control.Monad
    import Data.Sequence qualified as SQ
    import Data.Text qualified as T
    import Futhark.Analysis.Alias qualified as Alias
    import Futhark.Analysis.Metrics
    import Futhark.Analysis.SymbolTable qualified as ST
    import Futhark.IR
    import Futhark.IR.Aliases (Aliases, CanBeAliased (..))
    import Futhark.IR.GPU.Sizes
    import Futhark.IR.Mem (OpReturns (..), extReturns)
    import Futhark.IR.Prop.Aliases
    import Futhark.IR.SegOp
    import Futhark.IR.TypeCheck qualified as TC
    import Futhark.Optimise.Simplify.Engine qualified as Engine
    import Futhark.Optimise.Simplify.Rep
    import Futhark.Transform.Rename
    import Futhark.Transform.Substitute
    import Futhark.Util.Pretty
      ( commasep,
        parens,
        ppTuple',
        pretty,
        (<+>),
      )
    import Futhark.Util.Pretty qualified as PP

    -- | These dimensions (indexed from 0, outermost) of the corresponding
    -- 'SegSpace' should not be parallelised, but instead iterated
    -- sequentially.  For example, with a 'SegSeqDims' of @[0]@ and a
    -- 'SegSpace' with dimensions @[n][m]@, there will be an outer loop
    -- with @n@ iterations, while the @m@ dimension will be parallelised.
    --
    -- Semantically, this has no effect, but it may allow reductions in
    -- memory usage or other low-level optimisations.  Operationally, the
    -- guarantee is that for a SegSeqDims of e.g. @[i,j,k]@, threads
    -- running at any given moment will always have the same indexes along
    -- the dimensions specified by @[i,j,k]@.
    --
    -- At the moment, this is only supported for 'SegNoVirtFull'
    -- intra-block parallelism in GPU code, as we have not yet found it
    -- useful anywhere else.
    newtype SegSeqDims = SegSeqDims {SegSeqDims -> [Int]
    segSeqDims :: [Int]}
      deriving (SegSeqDims -> SegSeqDims -> Bool
    (SegSeqDims -> SegSeqDims -> Bool)
    -> (SegSeqDims -> SegSeqDims -> Bool) -> Eq SegSeqDims
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: SegSeqDims -> SegSeqDims -> Bool
    == :: SegSeqDims -> SegSeqDims -> Bool
    $c/= :: SegSeqDims -> SegSeqDims -> Bool
    /= :: SegSeqDims -> SegSeqDims -> Bool
    Eq, Eq SegSeqDims
    Eq SegSeqDims =>
    (SegSeqDims -> SegSeqDims -> Ordering)
    -> (SegSeqDims -> SegSeqDims -> Bool)
    -> (SegSeqDims -> SegSeqDims -> Bool)
    -> (SegSeqDims -> SegSeqDims -> Bool)
    -> (SegSeqDims -> SegSeqDims -> Bool)
    -> (SegSeqDims -> SegSeqDims -> SegSeqDims)
    -> (SegSeqDims -> SegSeqDims -> SegSeqDims)
    -> Ord SegSeqDims
    SegSeqDims -> SegSeqDims -> Bool
    SegSeqDims -> SegSeqDims -> Ordering
    SegSeqDims -> SegSeqDims -> SegSeqDims
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    $ccompare :: SegSeqDims -> SegSeqDims -> Ordering
    compare :: SegSeqDims -> SegSeqDims -> Ordering
    $c< :: SegSeqDims -> SegSeqDims -> Bool
    < :: SegSeqDims -> SegSeqDims -> Bool
    $c<= :: SegSeqDims -> SegSeqDims -> Bool
    <= :: SegSeqDims -> SegSeqDims -> Bool
    $c> :: SegSeqDims -> SegSeqDims -> Bool
    > :: SegSeqDims -> SegSeqDims -> Bool
    $c>= :: SegSeqDims -> SegSeqDims -> Bool
    >= :: SegSeqDims -> SegSeqDims -> Bool
    $cmax :: SegSeqDims -> SegSeqDims -> SegSeqDims
    max :: SegSeqDims -> SegSeqDims -> SegSeqDims
    $cmin :: SegSeqDims -> SegSeqDims -> SegSeqDims
    min :: SegSeqDims -> SegSeqDims -> SegSeqDims
    Ord, Int -> SegSeqDims -> ShowS
    [SegSeqDims] -> ShowS
    SegSeqDims -> String
    (Int -> SegSeqDims -> ShowS)
    -> (SegSeqDims -> String)
    -> ([SegSeqDims] -> ShowS)
    -> Show SegSeqDims
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> SegSeqDims -> ShowS
    showsPrec :: Int -> SegSeqDims -> ShowS
    $cshow :: SegSeqDims -> String
    show :: SegSeqDims -> String
    $cshowList :: [SegSeqDims] -> ShowS
    showList :: [SegSeqDims] -> ShowS
    Show)

    -- | Do we need block-virtualisation when generating code for the
    -- segmented operation?  In most cases, we do, but for some simple
    -- kernels, we compute the full number of blocks in advance, and then
    -- virtualisation is an unnecessary (but generally very small)
    -- overhead.  This only really matters for fairly trivial but very
    -- wide @map@ kernels where each thread performs constant-time work on
    -- scalars.
    data SegVirt
      = SegVirt
      | SegNoVirt
      | -- | Not only do we not need virtualisation, but we _guarantee_
        -- that all physical threads participate in the work.  This can
        -- save some checks in code generation.
        SegNoVirtFull SegSeqDims
      deriving (SegVirt -> SegVirt -> Bool
    (SegVirt -> SegVirt -> Bool)
    -> (SegVirt -> SegVirt -> Bool) -> Eq SegVirt
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: SegVirt -> SegVirt -> Bool
    == :: SegVirt -> SegVirt -> Bool
    $c/= :: SegVirt -> SegVirt -> Bool
    /= :: SegVirt -> SegVirt -> Bool
    Eq, Eq SegVirt
    Eq SegVirt =>
    (SegVirt -> SegVirt -> Ordering)
    -> (SegVirt -> SegVirt -> Bool)
    -> (SegVirt -> SegVirt -> Bool)
    -> (SegVirt -> SegVirt -> Bool)
    -> (SegVirt -> SegVirt -> Bool)
    -> (SegVirt -> SegVirt -> SegVirt)
    -> (SegVirt -> SegVirt -> SegVirt)
    -> Ord SegVirt
    SegVirt -> SegVirt -> Bool
    SegVirt -> SegVirt -> Ordering
    SegVirt -> SegVirt -> SegVirt
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    $ccompare :: SegVirt -> SegVirt -> Ordering
    compare :: SegVirt -> SegVirt -> Ordering
    $c< :: SegVirt -> SegVirt -> Bool
    < :: SegVirt -> SegVirt -> Bool
    $c<= :: SegVirt -> SegVirt -> Bool
    <= :: SegVirt -> SegVirt -> Bool
    $c> :: SegVirt -> SegVirt -> Bool
    > :: SegVirt -> SegVirt -> Bool
    $c>= :: SegVirt -> SegVirt -> Bool
    >= :: SegVirt -> SegVirt -> Bool
    $cmax :: SegVirt -> SegVirt -> SegVirt
    max :: SegVirt -> SegVirt -> SegVirt
    $cmin :: SegVirt -> SegVirt -> SegVirt
    min :: SegVirt -> SegVirt -> SegVirt
    Ord, Int -> SegVirt -> ShowS
    [SegVirt] -> ShowS
    SegVirt -> String
    (Int -> SegVirt -> ShowS)
    -> (SegVirt -> String) -> ([SegVirt] -> ShowS) -> Show SegVirt
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> SegVirt -> ShowS
    showsPrec :: Int -> SegVirt -> ShowS
    $cshow :: SegVirt -> String
    show :: SegVirt -> String
    $cshowList :: [SegVirt] -> ShowS
    showList :: [SegVirt] -> ShowS
    Show)

    -- | The actual, physical grid dimensions used for the GPU kernel
    -- running this 'SegOp'.
    data KernelGrid = KernelGrid
      { KernelGrid -> Count NumBlocks SubExp
    gridNumBlocks :: Count NumBlocks SubExp,
        KernelGrid -> Count BlockSize SubExp
    gridBlockSize :: Count BlockSize SubExp
      }
      deriving (KernelGrid -> KernelGrid -> Bool
    (KernelGrid -> KernelGrid -> Bool)
    -> (KernelGrid -> KernelGrid -> Bool) -> Eq KernelGrid
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: KernelGrid -> KernelGrid -> Bool
    == :: KernelGrid -> KernelGrid -> Bool
    $c/= :: KernelGrid -> KernelGrid -> Bool
    /= :: KernelGrid -> KernelGrid -> Bool
    Eq, Eq KernelGrid
    Eq KernelGrid =>
    (KernelGrid -> KernelGrid -> Ordering)
    -> (KernelGrid -> KernelGrid -> Bool)
    -> (KernelGrid -> KernelGrid -> Bool)
    -> (KernelGrid -> KernelGrid -> Bool)
    -> (KernelGrid -> KernelGrid -> Bool)
    -> (KernelGrid -> KernelGrid -> KernelGrid)
    -> (KernelGrid -> KernelGrid -> KernelGrid)
    -> Ord KernelGrid
    KernelGrid -> KernelGrid -> Bool
    KernelGrid -> KernelGrid -> Ordering
    KernelGrid -> KernelGrid -> KernelGrid
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    $ccompare :: KernelGrid -> KernelGrid -> Ordering
    compare :: KernelGrid -> KernelGrid -> Ordering
    $c< :: KernelGrid -> KernelGrid -> Bool
    < :: KernelGrid -> KernelGrid -> Bool
    $c<= :: KernelGrid -> KernelGrid -> Bool
    <= :: KernelGrid -> KernelGrid -> Bool
    $c> :: KernelGrid -> KernelGrid -> Bool
    > :: KernelGrid -> KernelGrid -> Bool
    $c>= :: KernelGrid -> KernelGrid -> Bool
    >= :: KernelGrid -> KernelGrid -> Bool
    $cmax :: KernelGrid -> KernelGrid -> KernelGrid
    max :: KernelGrid -> KernelGrid -> KernelGrid
    $cmin :: KernelGrid -> KernelGrid -> KernelGrid
    min :: KernelGrid -> KernelGrid -> KernelGrid
    Ord, Int -> KernelGrid -> ShowS
    [KernelGrid] -> ShowS
    KernelGrid -> String
    (Int -> KernelGrid -> ShowS)
    -> (KernelGrid -> String)
    -> ([KernelGrid] -> ShowS)
    -> Show KernelGrid
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> KernelGrid -> ShowS
    showsPrec :: Int -> KernelGrid -> ShowS
    $cshow :: KernelGrid -> String
    show :: KernelGrid -> String
    $cshowList :: [KernelGrid] -> ShowS
    showList :: [KernelGrid] -> ShowS
    Show)

    -- | At which level the *body* of a t'SegOp' executes.
    data SegLevel
      = SegThread SegVirt (Maybe KernelGrid)
      | SegBlock SegVirt (Maybe KernelGrid)
      | SegThreadInBlock SegVirt
      deriving (SegLevel -> SegLevel -> Bool
    (SegLevel -> SegLevel -> Bool)
    -> (SegLevel -> SegLevel -> Bool) -> Eq SegLevel
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: SegLevel -> SegLevel -> Bool
    == :: SegLevel -> SegLevel -> Bool
    $c/= :: SegLevel -> SegLevel -> Bool
    /= :: SegLevel -> SegLevel -> Bool
    Eq, Eq SegLevel
    Eq SegLevel =>
    (SegLevel -> SegLevel -> Ordering)
    -> (SegLevel -> SegLevel -> Bool)
    -> (SegLevel -> SegLevel -> Bool)
    -> (SegLevel -> SegLevel -> Bool)
    -> (SegLevel -> SegLevel -> Bool)
    -> (SegLevel -> SegLevel -> SegLevel)
    -> (SegLevel -> SegLevel -> SegLevel)
    -> Ord SegLevel
    SegLevel -> SegLevel -> Bool
    SegLevel -> SegLevel -> Ordering
    SegLevel -> SegLevel -> SegLevel
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    $ccompare :: SegLevel -> SegLevel -> Ordering
    compare :: SegLevel -> SegLevel -> Ordering
    $c< :: SegLevel -> SegLevel -> Bool
    < :: SegLevel -> SegLevel -> Bool
    $c<= :: SegLevel -> SegLevel -> Bool
    <= :: SegLevel -> SegLevel -> Bool
    $c> :: SegLevel -> SegLevel -> Bool
    > :: SegLevel -> SegLevel -> Bool
    $c>= :: SegLevel -> SegLevel -> Bool
    >= :: SegLevel -> SegLevel -> Bool
    $cmax :: SegLevel -> SegLevel -> SegLevel
    max :: SegLevel -> SegLevel -> SegLevel
    $cmin :: SegLevel -> SegLevel -> SegLevel
    min :: SegLevel -> SegLevel -> SegLevel
    Ord, Int -> SegLevel -> ShowS
    [SegLevel] -> ShowS
    SegLevel -> String
    (Int -> SegLevel -> ShowS)
    -> (SegLevel -> String) -> ([SegLevel] -> ShowS) -> Show SegLevel
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> SegLevel -> ShowS
    showsPrec :: Int -> SegLevel -> ShowS
    $cshow :: SegLevel -> String
    show :: SegLevel -> String
    $cshowList :: [SegLevel] -> ShowS
    showList :: [SegLevel] -> ShowS
    Show)

    -- | The 'SegVirt' of the 'SegLevel'.
    segVirt :: SegLevel -> SegVirt
    segVirt :: SegLevel -> SegVirt
    segVirt (SegThread SegVirt
    v Maybe KernelGrid
    _) = SegVirt
    v
    segVirt (SegBlock SegVirt
    v Maybe KernelGrid
    _) = SegVirt
    v
    segVirt (SegThreadInBlock SegVirt
    v) = SegVirt
    v

    instance PP.Pretty SegVirt where
      pretty :: forall ann. SegVirt -> Doc ann
    pretty SegVirt
    SegNoVirt = Doc ann
    forall a. Monoid a => a
    mempty
      pretty (SegNoVirtFull SegSeqDims
    dims) = Doc ann
    "full" Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> [Int] -> Doc ann
    forall ann. [Int] -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty (SegSeqDims -> [Int]
    segSeqDims SegSeqDims
    dims)
      pretty SegVirt
    SegVirt = Doc ann
    "virtualise"

    instance PP.Pretty KernelGrid where
      pretty :: forall ann. KernelGrid -> Doc ann
    pretty (KernelGrid Count NumBlocks SubExp
    num_tblocks Count BlockSize SubExp
    tblock_size) =
        Doc ann
    "grid="
          Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Count NumBlocks SubExp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Count NumBlocks SubExp -> Doc ann
    pretty Count NumBlocks SubExp
    num_tblocks
          Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    PP.semi
            Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    "blocksize="
          Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Count BlockSize SubExp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Count BlockSize SubExp -> Doc ann
    pretty Count BlockSize SubExp
    tblock_size

    instance PP.Pretty SegLevel where
      pretty :: forall ann. SegLevel -> Doc ann
    pretty (SegThread SegVirt
    virt Maybe KernelGrid
    grid) =
        Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    PP.parens (Doc ann
    "thread" Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    PP.semi Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> SegVirt -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SegVirt -> Doc ann
    pretty SegVirt
    virt Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    PP.semi Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Maybe KernelGrid -> Doc ann
    forall ann. Maybe KernelGrid -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty Maybe KernelGrid
    grid)
      pretty (SegBlock SegVirt
    virt Maybe KernelGrid
    grid) =
        Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    PP.parens (Doc ann
    "block" Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    PP.semi Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> SegVirt -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SegVirt -> Doc ann
    pretty SegVirt
    virt Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    PP.semi Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Maybe KernelGrid -> Doc ann
    forall ann. Maybe KernelGrid -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty Maybe KernelGrid
    grid)
      pretty (SegThreadInBlock SegVirt
    virt) =
        Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    PP.parens (Doc ann
    "inblock" Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    PP.semi Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> SegVirt -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SegVirt -> Doc ann
    pretty SegVirt
    virt)

    instance Engine.Simplifiable KernelGrid where
      simplify :: forall rep.
    SimplifiableRep rep =>
    KernelGrid -> SimpleM rep KernelGrid
    simplify (KernelGrid Count NumBlocks SubExp
    num_tblocks Count BlockSize SubExp
    tblock_size) =
        Count NumBlocks SubExp -> Count BlockSize SubExp -> KernelGrid
    KernelGrid
          (Count NumBlocks SubExp -> Count BlockSize SubExp -> KernelGrid)
    -> SimpleM rep (Count NumBlocks SubExp)
    -> SimpleM rep (Count BlockSize SubExp -> KernelGrid)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (SubExp -> SimpleM rep SubExp)
    -> Count NumBlocks SubExp -> SimpleM rep (Count NumBlocks SubExp)
    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) -> Count NumBlocks a -> f (Count NumBlocks b)
    traverse SubExp -> SimpleM rep SubExp
    forall rep. SimplifiableRep rep => SubExp -> SimpleM rep SubExp
    forall e rep.
    (Simplifiable e, SimplifiableRep rep) =>
    e -> SimpleM rep e
    Engine.simplify Count NumBlocks SubExp
    num_tblocks
          SimpleM rep (Count BlockSize SubExp -> KernelGrid)
    -> SimpleM rep (Count BlockSize SubExp) -> SimpleM rep KernelGrid
    forall a b. SimpleM rep (a -> b) -> SimpleM rep a -> SimpleM rep b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (SubExp -> SimpleM rep SubExp)
    -> Count BlockSize SubExp -> SimpleM rep (Count BlockSize SubExp)
    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) -> Count BlockSize a -> f (Count BlockSize b)
    traverse SubExp -> SimpleM rep SubExp
    forall rep. SimplifiableRep rep => SubExp -> SimpleM rep SubExp
    forall e rep.
    (Simplifiable e, SimplifiableRep rep) =>
    e -> SimpleM rep e
    Engine.simplify Count BlockSize SubExp
    tblock_size

    instance Engine.Simplifiable SegLevel where
      simplify :: forall rep. SimplifiableRep rep => SegLevel -> SimpleM rep SegLevel
    simplify (SegThread SegVirt
    virt Maybe KernelGrid
    grid) =
        SegVirt -> Maybe KernelGrid -> SegLevel
    SegThread SegVirt
    virt (Maybe KernelGrid -> SegLevel)
    -> SimpleM rep (Maybe KernelGrid) -> SimpleM rep SegLevel
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Maybe KernelGrid -> SimpleM rep (Maybe KernelGrid)
    forall rep.
    SimplifiableRep rep =>
    Maybe KernelGrid -> SimpleM rep (Maybe KernelGrid)
    forall e rep.
    (Simplifiable e, SimplifiableRep rep) =>
    e -> SimpleM rep e
    Engine.simplify Maybe KernelGrid
    grid
      simplify (SegBlock SegVirt
    virt Maybe KernelGrid
    grid) =
        SegVirt -> Maybe KernelGrid -> SegLevel
    SegBlock SegVirt
    virt (Maybe KernelGrid -> SegLevel)
    -> SimpleM rep (Maybe KernelGrid) -> SimpleM rep SegLevel
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Maybe KernelGrid -> SimpleM rep (Maybe KernelGrid)
    forall rep.
    SimplifiableRep rep =>
    Maybe KernelGrid -> SimpleM rep (Maybe KernelGrid)
    forall e rep.
    (Simplifiable e, SimplifiableRep rep) =>
    e -> SimpleM rep e
    Engine.simplify Maybe KernelGrid
    grid
      simplify (SegThreadInBlock SegVirt
    virt) =
        SegLevel -> SimpleM rep SegLevel
    forall a. a -> SimpleM rep a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (SegLevel -> SimpleM rep SegLevel)
    -> SegLevel -> SimpleM rep SegLevel
    forall a b. (a -> b) -> a -> b
    $ SegVirt -> SegLevel
    SegThreadInBlock SegVirt
    virt

    instance Substitute KernelGrid where
      substituteNames :: Map VName VName -> KernelGrid -> KernelGrid
    substituteNames Map VName VName
    substs (KernelGrid Count NumBlocks SubExp
    num_tblocks Count BlockSize SubExp
    tblock_size) =
        Count NumBlocks SubExp -> Count BlockSize SubExp -> KernelGrid
    KernelGrid
          (Map VName VName -> Count NumBlocks SubExp -> Count NumBlocks SubExp
    forall a. Substitute a => Map VName VName -> a -> a
    substituteNames Map VName VName
    substs Count NumBlocks SubExp
    num_tblocks)
          (Map VName VName -> Count BlockSize SubExp -> Count BlockSize SubExp
    forall a. Substitute a => Map VName VName -> a -> a
    substituteNames Map VName VName
    substs Count BlockSize SubExp
    tblock_size)

    instance Substitute SegLevel where
      substituteNames :: Map VName VName -> SegLevel -> SegLevel
    substituteNames Map VName VName
    substs (SegThread SegVirt
    virt Maybe KernelGrid
    grid) =
        SegVirt -> Maybe KernelGrid -> SegLevel
    SegThread SegVirt
    virt (Map VName VName -> Maybe KernelGrid -> Maybe KernelGrid
    forall a. Substitute a => Map VName VName -> a -> a
    substituteNames Map VName VName
    substs Maybe KernelGrid
    grid)
      substituteNames Map VName VName
    substs (SegBlock SegVirt
    virt Maybe KernelGrid
    grid) =
        SegVirt -> Maybe KernelGrid -> SegLevel
    SegBlock SegVirt
    virt (Map VName VName -> Maybe KernelGrid -> Maybe KernelGrid
    forall a. Substitute a => Map VName VName -> a -> a
    substituteNames Map VName VName
    substs Maybe KernelGrid
    grid)
      substituteNames Map VName VName
    _ (SegThreadInBlock SegVirt
    virt) =
        SegVirt -> SegLevel
    SegThreadInBlock SegVirt
    virt

    instance Rename SegLevel where
      rename :: SegLevel -> RenameM SegLevel
    rename = SegLevel -> RenameM SegLevel
    forall a. Substitute a => a -> RenameM a
    substituteRename

    instance FreeIn KernelGrid where
      freeIn' :: KernelGrid -> FV
    freeIn' (KernelGrid Count NumBlocks SubExp
    num_tblocks Count BlockSize SubExp
    tblock_size) =
        (Count NumBlocks SubExp, Count BlockSize SubExp) -> FV
    forall a. FreeIn a => a -> FV
    freeIn' (Count NumBlocks SubExp
    num_tblocks, Count BlockSize SubExp
    tblock_size)

    instance FreeIn SegLevel where
      freeIn' :: SegLevel -> FV
    freeIn' (SegThread SegVirt
    _virt Maybe KernelGrid
    grid) = Maybe KernelGrid -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Maybe KernelGrid
    grid
      freeIn' (SegBlock SegVirt
    _virt Maybe KernelGrid
    grid) = Maybe KernelGrid -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Maybe KernelGrid
    grid
      freeIn' (SegThreadInBlock SegVirt
    _virt) = FV
    forall a. Monoid a => a
    mempty

    -- | A simple size-level query or computation.
    data SizeOp
      = -- | Produce some runtime-configurable size.
        GetSize Name SizeClass
      | -- | The maximum size of some class.
        GetSizeMax SizeClass
      | -- | Compare size (likely a threshold) with some integer value.
        CmpSizeLe Name SizeClass SubExp
      | -- | @CalcNumBlocks w max_num_tblocks tblock_size@ calculates the
        -- number of GPU threadblocks to use for an input of the given size.
        -- The @Name@ is a size name.  Note that @w@ is an i64 to avoid
        -- overflow issues.
        CalcNumBlocks SubExp Name SubExp
      deriving (SizeOp -> SizeOp -> Bool
    (SizeOp -> SizeOp -> Bool)
    -> (SizeOp -> SizeOp -> Bool) -> Eq SizeOp
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: SizeOp -> SizeOp -> Bool
    == :: SizeOp -> SizeOp -> Bool
    $c/= :: SizeOp -> SizeOp -> Bool
    /= :: SizeOp -> SizeOp -> Bool
    Eq, Eq SizeOp
    Eq SizeOp =>
    (SizeOp -> SizeOp -> Ordering)
    -> (SizeOp -> SizeOp -> Bool)
    -> (SizeOp -> SizeOp -> Bool)
    -> (SizeOp -> SizeOp -> Bool)
    -> (SizeOp -> SizeOp -> Bool)
    -> (SizeOp -> SizeOp -> SizeOp)
    -> (SizeOp -> SizeOp -> SizeOp)
    -> Ord SizeOp
    SizeOp -> SizeOp -> Bool
    SizeOp -> SizeOp -> Ordering
    SizeOp -> SizeOp -> SizeOp
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    $ccompare :: SizeOp -> SizeOp -> Ordering
    compare :: SizeOp -> SizeOp -> Ordering
    $c< :: SizeOp -> SizeOp -> Bool
    < :: SizeOp -> SizeOp -> Bool
    $c<= :: SizeOp -> SizeOp -> Bool
    <= :: SizeOp -> SizeOp -> Bool
    $c> :: SizeOp -> SizeOp -> Bool
    > :: SizeOp -> SizeOp -> Bool
    $c>= :: SizeOp -> SizeOp -> Bool
    >= :: SizeOp -> SizeOp -> Bool
    $cmax :: SizeOp -> SizeOp -> SizeOp
    max :: SizeOp -> SizeOp -> SizeOp
    $cmin :: SizeOp -> SizeOp -> SizeOp
    min :: SizeOp -> SizeOp -> SizeOp
    Ord, Int -> SizeOp -> ShowS
    [SizeOp] -> ShowS
    SizeOp -> String
    (Int -> SizeOp -> ShowS)
    -> (SizeOp -> String) -> ([SizeOp] -> ShowS) -> Show SizeOp
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> SizeOp -> ShowS
    showsPrec :: Int -> SizeOp -> ShowS
    $cshow :: SizeOp -> String
    show :: SizeOp -> String
    $cshowList :: [SizeOp] -> ShowS
    showList :: [SizeOp] -> ShowS
    Show)

    instance Substitute SizeOp where
      substituteNames :: Map VName VName -> SizeOp -> SizeOp
    substituteNames Map VName VName
    substs (CmpSizeLe Name
    name SizeClass
    sclass SubExp
    x) =
        Name -> SizeClass -> SubExp -> SizeOp
    CmpSizeLe Name
    name SizeClass
    sclass (Map VName VName -> SubExp -> SubExp
    forall a. Substitute a => Map VName VName -> a -> a
    substituteNames Map VName VName
    substs SubExp
    x)
      substituteNames Map VName VName
    substs (CalcNumBlocks SubExp
    w Name
    max_num_tblocks SubExp
    tblock_size) =
        SubExp -> Name -> SubExp -> SizeOp
    CalcNumBlocks
          (Map VName VName -> SubExp -> SubExp
    forall a. Substitute a => Map VName VName -> a -> a
    substituteNames Map VName VName
    substs SubExp
    w)
          Name
    max_num_tblocks
          (Map VName VName -> SubExp -> SubExp
    forall a. Substitute a => Map VName VName -> a -> a
    substituteNames Map VName VName
    substs SubExp
    tblock_size)
      substituteNames Map VName VName
    _ SizeOp
    op = SizeOp
    op

    instance Rename SizeOp where
      rename :: SizeOp -> RenameM SizeOp
    rename (CmpSizeLe Name
    name SizeClass
    sclass SubExp
    x) =
        Name -> SizeClass -> SubExp -> SizeOp
    CmpSizeLe Name
    name SizeClass
    sclass (SubExp -> SizeOp) -> RenameM SubExp -> RenameM SizeOp
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> SubExp -> RenameM SubExp
    forall a. Rename a => a -> RenameM a
    rename SubExp
    x
      rename (CalcNumBlocks SubExp
    w Name
    max_num_tblocks SubExp
    tblock_size) =
        SubExp -> Name -> SubExp -> SizeOp
    CalcNumBlocks (SubExp -> Name -> SubExp -> SizeOp)
    -> RenameM SubExp -> RenameM (Name -> SubExp -> SizeOp)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> SubExp -> RenameM SubExp
    forall a. Rename a => a -> RenameM a
    rename SubExp
    w RenameM (Name -> SubExp -> SizeOp)
    -> RenameM Name -> RenameM (SubExp -> SizeOp)
    forall a b. RenameM (a -> b) -> RenameM a -> RenameM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> Name -> RenameM Name
    forall a. a -> RenameM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Name
    max_num_tblocks RenameM (SubExp -> SizeOp) -> RenameM SubExp -> RenameM SizeOp
    forall a b. RenameM (a -> b) -> RenameM a -> RenameM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SubExp -> RenameM SubExp
    forall a. Rename a => a -> RenameM a
    rename SubExp
    tblock_size
      rename SizeOp
    x = SizeOp -> RenameM SizeOp
    forall a. a -> RenameM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure SizeOp
    x

    instance FreeIn SizeOp where
      freeIn' :: SizeOp -> FV
    freeIn' (CmpSizeLe Name
    _ SizeClass
    _ SubExp
    x) = SubExp -> FV
    forall a. FreeIn a => a -> FV
    freeIn' SubExp
    x
      freeIn' (CalcNumBlocks SubExp
    w Name
    _ SubExp
    tblock_size) = SubExp -> FV
    forall a. FreeIn a => a -> FV
    freeIn' SubExp
    w FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> SubExp -> FV
    forall a. FreeIn a => a -> FV
    freeIn' SubExp
    tblock_size
      freeIn' SizeOp
    _ = FV
    forall a. Monoid a => a
    mempty

    instance PP.Pretty SizeOp where
      pretty :: forall ann. SizeOp -> Doc ann
    pretty (GetSize Name
    name SizeClass
    size_class) =
        Doc ann
    "get_size" Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    parens ([Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    commasep [Name -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Name -> Doc ann
    pretty Name
    name, SizeClass -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SizeClass -> Doc ann
    pretty SizeClass
    size_class])
      pretty (GetSizeMax SizeClass
    size_class) =
        Doc ann
    "get_size_max" Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    parens ([Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    commasep [SizeClass -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SizeClass -> Doc ann
    pretty SizeClass
    size_class])
      pretty (CmpSizeLe Name
    name SizeClass
    size_class SubExp
    x) =
        Doc ann
    "cmp_size"
          Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    parens ([Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    commasep [Name -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Name -> Doc ann
    pretty Name
    name, SizeClass -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SizeClass -> Doc ann
    pretty SizeClass
    size_class])
            Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    "<="
            Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> SubExp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SubExp -> Doc ann
    pretty SubExp
    x
      pretty (CalcNumBlocks SubExp
    w Name
    max_num_tblocks SubExp
    tblock_size) =
        Doc ann
    "calc_num_tblocks" Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    parens ([Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    commasep [SubExp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SubExp -> Doc ann
    pretty SubExp
    w, Name -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Name -> Doc ann
    pretty Name
    max_num_tblocks, SubExp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SubExp -> Doc ann
    pretty SubExp
    tblock_size])

    instance OpMetrics SizeOp where
      opMetrics :: SizeOp -> MetricsM ()
    opMetrics GetSize {} = Text -> MetricsM ()
    seen Text
    "GetSize"
      opMetrics GetSizeMax {} = Text -> MetricsM ()
    seen Text
    "GetSizeMax"
      opMetrics CmpSizeLe {} = Text -> MetricsM ()
    seen Text
    "CmpSizeLe"
      opMetrics CalcNumBlocks {} = Text -> MetricsM ()
    seen Text
    "CalcNumBlocks"

    typeCheckSizeOp :: (TC.Checkable rep) => SizeOp -> TC.TypeM rep ()
    typeCheckSizeOp :: forall rep. Checkable rep => SizeOp -> TypeM rep ()
    typeCheckSizeOp GetSize {} = () -> TypeM rep ()
    forall a. a -> TypeM rep a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
    typeCheckSizeOp GetSizeMax {} = () -> TypeM rep ()
    forall a. a -> TypeM rep a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
    typeCheckSizeOp (CmpSizeLe Name
    _ SizeClass
    _ SubExp
    x) = [Type] -> SubExp -> TypeM rep ()
    forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
    TC.require [PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64] SubExp
    x
    typeCheckSizeOp (CalcNumBlocks SubExp
    w Name
    _ SubExp
    tblock_size) = do
      [Type] -> SubExp -> TypeM rep ()
    forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
    TC.require [PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64] SubExp
    w
      [Type] -> SubExp -> TypeM rep ()
    forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
    TC.require [PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64] SubExp
    tblock_size

    -- | A host-level operation; parameterised by what else it can do.
    data HostOp op rep
      = -- | A segmented operation.
        SegOp (SegOp SegLevel rep)
      | SizeOp SizeOp
      | OtherOp (op rep)
      | -- | Code to run sequentially on the GPU,
        -- in a single thread.
        GPUBody [Type] (Body rep)
      deriving (HostOp op rep -> HostOp op rep -> Bool
    (HostOp op rep -> HostOp op rep -> Bool)
    -> (HostOp op rep -> HostOp op rep -> Bool) -> Eq (HostOp op rep)
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    forall (op :: * -> *) rep.
    (RepTypes rep, Eq (op rep)) =>
    HostOp op rep -> HostOp op rep -> Bool
    $c== :: forall (op :: * -> *) rep.
    (RepTypes rep, Eq (op rep)) =>
    HostOp op rep -> HostOp op rep -> Bool
    == :: HostOp op rep -> HostOp op rep -> Bool
    $c/= :: forall (op :: * -> *) rep.
    (RepTypes rep, Eq (op rep)) =>
    HostOp op rep -> HostOp op rep -> Bool
    /= :: HostOp op rep -> HostOp op rep -> Bool
    Eq, Eq (HostOp op rep)
    Eq (HostOp op rep) =>
    (HostOp op rep -> HostOp op rep -> Ordering)
    -> (HostOp op rep -> HostOp op rep -> Bool)
    -> (HostOp op rep -> HostOp op rep -> Bool)
    -> (HostOp op rep -> HostOp op rep -> Bool)
    -> (HostOp op rep -> HostOp op rep -> Bool)
    -> (HostOp op rep -> HostOp op rep -> HostOp op rep)
    -> (HostOp op rep -> HostOp op rep -> HostOp op rep)
    -> Ord (HostOp op rep)
    HostOp op rep -> HostOp op rep -> Bool
    HostOp op rep -> HostOp op rep -> Ordering
    HostOp op rep -> HostOp op rep -> HostOp op rep
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    forall (op :: * -> *) rep.
    (RepTypes rep, Ord (op rep)) =>
    Eq (HostOp op rep)
    forall (op :: * -> *) rep.
    (RepTypes rep, Ord (op rep)) =>
    HostOp op rep -> HostOp op rep -> Bool
    forall (op :: * -> *) rep.
    (RepTypes rep, Ord (op rep)) =>
    HostOp op rep -> HostOp op rep -> Ordering
    forall (op :: * -> *) rep.
    (RepTypes rep, Ord (op rep)) =>
    HostOp op rep -> HostOp op rep -> HostOp op rep
    $ccompare :: forall (op :: * -> *) rep.
    (RepTypes rep, Ord (op rep)) =>
    HostOp op rep -> HostOp op rep -> Ordering
    compare :: HostOp op rep -> HostOp op rep -> Ordering
    $c< :: forall (op :: * -> *) rep.
    (RepTypes rep, Ord (op rep)) =>
    HostOp op rep -> HostOp op rep -> Bool
    < :: HostOp op rep -> HostOp op rep -> Bool
    $c<= :: forall (op :: * -> *) rep.
    (RepTypes rep, Ord (op rep)) =>
    HostOp op rep -> HostOp op rep -> Bool
    <= :: HostOp op rep -> HostOp op rep -> Bool
    $c> :: forall (op :: * -> *) rep.
    (RepTypes rep, Ord (op rep)) =>
    HostOp op rep -> HostOp op rep -> Bool
    > :: HostOp op rep -> HostOp op rep -> Bool
    $c>= :: forall (op :: * -> *) rep.
    (RepTypes rep, Ord (op rep)) =>
    HostOp op rep -> HostOp op rep -> Bool
    >= :: HostOp op rep -> HostOp op rep -> Bool
    $cmax :: forall (op :: * -> *) rep.
    (RepTypes rep, Ord (op rep)) =>
    HostOp op rep -> HostOp op rep -> HostOp op rep
    max :: HostOp op rep -> HostOp op rep -> HostOp op rep
    $cmin :: forall (op :: * -> *) rep.
    (RepTypes rep, Ord (op rep)) =>
    HostOp op rep -> HostOp op rep -> HostOp op rep
    min :: HostOp op rep -> HostOp op rep -> HostOp op rep
    Ord, Int -> HostOp op rep -> ShowS
    [HostOp op rep] -> ShowS
    HostOp op rep -> String
    (Int -> HostOp op rep -> ShowS)
    -> (HostOp op rep -> String)
    -> ([HostOp op rep] -> ShowS)
    -> Show (HostOp op rep)
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    forall (op :: * -> *) rep.
    (RepTypes rep, Show (op rep)) =>
    Int -> HostOp op rep -> ShowS
    forall (op :: * -> *) rep.
    (RepTypes rep, Show (op rep)) =>
    [HostOp op rep] -> ShowS
    forall (op :: * -> *) rep.
    (RepTypes rep, Show (op rep)) =>
    HostOp op rep -> String
    $cshowsPrec :: forall (op :: * -> *) rep.
    (RepTypes rep, Show (op rep)) =>
    Int -> HostOp op rep -> ShowS
    showsPrec :: Int -> HostOp op rep -> ShowS
    $cshow :: forall (op :: * -> *) rep.
    (RepTypes rep, Show (op rep)) =>
    HostOp op rep -> String
    show :: HostOp op rep -> String
    $cshowList :: forall (op :: * -> *) rep.
    (RepTypes rep, Show (op rep)) =>
    [HostOp op rep] -> ShowS
    showList :: [HostOp op rep] -> ShowS
    Show)

    -- | A helper for defining 'TraverseOpStms'.
    traverseHostOpStms ::
      (Monad m) =>
      OpStmsTraverser m (op rep) rep ->
      OpStmsTraverser m (HostOp op rep) rep
    traverseHostOpStms :: forall (m :: * -> *) (op :: * -> *) rep.
    Monad m =>
    OpStmsTraverser m (op rep) rep
    -> OpStmsTraverser m (HostOp op rep) rep
    traverseHostOpStms OpStmsTraverser m (op rep) rep
    _ Scope rep -> Stms rep -> m (Stms rep)
    f (SegOp SegOp SegLevel rep
    segop) = SegOp SegLevel rep -> HostOp op rep
    forall (op :: * -> *) rep. SegOp SegLevel rep -> HostOp op rep
    SegOp (SegOp SegLevel rep -> HostOp op rep)
    -> m (SegOp SegLevel rep) -> m (HostOp op rep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> OpStmsTraverser m (SegOp SegLevel rep) rep
    forall (m :: * -> *) lvl rep.
    Monad m =>
    OpStmsTraverser m (SegOp lvl rep) rep
    traverseSegOpStms Scope rep -> Stms rep -> m (Stms rep)
    f SegOp SegLevel rep
    segop
    traverseHostOpStms OpStmsTraverser m (op rep) rep
    _ Scope rep -> Stms rep -> m (Stms rep)
    _ (SizeOp SizeOp
    sizeop) = HostOp op rep -> m (HostOp op rep)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (HostOp op rep -> m (HostOp op rep))
    -> HostOp op rep -> m (HostOp op rep)
    forall a b. (a -> b) -> a -> b
    $ SizeOp -> HostOp op rep
    forall (op :: * -> *) rep. SizeOp -> HostOp op rep
    SizeOp SizeOp
    sizeop
    traverseHostOpStms OpStmsTraverser m (op rep) rep
    onOtherOp Scope rep -> Stms rep -> m (Stms rep)
    f (OtherOp op rep
    other) = op rep -> HostOp op rep
    forall (op :: * -> *) rep. op rep -> HostOp op rep
    OtherOp (op rep -> HostOp op rep) -> m (op rep) -> m (HostOp op rep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> OpStmsTraverser m (op rep) rep
    onOtherOp Scope rep -> Stms rep -> m (Stms rep)
    f op rep
    other
    traverseHostOpStms OpStmsTraverser m (op rep) rep
    _ Scope rep -> Stms rep -> m (Stms rep)
    f (GPUBody [Type]
    ts Body rep
    body) = do
      stms <- Scope rep -> Stms rep -> m (Stms rep)
    f Scope rep
    forall a. Monoid a => a
    mempty (Stms rep -> m (Stms rep)) -> Stms rep -> m (Stms rep)
    forall a b. (a -> b) -> a -> b
    $ Body rep -> Stms rep
    forall rep. Body rep -> Stms rep
    bodyStms Body rep
    body
      pure $ GPUBody ts $ body {bodyStms = stms}

    instance (ASTRep rep, Substitute (op rep)) => Substitute (HostOp op rep) where
      substituteNames :: Map VName VName -> HostOp op rep -> HostOp op rep
    substituteNames Map VName VName
    substs (SegOp SegOp SegLevel rep
    op) =
        SegOp SegLevel rep -> HostOp op rep
    forall (op :: * -> *) rep. SegOp SegLevel rep -> HostOp op rep
    SegOp (SegOp SegLevel rep -> HostOp op rep)
    -> SegOp SegLevel rep -> HostOp op rep
    forall a b. (a -> b) -> a -> b
    $ Map VName VName -> SegOp SegLevel rep -> SegOp SegLevel rep
    forall a. Substitute a => Map VName VName -> a -> a
    substituteNames Map VName VName
    substs SegOp SegLevel rep
    op
      substituteNames Map VName VName
    substs (OtherOp op rep
    op) =
        op rep -> HostOp op rep
    forall (op :: * -> *) rep. op rep -> HostOp op rep
    OtherOp (op rep -> HostOp op rep) -> op rep -> HostOp op rep
    forall a b. (a -> b) -> a -> b
    $ Map VName VName -> op rep -> op rep
    forall a. Substitute a => Map VName VName -> a -> a
    substituteNames Map VName VName
    substs op rep
    op
      substituteNames Map VName VName
    substs (SizeOp SizeOp
    op) =
        SizeOp -> HostOp op rep
    forall (op :: * -> *) rep. SizeOp -> HostOp op rep
    SizeOp (SizeOp -> HostOp op rep) -> SizeOp -> HostOp op rep
    forall a b. (a -> b) -> a -> b
    $ Map VName VName -> SizeOp -> SizeOp
    forall a. Substitute a => Map VName VName -> a -> a
    substituteNames Map VName VName
    substs SizeOp
    op
      substituteNames Map VName VName
    substs (GPUBody [Type]
    ts Body rep
    body) =
        [Type] -> Body rep -> HostOp op rep
    forall (op :: * -> *) rep. [Type] -> Body rep -> HostOp op rep
    GPUBody (Map VName VName -> [Type] -> [Type]
    forall a. Substitute a => Map VName VName -> a -> a
    substituteNames Map VName VName
    substs [Type]
    ts) (Map VName VName -> Body rep -> Body rep
    forall a. Substitute a => Map VName VName -> a -> a
    substituteNames Map VName VName
    substs Body rep
    body)

    instance (ASTRep rep, Rename (op rep)) => Rename (HostOp op rep) where
      rename :: HostOp op rep -> RenameM (HostOp op rep)
    rename (SegOp SegOp SegLevel rep
    op) = SegOp SegLevel rep -> HostOp op rep
    forall (op :: * -> *) rep. SegOp SegLevel rep -> HostOp op rep
    SegOp (SegOp SegLevel rep -> HostOp op rep)
    -> RenameM (SegOp SegLevel rep) -> RenameM (HostOp op rep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> SegOp SegLevel rep -> RenameM (SegOp SegLevel rep)
    forall a. Rename a => a -> RenameM a
    rename SegOp SegLevel rep
    op
      rename (OtherOp op rep
    op) = op rep -> HostOp op rep
    forall (op :: * -> *) rep. op rep -> HostOp op rep
    OtherOp (op rep -> HostOp op rep)
    -> RenameM (op rep) -> RenameM (HostOp op rep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> op rep -> RenameM (op rep)
    forall a. Rename a => a -> RenameM a
    rename op rep
    op
      rename (SizeOp SizeOp
    op) = SizeOp -> HostOp op rep
    forall (op :: * -> *) rep. SizeOp -> HostOp op rep
    SizeOp (SizeOp -> HostOp op rep)
    -> RenameM SizeOp -> RenameM (HostOp op rep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> SizeOp -> RenameM SizeOp
    forall a. Rename a => a -> RenameM a
    rename SizeOp
    op
      rename (GPUBody [Type]
    ts Body rep
    body) = [Type] -> Body rep -> HostOp op rep
    forall (op :: * -> *) rep. [Type] -> Body rep -> HostOp op rep
    GPUBody ([Type] -> Body rep -> HostOp op rep)
    -> RenameM [Type] -> RenameM (Body rep -> HostOp op rep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> [Type] -> RenameM [Type]
    forall a. Rename a => a -> RenameM a
    rename [Type]
    ts RenameM (Body rep -> HostOp op rep)
    -> RenameM (Body rep) -> RenameM (HostOp op rep)
    forall a b. RenameM (a -> b) -> RenameM a -> RenameM b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> Body rep -> RenameM (Body rep)
    forall a. Rename a => a -> RenameM a
    rename Body rep
    body

    instance (IsOp op) => IsOp (HostOp op) where
      safeOp :: forall rep. ASTRep rep => HostOp op rep -> Bool
    safeOp (SegOp SegOp SegLevel rep
    op) = SegOp SegLevel rep -> Bool
    forall rep. ASTRep rep => SegOp SegLevel rep -> Bool
    forall (op :: * -> *) rep. (IsOp op, ASTRep rep) => op rep -> Bool
    safeOp SegOp SegLevel rep
    op
      safeOp (OtherOp op rep
    op) = op rep -> Bool
    forall rep. ASTRep rep => op rep -> Bool
    forall (op :: * -> *) rep. (IsOp op, ASTRep rep) => op rep -> Bool
    safeOp op rep
    op
      safeOp (SizeOp SizeOp
    _) = Bool
    True
      safeOp (GPUBody [Type]
    _ Body rep
    body) = (Stm rep -> Bool) -> Seq (Stm rep) -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (Exp rep -> Bool
    forall rep. ASTRep rep => Exp rep -> Bool
    safeExp (Exp rep -> Bool) -> (Stm rep -> Exp rep) -> Stm rep -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Stm rep -> Exp rep
    forall rep. Stm rep -> Exp rep
    stmExp) (Seq (Stm rep) -> Bool) -> Seq (Stm rep) -> Bool
    forall a b. (a -> b) -> a -> b
    $ Body rep -> Seq (Stm rep)
    forall rep. Body rep -> Stms rep
    bodyStms Body rep
    body

      cheapOp :: forall rep. ASTRep rep => HostOp op rep -> Bool
    cheapOp (SegOp SegOp SegLevel rep
    op) = SegOp SegLevel rep -> Bool
    forall rep. ASTRep rep => SegOp SegLevel rep -> Bool
    forall (op :: * -> *) rep. (IsOp op, ASTRep rep) => op rep -> Bool
    cheapOp SegOp SegLevel rep
    op
      cheapOp (OtherOp op rep
    op) = op rep -> Bool
    forall rep. ASTRep rep => op rep -> Bool
    forall (op :: * -> *) rep. (IsOp op, ASTRep rep) => op rep -> Bool
    cheapOp op rep
    op
      cheapOp (SizeOp SizeOp
    _) = Bool
    True
      cheapOp (GPUBody [Type]
    types Body rep
    body) =
        -- Current GPUBody usage only benefits from hoisting kernels that
        -- transfer scalars to device.
        Seq (Stm rep) -> Bool
    forall a. Seq a -> Bool
    SQ.null (Body rep -> Seq (Stm rep)
    forall rep. Body rep -> Stms rep
    bodyStms Body rep
    body) Bool -> Bool -> Bool
    && (Type -> Bool) -> [Type] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all ((Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    0) (Int -> Bool) -> (Type -> Int) -> Type -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Type -> Int
    forall shape u. ArrayShape shape => TypeBase shape u -> Int
    arrayRank) [Type]
    types

      opDependencies :: forall rep. ASTRep rep => HostOp op rep -> [Names]
    opDependencies (SegOp SegOp SegLevel rep
    op) = SegOp SegLevel rep -> [Names]
    forall rep. ASTRep rep => SegOp SegLevel rep -> [Names]
    forall (op :: * -> *) rep.
    (IsOp op, ASTRep rep) =>
    op rep -> [Names]
    opDependencies SegOp SegLevel rep
    op
      opDependencies (OtherOp op rep
    op) = op rep -> [Names]
    forall rep. ASTRep rep => op rep -> [Names]
    forall (op :: * -> *) rep.
    (IsOp op, ASTRep rep) =>
    op rep -> [Names]
    opDependencies op rep
    op
      opDependencies (SizeOp SizeOp
    op) = [SizeOp -> Names
    forall a. FreeIn a => a -> Names
    freeIn SizeOp
    op]
      opDependencies (GPUBody [Type]
    _ Body rep
    body) =
        Int -> Names -> [Names]
    forall a. Int -> a -> [a]
    replicate ([SubExpRes] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length ([SubExpRes] -> Int)
    -> (Body rep -> [SubExpRes]) -> Body rep -> Int
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Body rep -> [SubExpRes]
    forall rep. Body rep -> [SubExpRes]
    bodyResult (Body rep -> Int) -> Body rep -> Int
    forall a b. (a -> b) -> a -> b
    $ Body rep
    body) (Body rep -> Names
    forall a. FreeIn a => a -> Names
    freeIn Body rep
    body)

    instance (TypedOp op) => TypedOp (HostOp op) where
      opType :: forall rep (m :: * -> *).
    HasScope rep m =>
    HostOp op rep -> m [ExtType]
    opType (SegOp SegOp SegLevel rep
    op) = SegOp SegLevel rep -> m [ExtType]
    forall rep (m :: * -> *).
    HasScope rep m =>
    SegOp SegLevel rep -> m [ExtType]
    forall (op :: * -> *) rep (m :: * -> *).
    (TypedOp op, HasScope rep m) =>
    op rep -> m [ExtType]
    opType SegOp SegLevel rep
    op
      opType (OtherOp op rep
    op) = op rep -> m [ExtType]
    forall rep (m :: * -> *). HasScope rep m => op rep -> m [ExtType]
    forall (op :: * -> *) rep (m :: * -> *).
    (TypedOp op, HasScope rep m) =>
    op rep -> m [ExtType]
    opType op rep
    op
      opType (SizeOp (GetSize Name
    _ SizeClass
    _)) = [ExtType] -> m [ExtType]
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [PrimType -> ExtType
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64]
      opType (SizeOp (GetSizeMax SizeClass
    _)) = [ExtType] -> m [ExtType]
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [PrimType -> ExtType
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64]
      opType (SizeOp CmpSizeLe {}) = [ExtType] -> m [ExtType]
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [PrimType -> ExtType
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    Bool]
      opType (SizeOp (CalcNumBlocks {})) = [ExtType] -> m [ExtType]
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [PrimType -> ExtType
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64]
      opType (GPUBody [Type]
    ts Body rep
    _) =
        [ExtType] -> m [ExtType]
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([ExtType] -> m [ExtType]) -> [ExtType] -> m [ExtType]
    forall a b. (a -> b) -> a -> b
    $ [Type] -> [ExtType]
    forall u. [TypeBase (ShapeBase SubExp) u] -> [TypeBase ExtShape u]
    staticShapes ([Type] -> [ExtType]) -> [Type] -> [ExtType]
    forall a b. (a -> b) -> a -> b
    $ (Type -> Type) -> [Type] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map (Type -> SubExp -> Type
    forall d.
    ArrayShape (ShapeBase d) =>
    TypeBase (ShapeBase d) NoUniqueness
    -> d -> TypeBase (ShapeBase d) NoUniqueness
    `arrayOfRow` IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    1) [Type]
    ts

    instance (AliasedOp op) => AliasedOp (HostOp op) where
      opAliases :: forall rep. Aliased rep => HostOp op rep -> [Names]
    opAliases (SegOp SegOp SegLevel rep
    op) = SegOp SegLevel rep -> [Names]
    forall rep. Aliased rep => SegOp SegLevel rep -> [Names]
    forall (op :: * -> *) rep.
    (AliasedOp op, Aliased rep) =>
    op rep -> [Names]
    opAliases SegOp SegLevel rep
    op
      opAliases (OtherOp op rep
    op) = op rep -> [Names]
    forall rep. Aliased rep => op rep -> [Names]
    forall (op :: * -> *) rep.
    (AliasedOp op, Aliased rep) =>
    op rep -> [Names]
    opAliases op rep
    op
      opAliases (SizeOp SizeOp
    _) = [Names
    forall a. Monoid a => a
    mempty]
      opAliases (GPUBody [Type]
    ts Body rep
    _) = (Type -> Names) -> [Type] -> [Names]
    forall a b. (a -> b) -> [a] -> [b]
    map (Names -> Type -> Names
    forall a b. a -> b -> a
    const Names
    forall a. Monoid a => a
    mempty) [Type]
    ts

      consumedInOp :: forall rep. Aliased rep => HostOp op rep -> Names
    consumedInOp (SegOp SegOp SegLevel rep
    op) = SegOp SegLevel rep -> Names
    forall rep. Aliased rep => SegOp SegLevel rep -> Names
    forall (op :: * -> *) rep.
    (AliasedOp op, Aliased rep) =>
    op rep -> Names
    consumedInOp SegOp SegLevel rep
    op
      consumedInOp (OtherOp op rep
    op) = op rep -> Names
    forall rep. Aliased rep => op rep -> Names
    forall (op :: * -> *) rep.
    (AliasedOp op, Aliased rep) =>
    op rep -> Names
    consumedInOp op rep
    op
      consumedInOp (SizeOp SizeOp
    _) = Names
    forall a. Monoid a => a
    mempty
      consumedInOp (GPUBody [Type]
    _ Body rep
    body) = Body rep -> Names
    forall rep. Aliased rep => Body rep -> Names
    consumedInBody Body rep
    body

    instance (ASTRep rep, FreeIn (op rep)) => FreeIn (HostOp op rep) where
      freeIn' :: HostOp op rep -> FV
    freeIn' (SegOp SegOp SegLevel rep
    op) = SegOp SegLevel rep -> FV
    forall a. FreeIn a => a -> FV
    freeIn' SegOp SegLevel rep
    op
      freeIn' (OtherOp op rep
    op) = op rep -> FV
    forall a. FreeIn a => a -> FV
    freeIn' op rep
    op
      freeIn' (SizeOp SizeOp
    op) = SizeOp -> FV
    forall a. FreeIn a => a -> FV
    freeIn' SizeOp
    op
      freeIn' (GPUBody [Type]
    ts Body rep
    body) = [Type] -> FV
    forall a. FreeIn a => a -> FV
    freeIn' [Type]
    ts FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> Body rep -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Body rep
    body

    instance (CanBeAliased op) => CanBeAliased (HostOp op) where
      addOpAliases :: forall rep.
    AliasableRep rep =>
    AliasTable -> HostOp op rep -> HostOp op (Aliases rep)
    addOpAliases AliasTable
    aliases (SegOp SegOp SegLevel rep
    op) = SegOp SegLevel (Aliases rep) -> HostOp op (Aliases rep)
    forall (op :: * -> *) rep. SegOp SegLevel rep -> HostOp op rep
    SegOp (SegOp SegLevel (Aliases rep) -> HostOp op (Aliases rep))
    -> SegOp SegLevel (Aliases rep) -> HostOp op (Aliases rep)
    forall a b. (a -> b) -> a -> b
    $ AliasTable -> SegOp SegLevel rep -> SegOp SegLevel (Aliases rep)
    forall rep.
    AliasableRep rep =>
    AliasTable -> SegOp SegLevel rep -> SegOp SegLevel (Aliases rep)
    forall (op :: * -> *) rep.
    (CanBeAliased op, AliasableRep rep) =>
    AliasTable -> op rep -> op (Aliases rep)
    addOpAliases AliasTable
    aliases SegOp SegLevel rep
    op
      addOpAliases AliasTable
    aliases (GPUBody [Type]
    ts Body rep
    body) = [Type] -> Body (Aliases rep) -> HostOp op (Aliases rep)
    forall (op :: * -> *) rep. [Type] -> Body rep -> HostOp op rep
    GPUBody [Type]
    ts (Body (Aliases rep) -> HostOp op (Aliases rep))
    -> Body (Aliases rep) -> HostOp op (Aliases rep)
    forall a b. (a -> b) -> a -> b
    $ AliasTable -> Body rep -> Body (Aliases rep)
    forall rep.
    AliasableRep rep =>
    AliasTable -> Body rep -> Body (Aliases rep)
    Alias.analyseBody AliasTable
    aliases Body rep
    body
      addOpAliases AliasTable
    aliases (OtherOp op rep
    op) = op (Aliases rep) -> HostOp op (Aliases rep)
    forall (op :: * -> *) rep. op rep -> HostOp op rep
    OtherOp (op (Aliases rep) -> HostOp op (Aliases rep))
    -> op (Aliases rep) -> HostOp op (Aliases rep)
    forall a b. (a -> b) -> a -> b
    $ AliasTable -> op rep -> op (Aliases rep)
    forall rep.
    AliasableRep rep =>
    AliasTable -> op rep -> op (Aliases rep)
    forall (op :: * -> *) rep.
    (CanBeAliased op, AliasableRep rep) =>
    AliasTable -> op rep -> op (Aliases rep)
    addOpAliases AliasTable
    aliases op rep
    op
      addOpAliases AliasTable
    _ (SizeOp SizeOp
    op) = SizeOp -> HostOp op (Aliases rep)
    forall (op :: * -> *) rep. SizeOp -> HostOp op rep
    SizeOp SizeOp
    op

    instance (CanBeWise op) => CanBeWise (HostOp op) where
      addOpWisdom :: forall rep. Informing rep => HostOp op rep -> HostOp op (Wise rep)
    addOpWisdom (SegOp SegOp SegLevel rep
    op) = SegOp SegLevel (Wise rep) -> HostOp op (Wise rep)
    forall (op :: * -> *) rep. SegOp SegLevel rep -> HostOp op rep
    SegOp (SegOp SegLevel (Wise rep) -> HostOp op (Wise rep))
    -> SegOp SegLevel (Wise rep) -> HostOp op (Wise rep)
    forall a b. (a -> b) -> a -> b
    $ SegOp SegLevel rep -> SegOp SegLevel (Wise rep)
    forall rep.
    Informing rep =>
    SegOp SegLevel rep -> SegOp SegLevel (Wise rep)
    forall (op :: * -> *) rep.
    (CanBeWise op, Informing rep) =>
    op rep -> op (Wise rep)
    addOpWisdom SegOp SegLevel rep
    op
      addOpWisdom (OtherOp op rep
    op) = op (Wise rep) -> HostOp op (Wise rep)
    forall (op :: * -> *) rep. op rep -> HostOp op rep
    OtherOp (op (Wise rep) -> HostOp op (Wise rep))
    -> op (Wise rep) -> HostOp op (Wise rep)
    forall a b. (a -> b) -> a -> b
    $ op rep -> op (Wise rep)
    forall rep. Informing rep => op rep -> op (Wise rep)
    forall (op :: * -> *) rep.
    (CanBeWise op, Informing rep) =>
    op rep -> op (Wise rep)
    addOpWisdom op rep
    op
      addOpWisdom (SizeOp SizeOp
    op) = SizeOp -> HostOp op (Wise rep)
    forall (op :: * -> *) rep. SizeOp -> HostOp op rep
    SizeOp SizeOp
    op
      addOpWisdom (GPUBody [Type]
    ts Body rep
    body) = [Type] -> Body (Wise rep) -> HostOp op (Wise rep)
    forall (op :: * -> *) rep. [Type] -> Body rep -> HostOp op rep
    GPUBody [Type]
    ts (Body (Wise rep) -> HostOp op (Wise rep))
    -> Body (Wise rep) -> HostOp op (Wise rep)
    forall a b. (a -> b) -> a -> b
    $ Body rep -> Body (Wise rep)
    forall rep. Informing rep => Body rep -> Body (Wise rep)
    informBody Body rep
    body

    instance OpReturns (HostOp NoOp) where
      opReturns :: forall rep (inner :: * -> *) (m :: * -> *).
    (Mem rep inner, Monad m, HasScope rep m) =>
    HostOp NoOp rep -> m [ExpReturns]
    opReturns (SegOp SegOp SegLevel rep
    op) = SegOp SegLevel rep -> m [ExpReturns]
    forall rep (inner :: * -> *) (m :: * -> *) lvl.
    (Mem rep inner, Monad m, HasScope rep m) =>
    SegOp lvl rep -> m [ExpReturns]
    segOpReturns SegOp SegLevel rep
    op
      opReturns HostOp NoOp rep
    k = [ExtType] -> [ExpReturns]
    extReturns ([ExtType] -> [ExpReturns]) -> m [ExtType] -> m [ExpReturns]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> HostOp NoOp rep -> m [ExtType]
    forall rep (m :: * -> *).
    HasScope rep m =>
    HostOp NoOp rep -> m [ExtType]
    forall (op :: * -> *) rep (m :: * -> *).
    (TypedOp op, HasScope rep m) =>
    op rep -> m [ExtType]
    opType HostOp NoOp rep
    k

    instance (ASTRep rep, ST.IndexOp (op rep)) => ST.IndexOp (HostOp op rep) where
      indexOp :: forall rep.
    (ASTRep rep, IndexOp (Op rep)) =>
    SymbolTable rep
    -> Int -> HostOp op rep -> [TPrimExp Int64 VName] -> Maybe Indexed
    indexOp SymbolTable rep
    vtable Int
    k (SegOp SegOp SegLevel rep
    op) [TPrimExp Int64 VName]
    is = SymbolTable rep
    -> Int
    -> SegOp SegLevel rep
    -> [TPrimExp Int64 VName]
    -> Maybe Indexed
    forall rep.
    (ASTRep rep, IndexOp (Op rep)) =>
    SymbolTable rep
    -> Int
    -> SegOp SegLevel rep
    -> [TPrimExp Int64 VName]
    -> Maybe Indexed
    forall op rep.
    (IndexOp op, ASTRep rep, IndexOp (Op rep)) =>
    SymbolTable rep
    -> Int -> op -> [TPrimExp Int64 VName] -> Maybe Indexed
    ST.indexOp SymbolTable rep
    vtable Int
    k SegOp SegLevel rep
    op [TPrimExp Int64 VName]
    is
      indexOp SymbolTable rep
    vtable Int
    k (OtherOp op rep
    op) [TPrimExp Int64 VName]
    is = SymbolTable rep
    -> Int -> op rep -> [TPrimExp Int64 VName] -> Maybe Indexed
    forall rep.
    (ASTRep rep, IndexOp (Op rep)) =>
    SymbolTable rep
    -> Int -> op rep -> [TPrimExp Int64 VName] -> Maybe Indexed
    forall op rep.
    (IndexOp op, ASTRep rep, IndexOp (Op rep)) =>
    SymbolTable rep
    -> Int -> op -> [TPrimExp Int64 VName] -> Maybe Indexed
    ST.indexOp SymbolTable rep
    vtable Int
    k op rep
    op [TPrimExp Int64 VName]
    is
      indexOp SymbolTable rep
    _ Int
    _ HostOp op rep
    _ [TPrimExp Int64 VName]
    _ = Maybe Indexed
    forall a. Maybe a
    Nothing

    instance (PrettyRep rep, PP.Pretty (op rep)) => PP.Pretty (HostOp op rep) where
      pretty :: forall ann. HostOp op rep -> Doc ann
    pretty (SegOp SegOp SegLevel rep
    op) = SegOp SegLevel rep -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SegOp SegLevel rep -> Doc ann
    pretty SegOp SegLevel rep
    op
      pretty (OtherOp op rep
    op) = op rep -> Doc ann
    forall ann. op rep -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty op rep
    op
      pretty (SizeOp SizeOp
    op) = SizeOp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SizeOp -> Doc ann
    pretty SizeOp
    op
      pretty (GPUBody [Type]
    ts Body rep
    body) =
        Doc ann
    "gpu" Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    forall ann. Doc ann
    PP.colon Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> [Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    ppTuple' ((Type -> Doc ann) -> [Type] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map Type -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Type -> Doc ann
    pretty [Type]
    ts) Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann -> Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a -> Doc a
    PP.nestedBlock Doc ann
    "{" Doc ann
    "}" (Body rep -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Body rep -> Doc ann
    pretty Body rep
    body)

    instance (OpMetrics (Op rep), OpMetrics (op rep)) => OpMetrics (HostOp op rep) where
      opMetrics :: HostOp op rep -> MetricsM ()
    opMetrics (SegOp SegOp SegLevel rep
    op) = SegOp SegLevel rep -> MetricsM ()
    forall op. OpMetrics op => op -> MetricsM ()
    opMetrics SegOp SegLevel rep
    op
      opMetrics (OtherOp op rep
    op) = op rep -> MetricsM ()
    forall op. OpMetrics op => op -> MetricsM ()
    opMetrics op rep
    op
      opMetrics (SizeOp SizeOp
    op) = SizeOp -> MetricsM ()
    forall op. OpMetrics op => op -> MetricsM ()
    opMetrics SizeOp
    op
      opMetrics (GPUBody [Type]
    _ Body rep
    body) = Text -> MetricsM () -> MetricsM ()
    inside Text
    "GPUBody" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
    forall a b. (a -> b) -> a -> b
    $ Body rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Body rep -> MetricsM ()
    bodyMetrics Body rep
    body

    instance (RephraseOp op) => RephraseOp (HostOp op) where
      rephraseInOp :: forall (m :: * -> *) from to.
    Monad m =>
    Rephraser m from to -> HostOp op from -> m (HostOp op to)
    rephraseInOp Rephraser m from to
    r (SegOp SegOp SegLevel from
    op) = SegOp SegLevel to -> HostOp op to
    forall (op :: * -> *) rep. SegOp SegLevel rep -> HostOp op rep
    SegOp (SegOp SegLevel to -> HostOp op to)
    -> m (SegOp SegLevel to) -> m (HostOp op to)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Rephraser m from to -> SegOp SegLevel from -> m (SegOp SegLevel to)
    forall (op :: * -> *) (m :: * -> *) from to.
    (RephraseOp op, Monad m) =>
    Rephraser m from to -> op from -> m (op to)
    forall (m :: * -> *) from to.
    Monad m =>
    Rephraser m from to -> SegOp SegLevel from -> m (SegOp SegLevel to)
    rephraseInOp Rephraser m from to
    r SegOp SegLevel from
    op
      rephraseInOp Rephraser m from to
    r (OtherOp op from
    op) = op to -> HostOp op to
    forall (op :: * -> *) rep. op rep -> HostOp op rep
    OtherOp (op to -> HostOp op to) -> m (op to) -> m (HostOp op to)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Rephraser m from to -> op from -> m (op to)
    forall (op :: * -> *) (m :: * -> *) from to.
    (RephraseOp op, Monad m) =>
    Rephraser m from to -> op from -> m (op to)
    forall (m :: * -> *) from to.
    Monad m =>
    Rephraser m from to -> op from -> m (op to)
    rephraseInOp Rephraser m from to
    r op from
    op
      rephraseInOp Rephraser m from to
    _ (SizeOp SizeOp
    op) = HostOp op to -> m (HostOp op to)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (HostOp op to -> m (HostOp op to))
    -> HostOp op to -> m (HostOp op to)
    forall a b. (a -> b) -> a -> b
    $ SizeOp -> HostOp op to
    forall (op :: * -> *) rep. SizeOp -> HostOp op rep
    SizeOp SizeOp
    op
      rephraseInOp Rephraser m from to
    r (GPUBody [Type]
    ts Body from
    body) = [Type] -> Body to -> HostOp op to
    forall (op :: * -> *) rep. [Type] -> Body rep -> HostOp op rep
    GPUBody [Type]
    ts (Body to -> HostOp op to) -> m (Body to) -> m (HostOp op to)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Rephraser m from to -> Body from -> m (Body to)
    forall (m :: * -> *) from to.
    Monad m =>
    Rephraser m from to -> Body from -> m (Body to)
    rephraseBody Rephraser m from to
    r Body from
    body

    checkGrid :: (TC.Checkable rep) => KernelGrid -> TC.TypeM rep ()
    checkGrid :: forall rep. Checkable rep => KernelGrid -> TypeM rep ()
    checkGrid KernelGrid
    grid = do
      [Type] -> SubExp -> TypeM rep ()
    forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
    TC.require [PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64] (SubExp -> TypeM rep ()) -> SubExp -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ Count NumBlocks SubExp -> SubExp
    forall {k} (u :: k) e. Count u e -> e
    unCount (Count NumBlocks SubExp -> SubExp)
    -> Count NumBlocks SubExp -> SubExp
    forall a b. (a -> b) -> a -> b
    $ KernelGrid -> Count NumBlocks SubExp
    gridNumBlocks KernelGrid
    grid
      [Type] -> SubExp -> TypeM rep ()
    forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
    TC.require [PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64] (SubExp -> TypeM rep ()) -> SubExp -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ Count BlockSize SubExp -> SubExp
    forall {k} (u :: k) e. Count u e -> e
    unCount (Count BlockSize SubExp -> SubExp)
    -> Count BlockSize SubExp -> SubExp
    forall a b. (a -> b) -> a -> b
    $ KernelGrid -> Count BlockSize SubExp
    gridBlockSize KernelGrid
    grid

    checkSegLevel ::
      (TC.Checkable rep) =>
      Maybe SegLevel ->
      SegLevel ->
      TC.TypeM rep ()
    checkSegLevel :: forall rep.
    Checkable rep =>
    Maybe SegLevel -> SegLevel -> TypeM rep ()
    checkSegLevel (Just SegBlock {}) (SegThreadInBlock SegVirt
    _virt) =
      () -> TypeM rep ()
    forall a. a -> TypeM rep a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
    checkSegLevel Maybe SegLevel
    _ (SegThreadInBlock SegVirt
    _virt) =
      ErrorCase rep -> TypeM rep ()
    forall rep a. ErrorCase rep -> TypeM rep a
    TC.bad (ErrorCase rep -> TypeM rep ()) -> ErrorCase rep -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ Text -> ErrorCase rep
    forall rep. Text -> ErrorCase rep
    TC.TypeError Text
    "inblock SegOp not in block SegOp."
    checkSegLevel (Just SegThread {}) SegLevel
    _ =
      ErrorCase rep -> TypeM rep ()
    forall rep a. ErrorCase rep -> TypeM rep a
    TC.bad (ErrorCase rep -> TypeM rep ()) -> ErrorCase rep -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ Text -> ErrorCase rep
    forall rep. Text -> ErrorCase rep
    TC.TypeError Text
    "SegOps cannot occur when already at thread level."
    checkSegLevel (Just SegThreadInBlock {}) SegLevel
    _ =
      ErrorCase rep -> TypeM rep ()
    forall rep a. ErrorCase rep -> TypeM rep a
    TC.bad (ErrorCase rep -> TypeM rep ()) -> ErrorCase rep -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ Text -> ErrorCase rep
    forall rep. Text -> ErrorCase rep
    TC.TypeError Text
    "SegOps cannot occur when already at inblock level."
    checkSegLevel Maybe SegLevel
    _ (SegThread SegVirt
    _virt Maybe KernelGrid
    Nothing) =
      () -> TypeM rep ()
    forall a. a -> TypeM rep a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
    checkSegLevel (Just SegLevel
    _) SegThread {} =
      ErrorCase rep -> TypeM rep ()
    forall rep a. ErrorCase rep -> TypeM rep a
    TC.bad (ErrorCase rep -> TypeM rep ()) -> ErrorCase rep -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ Text -> ErrorCase rep
    forall rep. Text -> ErrorCase rep
    TC.TypeError Text
    "thread-level SegOp cannot be nested"
    checkSegLevel Maybe SegLevel
    Nothing (SegThread SegVirt
    _virt Maybe KernelGrid
    grid) =
      (KernelGrid -> TypeM rep ()) -> Maybe KernelGrid -> TypeM rep ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ KernelGrid -> TypeM rep ()
    forall rep. Checkable rep => KernelGrid -> TypeM rep ()
    checkGrid Maybe KernelGrid
    grid
    checkSegLevel (Just SegLevel
    _) SegBlock {} =
      ErrorCase rep -> TypeM rep ()
    forall rep a. ErrorCase rep -> TypeM rep a
    TC.bad (ErrorCase rep -> TypeM rep ()) -> ErrorCase rep -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ Text -> ErrorCase rep
    forall rep. Text -> ErrorCase rep
    TC.TypeError Text
    "block-level SegOp cannot be nested"
    checkSegLevel Maybe SegLevel
    Nothing (SegBlock SegVirt
    _virt Maybe KernelGrid
    grid) =
      (KernelGrid -> TypeM rep ()) -> Maybe KernelGrid -> TypeM rep ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ KernelGrid -> TypeM rep ()
    forall rep. Checkable rep => KernelGrid -> TypeM rep ()
    checkGrid Maybe KernelGrid
    grid

    typeCheckHostOp ::
      (TC.Checkable rep) =>
      (SegLevel -> Op (Aliases rep) -> TC.TypeM rep ()) ->
      Maybe SegLevel ->
      (op (Aliases rep) -> TC.TypeM rep ()) ->
      HostOp op (Aliases rep) ->
      TC.TypeM rep ()
    typeCheckHostOp :: forall rep (op :: * -> *).
    Checkable rep =>
    (SegLevel -> Op (Aliases rep) -> TypeM rep ())
    -> Maybe SegLevel
    -> (op (Aliases rep) -> TypeM rep ())
    -> HostOp op (Aliases rep)
    -> TypeM rep ()
    typeCheckHostOp SegLevel -> Op (Aliases rep) -> TypeM rep ()
    checker Maybe SegLevel
    lvl op (Aliases rep) -> TypeM rep ()
    _ (SegOp SegOp SegLevel (Aliases rep)
    op) =
      (Op (Aliases rep) -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
    forall rep a.
    (Op (Aliases rep) -> TypeM rep ()) -> TypeM rep a -> TypeM rep a
    TC.checkOpWith (SegLevel -> Op (Aliases rep) -> TypeM rep ()
    checker (SegLevel -> Op (Aliases rep) -> TypeM rep ())
    -> SegLevel -> Op (Aliases rep) -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ SegOp SegLevel (Aliases rep) -> SegLevel
    forall lvl rep. SegOp lvl rep -> lvl
    segLevel SegOp SegLevel (Aliases rep)
    op) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $
        (SegLevel -> TypeM rep ())
    -> SegOp SegLevel (Aliases rep) -> TypeM rep ()
    forall rep lvl.
    Checkable rep =>
    (lvl -> TypeM rep ()) -> SegOp lvl (Aliases rep) -> TypeM rep ()
    typeCheckSegOp (Maybe SegLevel -> SegLevel -> TypeM rep ()
    forall rep.
    Checkable rep =>
    Maybe SegLevel -> SegLevel -> TypeM rep ()
    checkSegLevel Maybe SegLevel
    lvl) SegOp SegLevel (Aliases rep)
    op
    typeCheckHostOp SegLevel -> Op (Aliases rep) -> TypeM rep ()
    _ Just {} op (Aliases rep) -> TypeM rep ()
    _ GPUBody {} =
      ErrorCase rep -> TypeM rep ()
    forall rep a. ErrorCase rep -> TypeM rep a
    TC.bad (ErrorCase rep -> TypeM rep ()) -> ErrorCase rep -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ Text -> ErrorCase rep
    forall rep. Text -> ErrorCase rep
    TC.TypeError Text
    "GPUBody may not be nested in SegOps."
    typeCheckHostOp SegLevel -> Op (Aliases rep) -> TypeM rep ()
    _ Maybe SegLevel
    _ op (Aliases rep) -> TypeM rep ()
    f (OtherOp op (Aliases rep)
    op) = op (Aliases rep) -> TypeM rep ()
    f op (Aliases rep)
    op
    typeCheckHostOp SegLevel -> Op (Aliases rep) -> TypeM rep ()
    _ Maybe SegLevel
    _ op (Aliases rep) -> TypeM rep ()
    _ (SizeOp SizeOp
    op) = SizeOp -> TypeM rep ()
    forall rep. Checkable rep => SizeOp -> TypeM rep ()
    typeCheckSizeOp SizeOp
    op
    typeCheckHostOp SegLevel -> Op (Aliases rep) -> TypeM rep ()
    _ Maybe SegLevel
    Nothing op (Aliases rep) -> TypeM rep ()
    _ (GPUBody [Type]
    ts Body (Aliases rep)
    body) = do
      (Type -> TypeM rep ()) -> [Type] -> TypeM rep ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ Type -> TypeM rep ()
    forall rep u.
    Checkable rep =>
    TypeBase (ShapeBase SubExp) u -> TypeM rep ()
    TC.checkType [Type]
    ts
      TypeM rep [Names] -> TypeM rep ()
    forall (f :: * -> *) a. Functor f => f a -> f ()
    void (TypeM rep [Names] -> TypeM rep ())
    -> TypeM rep [Names] -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ Body (Aliases rep) -> TypeM rep [Names]
    forall rep.
    Checkable rep =>
    Body (Aliases rep) -> TypeM rep [Names]
    TC.checkBody Body (Aliases rep)
    body
      body_ts <-
        ExtendedScope (Aliases rep) (TypeM rep) [Type]
    -> Scope (Aliases rep) -> TypeM rep [Type]
    forall rep (m :: * -> *) a.
    ExtendedScope rep m a -> Scope rep -> m a
    extendedScope
          ((SubExpRes -> ExtendedScope (Aliases rep) (TypeM rep) Type)
    -> [SubExpRes] -> ExtendedScope (Aliases rep) (TypeM rep) [Type]
    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) -> [a] -> f [b]
    traverse SubExpRes -> ExtendedScope (Aliases rep) (TypeM rep) Type
    forall t (m :: * -> *). HasScope t m => SubExpRes -> m Type
    subExpResType (Body (Aliases rep) -> [SubExpRes]
    forall rep. Body rep -> [SubExpRes]
    bodyResult Body (Aliases rep)
    body))
          (Stms (Aliases rep) -> Scope (Aliases rep)
    forall rep a. Scoped rep a => a -> Scope rep
    scopeOf (Body (Aliases rep) -> Stms (Aliases rep)
    forall rep. Body rep -> Stms rep
    bodyStms Body (Aliases rep)
    body))
      unless (body_ts == ts) . TC.bad . TC.TypeError . T.unlines $
        [ "Expected type: " <> prettyTuple ts,
          "Got body type: " <> prettyTuple body_ts
        ]