    {-# LANGUAGE TypeFamilies #-}

    -- | Compile a 'GPUMem' program to imperative code with kernels.
    -- This is mostly (but not entirely) the same process no matter if we
    -- are targeting OpenCL or CUDA.  The important distinctions (the host
    -- level code) are introduced later.
    module Futhark.CodeGen.ImpGen.GPU
      ( compileProgOpenCL,
        compileProgCUDA,
        compileProgHIP,
        Warnings,
      )
    where

    import Control.Monad
    import Data.List qualified as L
    import Data.Map qualified as M
    import Data.Maybe
    import Futhark.CodeGen.ImpCode.GPU qualified as Imp
    import Futhark.CodeGen.ImpGen hiding (compileProg)
    import Futhark.CodeGen.ImpGen qualified
    import Futhark.CodeGen.ImpGen.GPU.Base
    import Futhark.CodeGen.ImpGen.GPU.SegHist
    import Futhark.CodeGen.ImpGen.GPU.SegMap
    import Futhark.CodeGen.ImpGen.GPU.SegRed
    import Futhark.CodeGen.ImpGen.GPU.SegScan
    import Futhark.Error
    import Futhark.IR.GPUMem
    import Futhark.MonadFreshNames
    import Futhark.Util.IntegralExp (divUp, nextMul)
    import Prelude hiding (quot, rem)

    callKernelOperations :: Operations GPUMem HostEnv Imp.HostOp
    callKernelOperations :: Operations GPUMem HostEnv HostOp
    callKernelOperations =
      Operations
        { opsExpCompiler :: ExpCompiler GPUMem HostEnv HostOp
    opsExpCompiler = ExpCompiler GPUMem HostEnv HostOp
    expCompiler,
          opsCopyCompiler :: CopyCompiler GPUMem HostEnv HostOp
    opsCopyCompiler = CopyCompiler GPUMem HostEnv HostOp
    forall rep r op. CopyCompiler rep r op
    lmadCopy,
          opsOpCompiler :: OpCompiler GPUMem HostEnv HostOp
    opsOpCompiler = OpCompiler GPUMem HostEnv HostOp
    Pat LetDecMem -> Op GPUMem -> CallKernelGen ()
    opCompiler,
          opsStmsCompiler :: StmsCompiler GPUMem HostEnv HostOp
    opsStmsCompiler = StmsCompiler GPUMem HostEnv HostOp
    forall rep (inner :: * -> *) op r.
    (Mem rep inner, FreeIn op) =>
    Names -> Stms rep -> ImpM rep r op () -> ImpM rep r op ()
    defCompileStms,
          opsAllocCompilers :: Map Space (AllocCompiler GPUMem HostEnv HostOp)
    opsAllocCompilers = Map Space (AllocCompiler GPUMem HostEnv HostOp)
    forall a. Monoid a => a
    mempty
        }

    openclAtomics, cudaAtomics :: AtomicBinOp
    (AtomicBinOp
    openclAtomics, AtomicBinOp
    cudaAtomics) = ((BinOp
     -> [(BinOp,
          VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp)]
     -> Maybe
          (VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp))
    -> [(BinOp,
         VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp)]
    -> AtomicBinOp
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip BinOp
    -> [(BinOp,
         VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp)]
    -> Maybe
         (VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp)
    forall a b. Eq a => a -> [(a, b)] -> Maybe b
    lookup [(BinOp,
      VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp)]
    opencl, (BinOp
     -> [(BinOp,
          VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp)]
     -> Maybe
          (VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp))
    -> [(BinOp,
         VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp)]
    -> AtomicBinOp
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip BinOp
    -> [(BinOp,
         VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp)]
    -> Maybe
         (VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp)
    forall a b. Eq a => a -> [(a, b)] -> Maybe b
    lookup [(BinOp,
      VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp)]
    cuda)
      where
        opencl64 :: [(BinOp,
      VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp)]
    opencl64 =
          [ (IntType -> Overflow -> BinOp
    Add IntType
    Int64 Overflow
    OverflowUndef, IntType
    -> VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp
    Imp.AtomicAdd IntType
    Int64),
            (FloatType -> BinOp
    FAdd FloatType
    Float64, FloatType
    -> VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp
    Imp.AtomicFAdd FloatType
    Float64),
            (IntType -> BinOp
    SMax IntType
    Int64, IntType
    -> VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp
    Imp.AtomicSMax IntType
    Int64),
            (IntType -> BinOp
    SMin IntType
    Int64, IntType
    -> VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp
    Imp.AtomicSMin IntType
    Int64),
            (IntType -> BinOp
    UMax IntType
    Int64, IntType
    -> VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp
    Imp.AtomicUMax IntType
    Int64),
            (IntType -> BinOp
    UMin IntType
    Int64, IntType
    -> VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp
    Imp.AtomicUMin IntType
    Int64),
            (IntType -> BinOp
    And IntType
    Int64, IntType
    -> VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp
    Imp.AtomicAnd IntType
    Int64),
            (IntType -> BinOp
    Or IntType
    Int64, IntType
    -> VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp
    Imp.AtomicOr IntType
    Int64),
            (IntType -> BinOp
    Xor IntType
    Int64, IntType
    -> VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp
    Imp.AtomicXor IntType
    Int64)
          ]
        opencl32 :: [(BinOp,
      VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp)]
    opencl32 =
          [ (IntType -> Overflow -> BinOp
    Add IntType
    Int32 Overflow
    OverflowUndef, IntType
    -> VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp
    Imp.AtomicAdd IntType
    Int32),
            (FloatType -> BinOp
    FAdd FloatType
    Float32, FloatType
    -> VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp
    Imp.AtomicFAdd FloatType
    Float32),
            (IntType -> BinOp
    SMax IntType
    Int32, IntType
    -> VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp
    Imp.AtomicSMax IntType
    Int32),
            (IntType -> BinOp
    SMin IntType
    Int32, IntType
    -> VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp
    Imp.AtomicSMin IntType
    Int32),
            (IntType -> BinOp
    UMax IntType
    Int32, IntType
    -> VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp
    Imp.AtomicUMax IntType
    Int32),
            (IntType -> BinOp
    UMin IntType
    Int32, IntType
    -> VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp
    Imp.AtomicUMin IntType
    Int32),
            (IntType -> BinOp
    And IntType
    Int32, IntType
    -> VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp
    Imp.AtomicAnd IntType
    Int32),
            (IntType -> BinOp
    Or IntType
    Int32, IntType
    -> VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp
    Imp.AtomicOr IntType
    Int32),
            (IntType -> BinOp
    Xor IntType
    Int32, IntType
    -> VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp
    Imp.AtomicXor IntType
    Int32)
          ]
        opencl :: [(BinOp,
      VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp)]
    opencl = [(BinOp,
      VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp)]
    opencl32 [(BinOp,
      VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp)]
    -> [(BinOp,
         VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp)]
    -> [(BinOp,
         VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp)]
    forall a. [a] -> [a] -> [a]
    ++ [(BinOp,
      VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp)]
    opencl64
        cuda :: [(BinOp,
      VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp)]
    cuda = [(BinOp,
      VName -> VName -> Count Elements (TExp Int64) -> Exp -> AtomicOp)]
    opencl

    compileProg ::
      (MonadFreshNames m) =>
      HostEnv ->
      Prog GPUMem ->
      m (Warnings, Imp.Program)
    compileProg :: forall (m :: * -> *).
    MonadFreshNames m =>
    HostEnv -> Prog GPUMem -> m (Warnings, Program)
    compileProg HostEnv
    env =
      HostEnv
    -> Operations GPUMem HostEnv HostOp
    -> Space
    -> Prog GPUMem
    -> m (Warnings, Program)
    forall rep (inner :: * -> *) op (m :: * -> *) r.
    (Mem rep inner, FreeIn op, MonadFreshNames m) =>
    r
    -> Operations rep r op
    -> Space
    -> Prog rep
    -> m (Warnings, Definitions op)
    Futhark.CodeGen.ImpGen.compileProg HostEnv
    env Operations GPUMem HostEnv HostOp
    callKernelOperations Space
    device_space
      where
        device_space :: Space
    device_space = SpaceId -> Space
    Imp.Space SpaceId
    "device"

    -- | Compile a 'GPUMem' program to low-level parallel code, with
    -- either CUDA or OpenCL characteristics.
    compileProgOpenCL,
      compileProgCUDA,
      compileProgHIP ::
        (MonadFreshNames m) => Prog GPUMem -> m (Warnings, Imp.Program)
    compileProgOpenCL :: forall (m :: * -> *).
    MonadFreshNames m =>
    Prog GPUMem -> m (Warnings, Program)
    compileProgOpenCL = HostEnv -> Prog GPUMem -> m (Warnings, Program)
    forall (m :: * -> *).
    MonadFreshNames m =>
    HostEnv -> Prog GPUMem -> m (Warnings, Program)
    compileProg (HostEnv -> Prog GPUMem -> m (Warnings, Program))
    -> HostEnv -> Prog GPUMem -> m (Warnings, Program)
    forall a b. (a -> b) -> a -> b
    $ AtomicBinOp -> Target -> Map VName Locks -> HostEnv
    HostEnv AtomicBinOp
    openclAtomics Target
    OpenCL Map VName Locks
    forall a. Monoid a => a
    mempty
    compileProgCUDA :: forall (m :: * -> *).
    MonadFreshNames m =>
    Prog GPUMem -> m (Warnings, Program)
    compileProgCUDA = HostEnv -> Prog GPUMem -> m (Warnings, Program)
    forall (m :: * -> *).
    MonadFreshNames m =>
    HostEnv -> Prog GPUMem -> m (Warnings, Program)
    compileProg (HostEnv -> Prog GPUMem -> m (Warnings, Program))
    -> HostEnv -> Prog GPUMem -> m (Warnings, Program)
    forall a b. (a -> b) -> a -> b
    $ AtomicBinOp -> Target -> Map VName Locks -> HostEnv
    HostEnv AtomicBinOp
    cudaAtomics Target
    CUDA Map VName Locks
    forall a. Monoid a => a
    mempty
    compileProgHIP :: forall (m :: * -> *).
    MonadFreshNames m =>
    Prog GPUMem -> m (Warnings, Program)
    compileProgHIP = HostEnv -> Prog GPUMem -> m (Warnings, Program)
    forall (m :: * -> *).
    MonadFreshNames m =>
    HostEnv -> Prog GPUMem -> m (Warnings, Program)
    compileProg (HostEnv -> Prog GPUMem -> m (Warnings, Program))
    -> HostEnv -> Prog GPUMem -> m (Warnings, Program)
    forall a b. (a -> b) -> a -> b
    $ AtomicBinOp -> Target -> Map VName Locks -> HostEnv
    HostEnv AtomicBinOp
    cudaAtomics Target
    HIP Map VName Locks
    forall a. Monoid a => a
    mempty

    opCompiler ::
      Pat LetDecMem ->
      Op GPUMem ->
      CallKernelGen ()
    opCompiler :: Pat LetDecMem -> Op GPUMem -> CallKernelGen ()
    opCompiler Pat LetDecMem
    dest (Alloc SubExp
    e Space
    space) =
      Pat (LetDec GPUMem) -> SubExp -> Space -> CallKernelGen ()
    forall rep (inner :: * -> *) r op.
    Mem rep inner =>
    Pat (LetDec rep) -> SubExp -> Space -> ImpM rep r op ()
    compileAlloc Pat (LetDec GPUMem)
    Pat LetDecMem
    dest SubExp
    e Space
    space
    opCompiler (Pat [PatElem LetDecMem
    pe]) (Inner (SizeOp (GetSize Name
    key SizeClass
    size_class))) = do
      Maybe Name
    fname <- ImpM GPUMem HostEnv HostOp (Maybe Name)
    forall rep r op. ImpM rep r op (Maybe Name)
    askFunction
      HostOp -> CallKernelGen ()
    forall op rep r. op -> ImpM rep r op ()
    sOp (HostOp -> CallKernelGen ()) -> HostOp -> CallKernelGen ()
    forall a b. (a -> b) -> a -> b
    $
        VName -> Name -> SizeClass -> HostOp
    Imp.GetSize (PatElem LetDecMem -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem LetDecMem
    pe) (Maybe Name -> Name -> Name
    keyWithEntryPoint Maybe Name
    fname Name
    key) (SizeClass -> HostOp) -> SizeClass -> HostOp
    forall a b. (a -> b) -> a -> b
    $
          Maybe Name -> SizeClass -> SizeClass
    sizeClassWithEntryPoint Maybe Name
    fname SizeClass
    size_class
    opCompiler (Pat [PatElem LetDecMem
    pe]) (Inner (SizeOp (CmpSizeLe Name
    key SizeClass
    size_class SubExp
    x))) = do
      Maybe Name
    fname <- ImpM GPUMem HostEnv HostOp (Maybe Name)
    forall rep r op. ImpM rep r op (Maybe Name)
    askFunction
      let size_class' :: SizeClass
    size_class' = Maybe Name -> SizeClass -> SizeClass
    sizeClassWithEntryPoint Maybe Name
    fname SizeClass
    size_class
      HostOp -> CallKernelGen ()
    forall op rep r. op -> ImpM rep r op ()
    sOp (HostOp -> CallKernelGen ())
    -> (Exp -> HostOp) -> Exp -> CallKernelGen ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VName -> Name -> SizeClass -> Exp -> HostOp
    Imp.CmpSizeLe (PatElem LetDecMem -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem LetDecMem
    pe) (Maybe Name -> Name -> Name
    keyWithEntryPoint Maybe Name
    fname Name
    key) SizeClass
    size_class'
        (Exp -> CallKernelGen ())
    -> ImpM GPUMem HostEnv HostOp Exp -> CallKernelGen ()
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< SubExp -> ImpM GPUMem HostEnv HostOp Exp
    forall a rep r op. ToExp a => a -> ImpM rep r op Exp
    forall rep r op. SubExp -> ImpM rep r op Exp
    toExp SubExp
    x
    opCompiler (Pat [PatElem LetDecMem
    pe]) (Inner (SizeOp (GetSizeMax SizeClass
    size_class))) =
      HostOp -> CallKernelGen ()
    forall op rep r. op -> ImpM rep r op ()
    sOp (HostOp -> CallKernelGen ()) -> HostOp -> CallKernelGen ()
    forall a b. (a -> b) -> a -> b
    $ VName -> SizeClass -> HostOp
    Imp.GetSizeMax (PatElem LetDecMem -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem LetDecMem
    pe) SizeClass
    size_class
    opCompiler (Pat [PatElem LetDecMem
    pe]) (Inner (SizeOp (CalcNumBlocks SubExp
    w64 Name
    max_num_tblocks_key SubExp
    tblock_size))) = do
      Maybe Name
    fname <- ImpM GPUMem HostEnv HostOp (Maybe Name)
    forall rep r op. ImpM rep r op (Maybe Name)
    askFunction
      TV Int64
    max_num_tblocks :: TV Int64 <- SpaceId -> ImpM GPUMem HostEnv HostOp (TV Int64)
    forall {k} (t :: k) rep r op.
    MkTV t =>
    SpaceId -> ImpM rep r op (TV t)
    dPrim SpaceId
    "max_num_tblocks"
      HostOp -> CallKernelGen ()
    forall op rep r. op -> ImpM rep r op ()
    sOp (HostOp -> CallKernelGen ()) -> HostOp -> CallKernelGen ()
    forall a b. (a -> b) -> a -> b
    $
        VName -> Name -> SizeClass -> HostOp
    Imp.GetSize (TV Int64 -> VName
    forall {k} (t :: k). TV t -> VName
    tvVar TV Int64
    max_num_tblocks) (Maybe Name -> Name -> Name
    keyWithEntryPoint Maybe Name
    fname Name
    max_num_tblocks_key) (SizeClass -> HostOp) -> SizeClass -> HostOp
    forall a b. (a -> b) -> a -> b
    $
          Maybe Name -> SizeClass -> SizeClass
    sizeClassWithEntryPoint Maybe Name
    fname SizeClass
    SizeGrid

      -- If 'w' is small, we launch fewer blocks than we normally would.
      -- We don't want any idle blocks.
      --
      -- The calculations are done with 64-bit integers to avoid overflow
      -- issues.
      let num_tblocks_maybe_zero :: TExp Int64
    num_tblocks_maybe_zero =
            TExp Int64 -> TExp Int64 -> TExp Int64
    forall v. TPrimExp Int64 v -> TPrimExp Int64 v -> TPrimExp Int64 v
    sMin64 (SubExp -> TExp Int64
    pe64 SubExp
    w64 TExp Int64 -> TExp Int64 -> TExp Int64
    forall e. IntegralExp e => e -> e -> e
    `divUp` SubExp -> TExp Int64
    pe64 SubExp
    tblock_size) (TExp Int64 -> TExp Int64) -> TExp Int64 -> TExp Int64
    forall a b. (a -> b) -> a -> b
    $
              TExp Int64 -> TExp Int64
    forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
    sExt64 (TV Int64 -> TExp Int64
    forall {k} (t :: k). TV t -> TExp t
    tvExp TV Int64
    max_num_tblocks)
      -- We also don't want zero blocks.
      let num_tblocks :: TExp Int64
    num_tblocks = TExp Int64 -> TExp Int64 -> TExp Int64
    forall v. TPrimExp Int64 v -> TPrimExp Int64 v -> TPrimExp Int64 v
    sMax64 TExp Int64
    1 TExp Int64
    num_tblocks_maybe_zero
      VName -> TV Int32
    forall {k} (t :: k). MkTV t => VName -> TV t
    mkTV (PatElem LetDecMem -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem LetDecMem
    pe) TV Int32 -> TExp Int32 -> CallKernelGen ()
    forall {k} (t :: k) rep r op. TV t -> TExp t -> ImpM rep r op ()
    <-- TExp Int64 -> TExp Int32
    forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int32 v
    sExt32 TExp Int64
    num_tblocks
    opCompiler Pat LetDecMem
    dest (Inner (SegOp SegOp SegLevel GPUMem
    op)) =
      Pat LetDecMem -> SegOp SegLevel GPUMem -> CallKernelGen ()
    segOpCompiler Pat LetDecMem
    dest SegOp SegLevel GPUMem
    op
    opCompiler (Pat [PatElem LetDecMem]
    pes) (Inner (GPUBody [TypeBase Shape NoUniqueness]
    _ (Body BodyDec GPUMem
    _ Stms GPUMem
    stms Result
    res))) = do
      VName
    tid <- SpaceId -> ImpM GPUMem HostEnv HostOp VName
    forall (m :: * -> *). MonadFreshNames m => SpaceId -> m VName
    newVName SpaceId
    "tid"
      let one :: Count u SubExp
    one = SubExp -> Count u SubExp
    forall {k} (u :: k) e. e -> Count u e
    Count (IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    1)
      SpaceId
    -> VName -> KernelAttrs -> InKernelGen () -> CallKernelGen ()
    sKernelThread SpaceId
    "gpuseq" VName
    tid (Count NumBlocks SubExp -> Count BlockSize SubExp -> KernelAttrs
    defKernelAttrs Count NumBlocks SubExp
    forall {k} {u :: k}. Count u SubExp
    one Count BlockSize SubExp
    forall {k} {u :: k}. Count u SubExp
    one) (InKernelGen () -> CallKernelGen ())
    -> InKernelGen () -> CallKernelGen ()
    forall a b. (a -> b) -> a -> b
    $
        Names -> Stms GPUMem -> InKernelGen () -> InKernelGen ()
    forall rep r op.
    Names -> Stms rep -> ImpM rep r op () -> ImpM rep r op ()
    compileStms (Result -> Names
    forall a. FreeIn a => a -> Names
    freeIn Result
    res) Stms GPUMem
    stms (InKernelGen () -> InKernelGen ())
    -> InKernelGen () -> InKernelGen ()
    forall a b. (a -> b) -> a -> b
    $
          [(PatElem LetDecMem, SubExpRes)]
    -> ((PatElem LetDecMem, SubExpRes) -> InKernelGen ())
    -> InKernelGen ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([PatElem LetDecMem] -> Result -> [(PatElem LetDecMem, SubExpRes)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [PatElem LetDecMem]
    pes Result
    res) (((PatElem LetDecMem, SubExpRes) -> InKernelGen ())
     -> InKernelGen ())
    -> ((PatElem LetDecMem, SubExpRes) -> InKernelGen ())
    -> InKernelGen ()
    forall a b. (a -> b) -> a -> b
    $ \(PatElem LetDecMem
    pe, SubExpRes Certs
    _ SubExp
    se) ->
            VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> InKernelGen ()
    forall rep r op.
    VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
    copyDWIMFix (PatElem LetDecMem -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem LetDecMem
    pe) [TExp Int64
    0] SubExp
    se []
    opCompiler Pat LetDecMem
    pat Op GPUMem
    e =
      SpaceId -> CallKernelGen ()
    forall a. SpaceId -> a
    compilerBugS (SpaceId -> CallKernelGen ()) -> SpaceId -> CallKernelGen ()
    forall a b. (a -> b) -> a -> b
    $
        SpaceId
    "opCompiler: Invalid pattern\n  "
          SpaceId -> SpaceId -> SpaceId
    forall a. [a] -> [a] -> [a]
    ++ Pat LetDecMem -> SpaceId
    forall a. Pretty a => a -> SpaceId
    prettyString Pat LetDecMem
    pat
          SpaceId -> SpaceId -> SpaceId
    forall a. [a] -> [a] -> [a]
    ++ SpaceId
    "\nfor expression\n  "
          SpaceId -> SpaceId -> SpaceId
    forall a. [a] -> [a] -> [a]
    ++ MemOp (HostOp NoOp) GPUMem -> SpaceId
    forall a. Pretty a => a -> SpaceId
    prettyString Op GPUMem
    MemOp (HostOp NoOp) GPUMem
    e

    sizeClassWithEntryPoint :: Maybe Name -> Imp.SizeClass -> Imp.SizeClass
    sizeClassWithEntryPoint :: Maybe Name -> SizeClass -> SizeClass
    sizeClassWithEntryPoint Maybe Name
    fname (Imp.SizeThreshold KernelPath
    path Maybe Int64
    def) =
      KernelPath -> Maybe Int64 -> SizeClass
    Imp.SizeThreshold (((Name, Bool) -> (Name, Bool)) -> KernelPath -> KernelPath
    forall a b. (a -> b) -> [a] -> [b]
    map (Name, Bool) -> (Name, Bool)
    f KernelPath
    path) Maybe Int64
    def
      where
        f :: (Name, Bool) -> (Name, Bool)
    f (Name
    name, Bool
    x) = (Maybe Name -> Name -> Name
    keyWithEntryPoint Maybe Name
    fname Name
    name, Bool
    x)
    sizeClassWithEntryPoint Maybe Name
    _ SizeClass
    size_class = SizeClass
    size_class

    segOpCompiler ::
      Pat LetDecMem ->
      SegOp SegLevel GPUMem ->
      CallKernelGen ()
    segOpCompiler :: Pat LetDecMem -> SegOp SegLevel GPUMem -> CallKernelGen ()
    segOpCompiler Pat LetDecMem
    pat (SegMap SegLevel
    lvl SegSpace
    space [TypeBase Shape NoUniqueness]
    _ KernelBody GPUMem
    kbody) =
      Pat LetDecMem
    -> SegLevel -> SegSpace -> KernelBody GPUMem -> CallKernelGen ()
    compileSegMap Pat LetDecMem
    pat SegLevel
    lvl SegSpace
    space KernelBody GPUMem
    kbody
    segOpCompiler Pat LetDecMem
    pat (SegRed lvl :: SegLevel
    lvl@(SegThread SegVirt
    _ Maybe KernelGrid
    _) SegSpace
    space [SegBinOp GPUMem]
    reds [TypeBase Shape NoUniqueness]
    _ KernelBody GPUMem
    kbody) =
      Pat LetDecMem
    -> SegLevel
    -> SegSpace
    -> [SegBinOp GPUMem]
    -> KernelBody GPUMem
    -> CallKernelGen ()
    compileSegRed Pat LetDecMem
    pat SegLevel
    lvl SegSpace
    space [SegBinOp GPUMem]
    reds KernelBody GPUMem
    kbody
    segOpCompiler Pat LetDecMem
    pat (SegScan lvl :: SegLevel
    lvl@(SegThread SegVirt
    _ Maybe KernelGrid
    _) SegSpace
    space [SegBinOp GPUMem]
    scans [TypeBase Shape NoUniqueness]
    _ KernelBody GPUMem
    kbody) =
      Pat LetDecMem
    -> SegLevel
    -> SegSpace
    -> [SegBinOp GPUMem]
    -> KernelBody GPUMem
    -> CallKernelGen ()
    compileSegScan Pat LetDecMem
    pat SegLevel
    lvl SegSpace
    space [SegBinOp GPUMem]
    scans KernelBody GPUMem
    kbody
    segOpCompiler Pat LetDecMem
    pat (SegHist lvl :: SegLevel
    lvl@(SegThread SegVirt
    _ Maybe KernelGrid
    _) SegSpace
    space [HistOp GPUMem]
    ops [TypeBase Shape NoUniqueness]
    _ KernelBody GPUMem
    kbody) =
      Pat LetDecMem
    -> SegLevel
    -> SegSpace
    -> [HistOp GPUMem]
    -> KernelBody GPUMem
    -> CallKernelGen ()
    compileSegHist Pat LetDecMem
    pat SegLevel
    lvl SegSpace
    space [HistOp GPUMem]
    ops KernelBody GPUMem
    kbody
    segOpCompiler Pat LetDecMem
    pat SegOp SegLevel GPUMem
    segop =
      SpaceId -> CallKernelGen ()
    forall a. SpaceId -> a
    compilerBugS (SpaceId -> CallKernelGen ()) -> SpaceId -> CallKernelGen ()
    forall a b. (a -> b) -> a -> b
    $ SpaceId
    "segOpCompiler: unexpected " SpaceId -> SpaceId -> SpaceId
    forall a. [a] -> [a] -> [a]
    ++ SegLevel -> SpaceId
    forall a. Pretty a => a -> SpaceId
    prettyString (SegOp SegLevel GPUMem -> SegLevel
    forall lvl rep. SegOp lvl rep -> lvl
    segLevel SegOp SegLevel GPUMem
    segop) SpaceId -> SpaceId -> SpaceId
    forall a. [a] -> [a] -> [a]
    ++ SpaceId
    " for rhs of pattern " SpaceId -> SpaceId -> SpaceId
    forall a. [a] -> [a] -> [a]
    ++ Pat LetDecMem -> SpaceId
    forall a. Pretty a => a -> SpaceId
    prettyString Pat LetDecMem
    pat

    -- Create boolean expression that checks whether all kernels in the
    -- enclosed code do not use more shared memory than we have available.
    -- We look at *all* the kernels here, even those that might be
    -- otherwise protected by their own multi-versioning branches deeper
    -- down.  Currently the compiler will not generate multi-versioning
    -- that makes this a problem, but it might in the future.
    checkSharedMemoryReqs :: (VName -> Bool) -> Imp.HostCode -> CallKernelGen (Maybe (Imp.TExp Bool))
    checkSharedMemoryReqs :: (VName -> Bool) -> HostCode -> CallKernelGen (Maybe (TExp Bool))
    checkSharedMemoryReqs VName -> Bool
    in_scope HostCode
    code = do
      let alloc_sizes :: [Count Bytes (TExp Int64)]
    alloc_sizes = (Kernel -> Count Bytes (TExp Int64))
    -> [Kernel] -> [Count Bytes (TExp Int64)]
    forall a b. (a -> b) -> [a] -> [b]
    map ([Count Bytes (TExp Int64)] -> Count Bytes (TExp Int64)
    forall a. Num a => [a] -> a
    forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
    sum ([Count Bytes (TExp Int64)] -> Count Bytes (TExp Int64))
    -> (Kernel -> [Count Bytes (TExp Int64)])
    -> Kernel
    -> Count Bytes (TExp Int64)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Count Bytes (TExp Int64) -> Count Bytes (TExp Int64))
    -> [Count Bytes (TExp Int64)] -> [Count Bytes (TExp Int64)]
    forall a b. (a -> b) -> [a] -> [b]
    map Count Bytes (TExp Int64) -> Count Bytes (TExp Int64)
    forall {e}. IntegralExp e => e -> e
    alignedSize ([Count Bytes (TExp Int64)] -> [Count Bytes (TExp Int64)])
    -> (Kernel -> [Count Bytes (TExp Int64)])
    -> Kernel
    -> [Count Bytes (TExp Int64)]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Code KernelOp -> [Count Bytes (TExp Int64)]
    localAllocSizes (Code KernelOp -> [Count Bytes (TExp Int64)])
    -> (Kernel -> Code KernelOp)
    -> Kernel
    -> [Count Bytes (TExp Int64)]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Kernel -> Code KernelOp
    Imp.kernelBody) ([Kernel] -> [Count Bytes (TExp Int64)])
    -> [Kernel] -> [Count Bytes (TExp Int64)]
    forall a b. (a -> b) -> a -> b
    $ HostCode -> [Kernel]
    getGPU HostCode
    code

      -- If any of the sizes involve a variable that is not known at this
      -- point, then we cannot check the requirements.
      if Bool -> Bool
    not (Bool -> Bool) -> Bool -> Bool
    forall a b. (a -> b) -> a -> b
    $ (VName -> Bool) -> [VName] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all VName -> Bool
    in_scope ([VName] -> Bool) -> [VName] -> Bool
    forall a b. (a -> b) -> a -> b
    $ Names -> [VName]
    namesToList (Names -> [VName]) -> Names -> [VName]
    forall a b. (a -> b) -> a -> b
    $ [Count Bytes (TExp Int64)] -> Names
    forall a. FreeIn a => a -> Names
    freeIn [Count Bytes (TExp Int64)]
    alloc_sizes
        then Maybe (TExp Bool) -> CallKernelGen (Maybe (TExp Bool))
    forall a. a -> ImpM GPUMem HostEnv HostOp a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Maybe (TExp Bool)
    forall a. Maybe a
    Nothing
        else do
          TV Int64
    shared_memory_capacity :: TV Int64 <- SpaceId -> ImpM GPUMem HostEnv HostOp (TV Int64)
    forall {k} (t :: k) rep r op.
    MkTV t =>
    SpaceId -> ImpM rep r op (TV t)
    dPrim SpaceId
    "shared_memory_capacity"
          HostOp -> CallKernelGen ()
    forall op rep r. op -> ImpM rep r op ()
    sOp (HostOp -> CallKernelGen ()) -> HostOp -> CallKernelGen ()
    forall a b. (a -> b) -> a -> b
    $ VName -> SizeClass -> HostOp
    Imp.GetSizeMax (TV Int64 -> VName
    forall {k} (t :: k). TV t -> VName
    tvVar TV Int64
    shared_memory_capacity) SizeClass
    SizeSharedMemory

          let shared_memory_capacity_64 :: TExp Int64
    shared_memory_capacity_64 =
                TExp Int64 -> TExp Int64
    forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
    sExt64 (TExp Int64 -> TExp Int64) -> TExp Int64 -> TExp Int64
    forall a b. (a -> b) -> a -> b
    $ TV Int64 -> TExp Int64
    forall {k} (t :: k). TV t -> TExp t
    tvExp TV Int64
    shared_memory_capacity
              fits :: Count Bytes (TExp Int64) -> TExp Bool
    fits Count Bytes (TExp Int64)
    size =
                Count Bytes (TExp Int64) -> TExp Int64
    forall {k} (u :: k) e. Count u e -> e
    unCount Count Bytes (TExp Int64)
    size TExp Int64 -> TExp Int64 -> TExp Bool
    forall {k} v (t :: k).
    Eq v =>
    TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
    .<=. TExp Int64
    shared_memory_capacity_64
          Maybe (TExp Bool) -> CallKernelGen (Maybe (TExp Bool))
    forall a. a -> ImpM GPUMem HostEnv HostOp a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Maybe (TExp Bool) -> CallKernelGen (Maybe (TExp Bool)))
    -> Maybe (TExp Bool) -> CallKernelGen (Maybe (TExp Bool))
    forall a b. (a -> b) -> a -> b
    $ TExp Bool -> Maybe (TExp Bool)
    forall a. a -> Maybe a
    Just (TExp Bool -> Maybe (TExp Bool)) -> TExp Bool -> Maybe (TExp Bool)
    forall a b. (a -> b) -> a -> b
    $ (TExp Bool -> TExp Bool -> TExp Bool)
    -> TExp Bool -> [TExp Bool] -> TExp Bool
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    L.foldl' TExp Bool -> TExp Bool -> TExp Bool
    forall v.
    Eq v =>
    TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v
    (.&&.) TExp Bool
    forall v. TPrimExp Bool v
    true ((Count Bytes (TExp Int64) -> TExp Bool)
    -> [Count Bytes (TExp Int64)] -> [TExp Bool]
    forall a b. (a -> b) -> [a] -> [b]
    map Count Bytes (TExp Int64) -> TExp Bool
    fits [Count Bytes (TExp Int64)]
    alloc_sizes)
      where
        getGPU :: HostCode -> [Kernel]
    getGPU = (HostOp -> [Kernel]) -> HostCode -> [Kernel]
    forall m a. Monoid m => (a -> m) -> Code a -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap HostOp -> [Kernel]
    getKernel
        getKernel :: HostOp -> [Kernel]
    getKernel (Imp.CallKernel Kernel
    k) | Kernel -> Bool
    Imp.kernelCheckSharedMemory Kernel
    k = [Kernel
    k]
        getKernel HostOp
    _ = []

        localAllocSizes :: Code KernelOp -> [Count Bytes (TExp Int64)]
    localAllocSizes = (KernelOp -> [Count Bytes (TExp Int64)])
    -> Code KernelOp -> [Count Bytes (TExp Int64)]
    forall m a. Monoid m => (a -> m) -> Code a -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap KernelOp -> [Count Bytes (TExp Int64)]
    localAllocSize
        localAllocSize :: KernelOp -> [Count Bytes (TExp Int64)]
    localAllocSize (Imp.SharedAlloc VName
    _ Count Bytes (TExp Int64)
    size) = [Count Bytes (TExp Int64)
    size]
        localAllocSize KernelOp
    _ = []

        -- These allocations will actually be padded to an 8-byte aligned
        -- size, so we should take that into account when checking whether
        -- they fit.
        alignedSize :: e -> e
    alignedSize e
    x = e -> e -> e
    forall e. IntegralExp e => e -> e -> e
    nextMul e
    x e
    8

    withAcc ::
      Pat LetDecMem ->
      [(Shape, [VName], Maybe (Lambda GPUMem, [SubExp]))] ->
      Lambda GPUMem ->
      CallKernelGen ()
    withAcc :: Pat LetDecMem
    -> [(Shape, [VName], Maybe (Lambda GPUMem, [SubExp]))]
    -> Lambda GPUMem
    -> CallKernelGen ()
    withAcc Pat LetDecMem
    pat [(Shape, [VName], Maybe (Lambda GPUMem, [SubExp]))]
    inputs Lambda GPUMem
    lam = do
      AtomicBinOp
    atomics <- HostEnv -> AtomicBinOp
    hostAtomics (HostEnv -> AtomicBinOp)
    -> ImpM GPUMem HostEnv HostOp HostEnv
    -> ImpM GPUMem HostEnv HostOp AtomicBinOp
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> ImpM GPUMem HostEnv HostOp HostEnv
    forall rep r op. ImpM rep r op r
    askEnv
      AtomicBinOp
    -> [(VName, (Shape, [VName], Maybe (Lambda GPUMem, [SubExp])))]
    -> CallKernelGen ()
    locksForInputs AtomicBinOp
    atomics ([(VName, (Shape, [VName], Maybe (Lambda GPUMem, [SubExp])))]
     -> CallKernelGen ())
    -> [(VName, (Shape, [VName], Maybe (Lambda GPUMem, [SubExp])))]
    -> CallKernelGen ()
    forall a b. (a -> b) -> a -> b
    $ [VName]
    -> [(Shape, [VName], Maybe (Lambda GPUMem, [SubExp]))]
    -> [(VName, (Shape, [VName], Maybe (Lambda GPUMem, [SubExp])))]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [VName]
    accs [(Shape, [VName], Maybe (Lambda GPUMem, [SubExp]))]
    inputs
      where
        accs :: [VName]
    accs = (Param LetDecMem -> VName) -> [Param LetDecMem] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Param LetDecMem -> VName
    forall dec. Param dec -> VName
    paramName ([Param LetDecMem] -> [VName]) -> [Param LetDecMem] -> [VName]
    forall a b. (a -> b) -> a -> b
    $ Lambda GPUMem -> [LParam GPUMem]
    forall rep. Lambda rep -> [LParam rep]
    lambdaParams Lambda GPUMem
    lam
        locksForInputs :: AtomicBinOp
    -> [(VName, (Shape, [VName], Maybe (Lambda GPUMem, [SubExp])))]
    -> CallKernelGen ()
    locksForInputs AtomicBinOp
    _ [] =
          ExpCompiler GPUMem HostEnv HostOp
    forall rep (inner :: * -> *) r op.
    Mem rep inner =>
    Pat (LetDec rep) -> Exp rep -> ImpM rep r op ()
    defCompileExp Pat (LetDec GPUMem)
    Pat LetDecMem
    pat (Exp GPUMem -> CallKernelGen ()) -> Exp GPUMem -> CallKernelGen ()
    forall a b. (a -> b) -> a -> b
    $ [(Shape, [VName], Maybe (Lambda GPUMem, [SubExp]))]
    -> Lambda GPUMem -> Exp GPUMem
    forall rep. [WithAccInput rep] -> Lambda rep -> Exp rep
    WithAcc [(Shape, [VName], Maybe (Lambda GPUMem, [SubExp]))]
    inputs Lambda GPUMem
    lam
        locksForInputs AtomicBinOp
    atomics ((VName
    c, (Shape
    _, [VName]
    _, Maybe (Lambda GPUMem, [SubExp])
    op)) : [(VName, (Shape, [VName], Maybe (Lambda GPUMem, [SubExp])))]
    inputs')
          | Just (Lambda GPUMem
    op_lam, [SubExp]
    _) <- Maybe (Lambda GPUMem, [SubExp])
    op,
            AtomicLocking Locking -> DoAtomicUpdate GPUMem KernelEnv
    _ <- AtomicBinOp -> Lambda GPUMem -> AtomicUpdate GPUMem KernelEnv
    atomicUpdateLocking AtomicBinOp
    atomics Lambda GPUMem
    op_lam = do
              let num_locks :: Int
    num_locks = Int
    100151
              VName
    locks_arr <- SpaceId -> Int -> ImpM GPUMem HostEnv HostOp VName
    genZeroes SpaceId
    "withacc_locks" Int
    num_locks
              let locks :: Locks
    locks = VName -> Int -> Locks
    Locks VName
    locks_arr Int
    num_locks
                  extend :: HostEnv -> HostEnv
    extend HostEnv
    env = HostEnv
    env {hostLocks = M.insert c locks $ hostLocks env}
              (HostEnv -> HostEnv) -> CallKernelGen () -> CallKernelGen ()
    forall r rep op a. (r -> r) -> ImpM rep r op a -> ImpM rep r op a
    localEnv HostEnv -> HostEnv
    extend (CallKernelGen () -> CallKernelGen ())
    -> CallKernelGen () -> CallKernelGen ()
    forall a b. (a -> b) -> a -> b
    $ AtomicBinOp
    -> [(VName, (Shape, [VName], Maybe (Lambda GPUMem, [SubExp])))]
    -> CallKernelGen ()
    locksForInputs AtomicBinOp
    atomics [(VName, (Shape, [VName], Maybe (Lambda GPUMem, [SubExp])))]
    inputs'
          | Bool
    otherwise =
              AtomicBinOp
    -> [(VName, (Shape, [VName], Maybe (Lambda GPUMem, [SubExp])))]
    -> CallKernelGen ()
    locksForInputs AtomicBinOp
    atomics [(VName, (Shape, [VName], Maybe (Lambda GPUMem, [SubExp])))]
    inputs'

    expCompiler :: ExpCompiler GPUMem HostEnv Imp.HostOp
    -- We generate a simple kernel for iota and replicate.
    expCompiler :: ExpCompiler GPUMem HostEnv HostOp
    expCompiler (Pat [PatElem (LetDec GPUMem)
    pe]) (BasicOp (Iota SubExp
    n SubExp
    x SubExp
    s IntType
    et)) = do
      Exp
    x' <- SubExp -> ImpM GPUMem HostEnv HostOp Exp
    forall a rep r op. ToExp a => a -> ImpM rep r op Exp
    forall rep r op. SubExp -> ImpM rep r op Exp
    toExp SubExp
    x
      Exp
    s' <- SubExp -> ImpM GPUMem HostEnv HostOp Exp
    forall a rep r op. ToExp a => a -> ImpM rep r op Exp
    forall rep r op. SubExp -> ImpM rep r op Exp
    toExp SubExp
    s
      VName -> TExp Int64 -> Exp -> Exp -> IntType -> CallKernelGen ()
    sIota (PatElem LetDecMem -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem (LetDec GPUMem)
    PatElem LetDecMem
    pe) (SubExp -> TExp Int64
    pe64 SubExp
    n) Exp
    x' Exp
    s' IntType
    et
    expCompiler (Pat [PatElem (LetDec GPUMem)
    pe]) (BasicOp (Replicate Shape
    shape SubExp
    se))
      | Acc {} <- PatElem LetDecMem -> TypeBase Shape NoUniqueness
    forall dec. Typed dec => PatElem dec -> TypeBase Shape NoUniqueness
    patElemType PatElem (LetDec GPUMem)
    PatElem LetDecMem
    pe = () -> CallKernelGen ()
    forall a. a -> ImpM GPUMem HostEnv HostOp a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
      | Bool
    otherwise =
          if Shape -> Int
    forall a. ArrayShape a => a -> Int
    shapeRank Shape
    shape Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    0
            then VName
    -> [DimIndex (TExp Int64)]
    -> SubExp
    -> [DimIndex (TExp Int64)]
    -> CallKernelGen ()
    forall rep r op.
    VName
    -> [DimIndex (TExp Int64)]
    -> SubExp
    -> [DimIndex (TExp Int64)]
    -> ImpM rep r op ()
    copyDWIM (PatElem LetDecMem -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem (LetDec GPUMem)
    PatElem LetDecMem
    pe) [] SubExp
    se []
            else VName -> SubExp -> CallKernelGen ()
    sReplicate (PatElem LetDecMem -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem (LetDec GPUMem)
    PatElem LetDecMem
    pe) SubExp
    se
    -- Allocation in the "shared" space is just a placeholder.
    expCompiler Pat (LetDec GPUMem)
    _ (Op (Alloc SubExp
    _ (Space SpaceId
    "shared"))) =
      () -> CallKernelGen ()
    forall a. a -> ImpM GPUMem HostEnv HostOp a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
    expCompiler Pat (LetDec GPUMem)
    pat (WithAcc [(Shape, [VName], Maybe (Lambda GPUMem, [SubExp]))]
    inputs Lambda GPUMem
    lam) =
      Pat LetDecMem
    -> [(Shape, [VName], Maybe (Lambda GPUMem, [SubExp]))]
    -> Lambda GPUMem
    -> CallKernelGen ()
    withAcc Pat (LetDec GPUMem)
    Pat LetDecMem
    pat [(Shape, [VName], Maybe (Lambda GPUMem, [SubExp]))]
    inputs Lambda GPUMem
    lam
    -- This is a multi-versioning Match created by incremental flattening.
    -- We need to augment the conditional with a check that any local
    -- memory requirements in tbranch are compatible with the hardware.
    -- We do not check anything for defbody, as we assume that it will
    -- always be safe (and what would we do if none of the branches would
    -- work?).
    expCompiler Pat (LetDec GPUMem)
    dest (Match [SubExp]
    cond (Case (Body GPUMem)
    first_case : [Case (Body GPUMem)]
    cases) Body GPUMem
    defbranch sort :: MatchDec (BranchType GPUMem)
    sort@(MatchDec [BranchType GPUMem]
    _ MatchSort
    MatchEquiv)) = do
      Scope SOACS
    scope <- ImpM GPUMem HostEnv HostOp (Scope SOACS)
    forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
    askScope
      HostCode
    tcode <- CallKernelGen () -> ImpM GPUMem HostEnv HostOp HostCode
    forall rep r op. ImpM rep r op () -> ImpM rep r op (Code op)
    collect (CallKernelGen () -> ImpM GPUMem HostEnv HostOp HostCode)
    -> CallKernelGen () -> ImpM GPUMem HostEnv HostOp HostCode
    forall a b. (a -> b) -> a -> b
    $ Pat (LetDec GPUMem) -> Body GPUMem -> CallKernelGen ()
    forall rep r op. Pat (LetDec rep) -> Body rep -> ImpM rep r op ()
    compileBody Pat (LetDec GPUMem)
    dest (Body GPUMem -> CallKernelGen ())
    -> Body GPUMem -> CallKernelGen ()
    forall a b. (a -> b) -> a -> b
    $ Case (Body GPUMem) -> Body GPUMem
    forall body. Case body -> body
    caseBody Case (Body GPUMem)
    first_case
      HostCode
    fcode <- CallKernelGen () -> ImpM GPUMem HostEnv HostOp HostCode
    forall rep r op. ImpM rep r op () -> ImpM rep r op (Code op)
    collect (CallKernelGen () -> ImpM GPUMem HostEnv HostOp HostCode)
    -> CallKernelGen () -> ImpM GPUMem HostEnv HostOp HostCode
    forall a b. (a -> b) -> a -> b
    $ ExpCompiler GPUMem HostEnv HostOp
    expCompiler Pat (LetDec GPUMem)
    dest (Exp GPUMem -> CallKernelGen ()) -> Exp GPUMem -> CallKernelGen ()
    forall a b. (a -> b) -> a -> b
    $ [SubExp]
    -> [Case (Body GPUMem)]
    -> Body GPUMem
    -> MatchDec (BranchType GPUMem)
    -> Exp GPUMem
    forall rep.
    [SubExp]
    -> [Case (Body rep)]
    -> Body rep
    -> MatchDec (BranchType rep)
    -> Exp rep
    Match [SubExp]
    cond [Case (Body GPUMem)]
    cases Body GPUMem
    defbranch MatchDec (BranchType GPUMem)
    sort
      Maybe (TExp Bool)
    check <- (VName -> Bool) -> HostCode -> CallKernelGen (Maybe (TExp Bool))
    checkSharedMemoryReqs (VName -> Scope SOACS -> Bool
    forall k a. Ord k => k -> Map k a -> Bool
    `M.member` Scope SOACS
    scope) HostCode
    tcode
      let matches :: TExp Bool
    matches = [SubExp] -> [Maybe PrimValue] -> TExp Bool
    caseMatch [SubExp]
    cond (Case (Body GPUMem) -> [Maybe PrimValue]
    forall body. Case body -> [Maybe PrimValue]
    casePat Case (Body GPUMem)
    first_case)
      HostCode -> CallKernelGen ()
    forall op rep r. Code op -> ImpM rep r op ()
    emit (HostCode -> CallKernelGen ()) -> HostCode -> CallKernelGen ()
    forall a b. (a -> b) -> a -> b
    $ case Maybe (TExp Bool)
    check of
        Maybe (TExp Bool)
    Nothing -> HostCode
    fcode
        Just TExp Bool
    ok -> TExp Bool -> HostCode -> HostCode -> HostCode
    forall a. TExp Bool -> Code a -> Code a -> Code a
    Imp.If (TExp Bool
    matches TExp Bool -> TExp Bool -> TExp Bool
    forall v.
    Eq v =>
    TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v
    .&&. TExp Bool
    ok) HostCode
    tcode HostCode
    fcode
    expCompiler Pat (LetDec GPUMem)
    dest Exp GPUMem
    e =
      ExpCompiler GPUMem HostEnv HostOp
    forall rep (inner :: * -> *) r op.
    Mem rep inner =>
    Pat (LetDec rep) -> Exp rep -> ImpM rep r op ()
    defCompileExp Pat (LetDec GPUMem)
    dest Exp GPUMem
    e