    {-# LANGUAGE TypeFamilies #-}

    module Futhark.IR.GPUMem
      ( GPUMem,

        -- * Simplification
        simplifyProg,
        simplifyStms,
        simpleGPUMem,

        -- * Module re-exports
        module Futhark.IR.Mem,
        module Futhark.IR.GPU.Op,
      )
    where

    import Futhark.Analysis.PrimExp.Convert
    import Futhark.Analysis.UsageTable qualified as UT
    import Futhark.IR.Aliases (Aliases)
    import Futhark.IR.GPU.Op
    import Futhark.IR.GPU.Simplify (simplifyKernelOp)
    import Futhark.IR.Mem
    import Futhark.IR.Mem.Simplify
    import Futhark.IR.TypeCheck qualified as TC
    import Futhark.MonadFreshNames
    import Futhark.Optimise.Simplify.Engine qualified as Engine
    import Futhark.Pass
    import Futhark.Pass.ExplicitAllocations (BuilderOps (..), mkLetNamesB', mkLetNamesB'')

    data GPUMem

    instance RepTypes GPUMem where
      type LetDec GPUMem = LetDecMem
      type FParamInfo GPUMem = FParamMem
      type LParamInfo GPUMem = LParamMem
      type RetType GPUMem = RetTypeMem
      type BranchType GPUMem = BranchTypeMem
      type OpC GPUMem = MemOp (HostOp NoOp)

    instance ASTRep GPUMem where
      expTypesFromPat :: forall (m :: * -> *).
    (HasScope GPUMem m, Monad m) =>
    Pat (LetDec GPUMem) -> m [BranchType GPUMem]
    expTypesFromPat = [BranchTypeMem] -> m [BranchTypeMem]
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([BranchTypeMem] -> m [BranchTypeMem])
    -> (Pat LetDecMem -> [BranchTypeMem])
    -> Pat LetDecMem
    -> m [BranchTypeMem]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ((VName, BranchTypeMem) -> BranchTypeMem)
    -> [(VName, BranchTypeMem)] -> [BranchTypeMem]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName, BranchTypeMem) -> BranchTypeMem
    forall a b. (a, b) -> b
    snd ([(VName, BranchTypeMem)] -> [BranchTypeMem])
    -> (Pat LetDecMem -> [(VName, BranchTypeMem)])
    -> Pat LetDecMem
    -> [BranchTypeMem]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Pat LetDecMem -> [(VName, BranchTypeMem)]
    bodyReturnsFromPat

    instance PrettyRep GPUMem

    instance TC.Checkable GPUMem where
      checkOp :: Op (Aliases GPUMem) -> TypeM GPUMem ()
    checkOp = Maybe SegLevel
    -> MemOp (HostOp NoOp) (Aliases GPUMem) -> TypeM GPUMem ()
    typeCheckMemoryOp Maybe SegLevel
    forall a. Maybe a
    Nothing
        where
          -- GHC 9.2 goes into an infinite loop without the type annotation.
          typeCheckMemoryOp ::
            Maybe SegLevel ->
            MemOp (HostOp NoOp) (Aliases GPUMem) ->
            TC.TypeM GPUMem ()
          typeCheckMemoryOp :: Maybe SegLevel
    -> MemOp (HostOp NoOp) (Aliases GPUMem) -> TypeM GPUMem ()
    typeCheckMemoryOp Maybe SegLevel
    _ (Alloc SubExp
    size Space
    _) =
            [TypeBase (ShapeBase SubExp) NoUniqueness]
    -> SubExp -> TypeM GPUMem ()
    forall rep.
    Checkable rep =>
    [TypeBase (ShapeBase SubExp) NoUniqueness]
    -> SubExp -> TypeM rep ()
    TC.require [PrimType -> TypeBase (ShapeBase SubExp) NoUniqueness
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64] SubExp
    size
          typeCheckMemoryOp Maybe SegLevel
    lvl (Inner HostOp NoOp (Aliases GPUMem)
    op) =
            (SegLevel -> Op (Aliases GPUMem) -> TypeM GPUMem ())
    -> Maybe SegLevel
    -> (NoOp (Aliases GPUMem) -> TypeM GPUMem ())
    -> HostOp NoOp (Aliases GPUMem)
    -> TypeM GPUMem ()
    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 (Maybe SegLevel
    -> MemOp (HostOp NoOp) (Aliases GPUMem) -> TypeM GPUMem ()
    typeCheckMemoryOp (Maybe SegLevel
     -> MemOp (HostOp NoOp) (Aliases GPUMem) -> TypeM GPUMem ())
    -> (SegLevel -> Maybe SegLevel)
    -> SegLevel
    -> MemOp (HostOp NoOp) (Aliases GPUMem)
    -> TypeM GPUMem ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegLevel -> Maybe SegLevel
    forall a. a -> Maybe a
    Just) Maybe SegLevel
    lvl (TypeM GPUMem () -> NoOp (Aliases GPUMem) -> TypeM GPUMem ()
    forall a b. a -> b -> a
    const (TypeM GPUMem () -> NoOp (Aliases GPUMem) -> TypeM GPUMem ())
    -> TypeM GPUMem () -> NoOp (Aliases GPUMem) -> TypeM GPUMem ()
    forall a b. (a -> b) -> a -> b
    $ () -> TypeM GPUMem ()
    forall a. a -> TypeM GPUMem a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()) HostOp NoOp (Aliases GPUMem)
    op
      checkFParamDec :: VName -> FParamInfo GPUMem -> TypeM GPUMem ()
    checkFParamDec = VName -> FParamInfo GPUMem -> TypeM GPUMem ()
    VName -> FParamMem -> TypeM GPUMem ()
    forall rep u.
    Checkable rep =>
    VName -> MemInfo SubExp u MemBind -> TypeM rep ()
    checkMemInfo
      checkLParamDec :: VName -> LParamInfo GPUMem -> TypeM GPUMem ()
    checkLParamDec = VName -> LParamInfo GPUMem -> TypeM GPUMem ()
    VName -> LetDecMem -> TypeM GPUMem ()
    forall rep u.
    Checkable rep =>
    VName -> MemInfo SubExp u MemBind -> TypeM rep ()
    checkMemInfo
      checkLetBoundDec :: VName -> LetDec GPUMem -> TypeM GPUMem ()
    checkLetBoundDec = VName -> LetDec GPUMem -> TypeM GPUMem ()
    VName -> LetDecMem -> TypeM GPUMem ()
    forall rep u.
    Checkable rep =>
    VName -> MemInfo SubExp u MemBind -> TypeM rep ()
    checkMemInfo
      checkRetType :: [RetType GPUMem] -> TypeM GPUMem ()
    checkRetType = (RetType GPUMem -> TypeM GPUMem ())
    -> [RetType GPUMem] -> TypeM GPUMem ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ ((RetType GPUMem -> TypeM GPUMem ())
     -> [RetType GPUMem] -> TypeM GPUMem ())
    -> (RetType GPUMem -> TypeM GPUMem ())
    -> [RetType GPUMem]
    -> TypeM GPUMem ()
    forall a b. (a -> b) -> a -> b
    $ TypeBase (ShapeBase (Ext SubExp)) Uniqueness -> TypeM GPUMem ()
    forall rep u.
    Checkable rep =>
    TypeBase (ShapeBase (Ext SubExp)) u -> TypeM rep ()
    TC.checkExtType (TypeBase (ShapeBase (Ext SubExp)) Uniqueness -> TypeM GPUMem ())
    -> (RetTypeMem -> TypeBase (ShapeBase (Ext SubExp)) Uniqueness)
    -> RetTypeMem
    -> TypeM GPUMem ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . RetTypeMem -> TypeBase (ShapeBase (Ext SubExp)) Uniqueness
    forall t.
    DeclExtTyped t =>
    t -> TypeBase (ShapeBase (Ext SubExp)) Uniqueness
    declExtTypeOf
      primFParam :: VName -> PrimType -> TypeM GPUMem (FParam (Aliases GPUMem))
    primFParam VName
    name PrimType
    t = FParam (Aliases GPUMem) -> TypeM GPUMem (FParam (Aliases GPUMem))
    forall a. a -> TypeM GPUMem a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (FParam (Aliases GPUMem) -> TypeM GPUMem (FParam (Aliases GPUMem)))
    -> FParam (Aliases GPUMem)
    -> TypeM GPUMem (FParam (Aliases GPUMem))
    forall a b. (a -> b) -> a -> b
    $ Attrs -> VName -> FParamMem -> Param FParamMem
    forall dec. Attrs -> VName -> dec -> Param dec
    Param Attrs
    forall a. Monoid a => a
    mempty VName
    name (PrimType -> FParamMem
    forall d u ret. PrimType -> MemInfo d u ret
    MemPrim PrimType
    t)
      matchPat :: Pat (LetDec (Aliases GPUMem))
    -> Exp (Aliases GPUMem) -> TypeM GPUMem ()
    matchPat = Pat (LetDec (Aliases GPUMem))
    -> Exp (Aliases GPUMem) -> TypeM GPUMem ()
    forall rep (inner :: * -> *).
    (Mem rep inner, LetDec rep ~ LetDecMem, Checkable rep) =>
    Pat (LetDec (Aliases rep)) -> Exp (Aliases rep) -> TypeM rep ()
    matchPatToExp
      matchReturnType :: [RetType GPUMem] -> Result -> TypeM GPUMem ()
    matchReturnType = [RetType GPUMem] -> Result -> TypeM GPUMem ()
    [RetTypeMem] -> Result -> TypeM GPUMem ()
    forall rep (inner :: * -> *).
    (Mem rep inner, Checkable rep) =>
    [RetTypeMem] -> Result -> TypeM rep ()
    matchFunctionReturnType
      matchBranchType :: [BranchType GPUMem] -> Body (Aliases GPUMem) -> TypeM GPUMem ()
    matchBranchType = [BranchType GPUMem] -> Body (Aliases GPUMem) -> TypeM GPUMem ()
    [BranchTypeMem] -> Body (Aliases GPUMem) -> TypeM GPUMem ()
    forall rep (inner :: * -> *).
    (Mem rep inner, Checkable rep) =>
    [BranchTypeMem] -> Body (Aliases rep) -> TypeM rep ()
    matchBranchReturnType
      matchLoopResult :: [FParam (Aliases GPUMem)] -> Result -> TypeM GPUMem ()
    matchLoopResult = [FParam (Aliases GPUMem)] -> Result -> TypeM GPUMem ()
    forall rep (inner :: * -> *).
    (Mem rep inner, Checkable rep) =>
    [FParam (Aliases rep)] -> Result -> TypeM rep ()
    matchLoopResultMem

    instance BuilderOps GPUMem where
      mkExpDecB :: forall (m :: * -> *).
    (MonadBuilder m, Rep m ~ GPUMem) =>
    Pat (LetDec GPUMem) -> Exp GPUMem -> m (ExpDec GPUMem)
    mkExpDecB Pat (LetDec GPUMem)
    _ Exp GPUMem
    _ = () -> m ()
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
      mkBodyB :: forall (m :: * -> *).
    (MonadBuilder m, Rep m ~ GPUMem) =>
    Stms GPUMem -> Result -> m (Body GPUMem)
    mkBodyB Stms GPUMem
    stms Result
    res = Body GPUMem -> m (Body GPUMem)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Body GPUMem -> m (Body GPUMem)) -> Body GPUMem -> m (Body GPUMem)
    forall a b. (a -> b) -> a -> b
    $ BodyDec GPUMem -> Stms GPUMem -> Result -> Body GPUMem
    forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
    Body () Stms GPUMem
    stms Result
    res
      mkLetNamesB :: forall (m :: * -> *).
    (MonadBuilder m, Rep m ~ GPUMem) =>
    [VName] -> Exp GPUMem -> m (Stm GPUMem)
    mkLetNamesB = Space
    -> ExpDec (Rep m) -> [VName] -> Exp (Rep m) -> m (Stm (Rep m))
    forall (m :: * -> *) (inner :: * -> *).
    (LetDec (Rep m) ~ LetDecMem, Mem (Rep m) inner, MonadBuilder m,
     ExpDec (Rep m) ~ ()) =>
    Space
    -> ExpDec (Rep m) -> [VName] -> Exp (Rep m) -> m (Stm (Rep m))
    mkLetNamesB' (SpaceId -> Space
    Space SpaceId
    "device") ()

    instance BuilderOps (Engine.Wise GPUMem) where
      mkExpDecB :: forall (m :: * -> *).
    (MonadBuilder m, Rep m ~ Wise GPUMem) =>
    Pat (LetDec (Wise GPUMem))
    -> Exp (Wise GPUMem) -> m (ExpDec (Wise GPUMem))
    mkExpDecB Pat (LetDec (Wise GPUMem))
    pat Exp (Wise GPUMem)
    e = ExpDec (Wise GPUMem) -> m (ExpDec (Wise GPUMem))
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (ExpDec (Wise GPUMem) -> m (ExpDec (Wise GPUMem)))
    -> ExpDec (Wise GPUMem) -> m (ExpDec (Wise GPUMem))
    forall a b. (a -> b) -> a -> b
    $ Pat (LetDec (Wise GPUMem))
    -> ExpDec GPUMem -> Exp (Wise GPUMem) -> ExpDec (Wise GPUMem)
    forall rep.
    Informing rep =>
    Pat (LetDec (Wise rep))
    -> ExpDec rep -> Exp (Wise rep) -> ExpDec (Wise rep)
    Engine.mkWiseExpDec Pat (LetDec (Wise GPUMem))
    pat () Exp (Wise GPUMem)
    e
      mkBodyB :: forall (m :: * -> *).
    (MonadBuilder m, Rep m ~ Wise GPUMem) =>
    Stms (Wise GPUMem) -> Result -> m (Body (Wise GPUMem))
    mkBodyB Stms (Wise GPUMem)
    stms Result
    res = Body (Wise GPUMem) -> m (Body (Wise GPUMem))
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Body (Wise GPUMem) -> m (Body (Wise GPUMem)))
    -> Body (Wise GPUMem) -> m (Body (Wise GPUMem))
    forall a b. (a -> b) -> a -> b
    $ BodyDec GPUMem
    -> Stms (Wise GPUMem) -> Result -> Body (Wise GPUMem)
    forall rep.
    Informing rep =>
    BodyDec rep -> Stms (Wise rep) -> Result -> Body (Wise rep)
    Engine.mkWiseBody () Stms (Wise GPUMem)
    stms Result
    res
      mkLetNamesB :: forall (m :: * -> *).
    (MonadBuilder m, Rep m ~ Wise GPUMem) =>
    [VName] -> Exp (Wise GPUMem) -> m (Stm (Wise GPUMem))
    mkLetNamesB = Space -> [VName] -> Exp (Wise GPUMem) -> m (Stm (Wise GPUMem))
    forall rep (inner :: * -> *) (m :: * -> *).
    (Mem rep inner, LetDec rep ~ LetDecMem, OpReturns inner,
     ExpDec rep ~ (), Rep m ~ Wise rep, HasScope (Wise rep) m,
     MonadBuilder m, AliasedOp inner, RephraseOp (MemOp inner),
     CanBeWise inner, ASTConstraints (inner (Wise rep))) =>
    Space -> [VName] -> Exp (Wise rep) -> m (Stm (Wise rep))
    mkLetNamesB'' (SpaceId -> Space
    Space SpaceId
    "device")

    instance TraverseOpStms (Engine.Wise GPUMem) where
      traverseOpStms :: forall (m :: * -> *).
    Monad m =>
    OpStmsTraverser m (Op (Wise GPUMem)) (Wise GPUMem)
    traverseOpStms = OpStmsTraverser m (HostOp NoOp (Wise GPUMem)) (Wise GPUMem)
    -> OpStmsTraverser
         m (MemOp (HostOp NoOp) (Wise GPUMem)) (Wise GPUMem)
    forall (m :: * -> *) (inner :: * -> *) rep.
    Monad m =>
    OpStmsTraverser m (inner rep) rep
    -> OpStmsTraverser m (MemOp inner rep) rep
    traverseMemOpStms (OpStmsTraverser m (NoOp (Wise GPUMem)) (Wise GPUMem)
    -> OpStmsTraverser m (HostOp NoOp (Wise GPUMem)) (Wise GPUMem)
    forall (m :: * -> *) (op :: * -> *) rep.
    Monad m =>
    OpStmsTraverser m (op rep) rep
    -> OpStmsTraverser m (HostOp op rep) rep
    traverseHostOpStms ((NoOp (Wise GPUMem) -> m (NoOp (Wise GPUMem)))
    -> OpStmsTraverser m (NoOp (Wise GPUMem)) (Wise GPUMem)
    forall a b. a -> b -> a
    const NoOp (Wise GPUMem) -> m (NoOp (Wise GPUMem))
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure))

    simplifyProg :: Prog GPUMem -> PassM (Prog GPUMem)
    simplifyProg :: Prog GPUMem -> PassM (Prog GPUMem)
    simplifyProg = RuleBook (Wise GPUMem)
    -> SimpleOps GPUMem -> Prog GPUMem -> PassM (Prog GPUMem)
    forall rep (inner :: * -> *).
    SimplifyMemory rep inner =>
    RuleBook (Wise rep)
    -> SimpleOps rep -> Prog rep -> PassM (Prog rep)
    simplifyProgGeneric RuleBook (Wise GPUMem)
    forall rep (inner :: * -> *).
    SimplifyMemory rep inner =>
    RuleBook (Wise rep)
    memRuleBook SimpleOps GPUMem
    simpleGPUMem

    simplifyStms ::
      (HasScope GPUMem m, MonadFreshNames m) => Stms GPUMem -> m (Stms GPUMem)
    simplifyStms :: forall (m :: * -> *).
    (HasScope GPUMem m, MonadFreshNames m) =>
    Stms GPUMem -> m (Stms GPUMem)
    simplifyStms = RuleBook (Wise GPUMem)
    -> SimpleOps GPUMem -> Stms GPUMem -> m (Stms GPUMem)
    forall rep (m :: * -> *) (inner :: * -> *).
    (HasScope rep m, MonadFreshNames m, SimplifyMemory rep inner) =>
    RuleBook (Wise rep) -> SimpleOps rep -> Stms rep -> m (Stms rep)
    simplifyStmsGeneric RuleBook (Wise GPUMem)
    forall rep (inner :: * -> *).
    SimplifyMemory rep inner =>
    RuleBook (Wise rep)
    memRuleBook SimpleOps GPUMem
    simpleGPUMem

    simpleGPUMem :: Engine.SimpleOps GPUMem
    simpleGPUMem :: SimpleOps GPUMem
    simpleGPUMem =
      (HostOp NoOp (Wise GPUMem) -> UsageTable)
    -> SimplifyOp GPUMem (HostOp NoOp (Wise GPUMem))
    -> SimpleOps GPUMem
    forall rep (inner :: * -> *).
    SimplifyMemory rep inner =>
    (inner (Wise rep) -> UsageTable)
    -> SimplifyOp rep (inner (Wise rep)) -> SimpleOps rep
    simpleGeneric HostOp NoOp (Wise GPUMem) -> UsageTable
    usage (SimplifyOp GPUMem (HostOp NoOp (Wise GPUMem)) -> SimpleOps GPUMem)
    -> SimplifyOp GPUMem (HostOp NoOp (Wise GPUMem))
    -> SimpleOps GPUMem
    forall a b. (a -> b) -> a -> b
    $ SimplifyOp GPUMem (NoOp (Wise GPUMem))
    -> SimplifyOp GPUMem (HostOp NoOp (Wise GPUMem))
    forall rep (op :: * -> *).
    (SimplifiableRep rep, BodyDec rep ~ ()) =>
    SimplifyOp rep (op (Wise rep))
    -> HostOp op (Wise rep)
    -> SimpleM rep (HostOp op (Wise rep), Stms (Wise rep))
    simplifyKernelOp (SimplifyOp GPUMem (NoOp (Wise GPUMem))
     -> SimplifyOp GPUMem (HostOp NoOp (Wise GPUMem)))
    -> SimplifyOp GPUMem (NoOp (Wise GPUMem))
    -> SimplifyOp GPUMem (HostOp NoOp (Wise GPUMem))
    forall a b. (a -> b) -> a -> b
    $ SimpleM GPUMem (NoOp (Wise GPUMem), Stms (Wise GPUMem))
    -> SimplifyOp GPUMem (NoOp (Wise GPUMem))
    forall a b. a -> b -> a
    const (SimpleM GPUMem (NoOp (Wise GPUMem), Stms (Wise GPUMem))
     -> SimplifyOp GPUMem (NoOp (Wise GPUMem)))
    -> SimpleM GPUMem (NoOp (Wise GPUMem), Stms (Wise GPUMem))
    -> SimplifyOp GPUMem (NoOp (Wise GPUMem))
    forall a b. (a -> b) -> a -> b
    $ (NoOp (Wise GPUMem), Stms (Wise GPUMem))
    -> SimpleM GPUMem (NoOp (Wise GPUMem), Stms (Wise GPUMem))
    forall a. a -> SimpleM GPUMem a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (NoOp (Wise GPUMem)
    forall {k} (rep :: k). NoOp rep
    NoOp, Stms (Wise GPUMem)
    forall a. Monoid a => a
    mempty)
      where
        -- Slightly hackily and very inefficiently, we look at the inside
        -- of SegOps to figure out the sizes of shared memory allocations,
        -- and add usages for those sizes.  This is necessary so the
        -- simplifier will hoist those sizes out as far as possible (most
        -- importantly, past the versioning If, but see also #1569).
        usage :: HostOp NoOp (Wise GPUMem) -> UsageTable
    usage (SegOp (SegMap SegLevel
    _ SegSpace
    _ [TypeBase (ShapeBase SubExp) NoUniqueness]
    _ KernelBody (Wise GPUMem)
    kbody)) = KernelBody (Wise GPUMem) -> UsageTable
    localAllocs KernelBody (Wise GPUMem)
    kbody
        usage HostOp NoOp (Wise GPUMem)
    _ = UsageTable
    forall a. Monoid a => a
    mempty
        localAllocs :: KernelBody (Wise GPUMem) -> UsageTable
    localAllocs = (Stm (Wise GPUMem) -> UsageTable)
    -> Stms (Wise GPUMem) -> UsageTable
    forall m a. Monoid m => (a -> m) -> Seq a -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap Stm (Wise GPUMem) -> UsageTable
    stmSharedAlloc (Stms (Wise GPUMem) -> UsageTable)
    -> (KernelBody (Wise GPUMem) -> Stms (Wise GPUMem))
    -> KernelBody (Wise GPUMem)
    -> UsageTable
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . KernelBody (Wise GPUMem) -> Stms (Wise GPUMem)
    forall rep. KernelBody rep -> Stms rep
    kernelBodyStms
        stmSharedAlloc :: Stm (Wise GPUMem) -> UsageTable
    stmSharedAlloc = Exp (Wise GPUMem) -> UsageTable
    expSharedAlloc (Exp (Wise GPUMem) -> UsageTable)
    -> (Stm (Wise GPUMem) -> Exp (Wise GPUMem))
    -> Stm (Wise GPUMem)
    -> UsageTable
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Stm (Wise GPUMem) -> Exp (Wise GPUMem)
    forall rep. Stm rep -> Exp rep
    stmExp
        expSharedAlloc :: Exp (Wise GPUMem) -> UsageTable
    expSharedAlloc (Op (Alloc (Var VName
    v) Space
    _)) =
          VName -> UsageTable
    UT.sizeUsage VName
    v
        expSharedAlloc (Op (Inner (SegOp (SegMap SegLevel
    _ SegSpace
    _ [TypeBase (ShapeBase SubExp) NoUniqueness]
    _ KernelBody (Wise GPUMem)
    kbody)))) =
          KernelBody (Wise GPUMem) -> UsageTable
    localAllocs KernelBody (Wise GPUMem)
    kbody
        expSharedAlloc Exp (Wise GPUMem)
    _ =
          UsageTable
    forall a. Monoid a => a
    mempty