    {-# LANGUAGE TypeFamilies #-}

    -- | Code generation for 'SegMap' is quite straightforward.  The only
    -- trick is virtualisation in case the physical number of threads is
    -- not sufficient to cover the logical thread space.  This is handled
    -- by having actual threadblocks run a loop to imitate multiple threadblocks.
    module Futhark.CodeGen.ImpGen.GPU.SegMap (compileSegMap) where

    import Control.Monad
    import Futhark.CodeGen.ImpCode.GPU qualified as Imp
    import Futhark.CodeGen.ImpGen
    import Futhark.CodeGen.ImpGen.GPU.Base
    import Futhark.CodeGen.ImpGen.GPU.Block
    import Futhark.IR.GPUMem
    import Futhark.Util.IntegralExp (divUp)
    import Prelude hiding (quot, rem)

    -- | Compile 'SegMap' instance code.
    compileSegMap ::
      Pat LetDecMem ->
      SegLevel ->
      SegSpace ->
      KernelBody GPUMem ->
      CallKernelGen ()
    compileSegMap :: Pat LetDecMem
    -> SegLevel -> SegSpace -> KernelBody GPUMem -> CallKernelGen ()
    compileSegMap Pat LetDecMem
    pat SegLevel
    lvl SegSpace
    space KernelBody GPUMem
    kbody = do
      attrs <- SegLevel -> CallKernelGen KernelAttrs
    lvlKernelAttrs SegLevel
    lvl

      let (is, dims) = unzip $ unSegSpace space
          dims' = (SubExp -> TPrimExp Int64 VName)
    -> [SubExp] -> [TPrimExp Int64 VName]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> TPrimExp Int64 VName
    pe64 [SubExp]
    dims
          tblock_size' = SubExp -> TPrimExp Int64 VName
    pe64 (SubExp -> TPrimExp Int64 VName)
    -> Count BlockSize SubExp -> Count BlockSize (TPrimExp Int64 VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> KernelAttrs -> Count BlockSize SubExp
    kAttrBlockSize KernelAttrs
    attrs

      emit $ Imp.DebugPrint "\n# SegMap" Nothing
      case lvl of
        SegThread {} -> do
          virt_num_tblocks <- String -> TExp Int32 -> ImpM GPUMem HostEnv HostOp (TExp Int32)
    forall {k} (t :: k) rep r op.
    String -> TExp t -> ImpM rep r op (TExp t)
    dPrimVE String
    "virt_num_tblocks" (TExp Int32 -> ImpM GPUMem HostEnv HostOp (TExp Int32))
    -> TExp Int32 -> ImpM GPUMem HostEnv HostOp (TExp Int32)
    forall a b. (a -> b) -> a -> b
    $ TPrimExp Int64 VName -> TExp Int32
    forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int32 v
    sExt32 (TPrimExp Int64 VName -> TExp Int32)
    -> TPrimExp Int64 VName -> TExp Int32
    forall a b. (a -> b) -> a -> b
    $ [TPrimExp Int64 VName] -> TPrimExp Int64 VName
    forall a. Num a => [a] -> a
    forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
    product [TPrimExp Int64 VName]
    dims' TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall e. IntegralExp e => e -> e -> e
    `divUp` Count BlockSize (TPrimExp Int64 VName) -> TPrimExp Int64 VName
    forall {k} (u :: k) e. Count u e -> e
    unCount Count BlockSize (TPrimExp Int64 VName)
    tblock_size'
          sKernelThread "segmap" (segFlat space) attrs $
            virtualiseBlocks (segVirt lvl) virt_num_tblocks $ \TExp Int32
    tblock_id -> do
              local_tid <- KernelConstants -> TExp Int32
    kernelLocalThreadId (KernelConstants -> TExp Int32)
    -> (KernelEnv -> KernelConstants) -> KernelEnv -> TExp Int32
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . KernelEnv -> KernelConstants
    kernelConstants (KernelEnv -> TExp Int32)
    -> ImpM GPUMem KernelEnv KernelOp KernelEnv
    -> ImpM GPUMem KernelEnv KernelOp (TExp Int32)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> ImpM GPUMem KernelEnv KernelOp KernelEnv
    forall rep r op. ImpM rep r op r
    askEnv

              global_tid <-
                dPrimVE "global_tid" $
                  sExt64 tblock_id * sExt64 (unCount tblock_size')
                    + sExt64 local_tid

              dIndexSpace (zip is dims') global_tid

              sWhen (isActive $ unSegSpace space) $
                compileStms mempty (kernelBodyStms kbody) $
                  zipWithM_ (compileThreadResult space) (patElems pat) $
                    kernelBodyResult kbody
        SegBlock {} -> do
          pc <- Count BlockSize (TPrimExp Int64 VName)
    -> Stms GPUMem -> CallKernelGen Precomputed
    precomputeConstants Count BlockSize (TPrimExp Int64 VName)
    tblock_size' (Stms GPUMem -> CallKernelGen Precomputed)
    -> Stms GPUMem -> CallKernelGen Precomputed
    forall a b. (a -> b) -> a -> b
    $ KernelBody GPUMem -> Stms GPUMem
    forall rep. KernelBody rep -> Stms rep
    kernelBodyStms KernelBody GPUMem
    kbody
          virt_num_tblocks <- dPrimVE "virt_num_tblocks" $ sExt32 $ product dims'
          sKernelBlock "segmap_intrablock" (segFlat space) attrs $ do
            precomputedConstants pc $
              virtualiseBlocks (segVirt lvl) virt_num_tblocks $ \TExp Int32
    tblock_id -> do
                [(VName, TPrimExp Int64 VName)]
    -> TPrimExp Int64 VName -> InKernelGen ()
    forall rep r op.
    [(VName, TPrimExp Int64 VName)]
    -> TPrimExp Int64 VName -> ImpM rep r op ()
    dIndexSpace ([VName]
    -> [TPrimExp Int64 VName] -> [(VName, TPrimExp Int64 VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [VName]
    is [TPrimExp Int64 VName]
    dims') (TPrimExp Int64 VName -> InKernelGen ())
    -> TPrimExp Int64 VName -> InKernelGen ()
    forall a b. (a -> b) -> a -> b
    $ TExp Int32 -> TPrimExp Int64 VName
    forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
    sExt64 TExp Int32
    tblock_id

                Names -> Stms GPUMem -> InKernelGen () -> InKernelGen ()
    forall rep r op.
    Names -> Stms rep -> ImpM rep r op () -> ImpM rep r op ()
    compileStms Names
    forall a. Monoid a => a
    mempty (KernelBody GPUMem -> Stms GPUMem
    forall rep. KernelBody rep -> Stms rep
    kernelBodyStms KernelBody GPUMem
    kbody) (InKernelGen () -> InKernelGen ())
    -> InKernelGen () -> InKernelGen ()
    forall a b. (a -> b) -> a -> b
    $
                  (PatElem LetDecMem -> KernelResult -> InKernelGen ())
    -> [PatElem LetDecMem] -> [KernelResult] -> InKernelGen ()
    forall (m :: * -> *) a b c.
    Applicative m =>
    (a -> b -> m c) -> [a] -> [b] -> m ()
    zipWithM_ (SegSpace -> PatElem LetDecMem -> KernelResult -> InKernelGen ()
    compileBlockResult SegSpace
    space) (Pat LetDecMem -> [PatElem LetDecMem]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat LetDecMem
    pat) ([KernelResult] -> InKernelGen ())
    -> [KernelResult] -> InKernelGen ()
    forall a b. (a -> b) -> a -> b
    $
                    KernelBody GPUMem -> [KernelResult]
    forall rep. KernelBody rep -> [KernelResult]
    kernelBodyResult KernelBody GPUMem
    kbody
        SegThreadInBlock {} ->
          String -> CallKernelGen ()
    forall a. HasCallStack => String -> a
    error String
    "compileSegMap: SegThreadInBlock"
      emit $ Imp.DebugPrint "" Nothing