    {-# LANGUAGE TypeFamilies #-}
    {-# OPTIONS_GHC -fno-warn-orphans #-}

    module Futhark.Pass.ExplicitAllocations.SegOp
      ( allocInKernelBody,
        allocInBinOpLambda,
      )
    where

    import Control.Monad
    import Futhark.IR.GPUMem
    import Futhark.IR.Mem.LMAD qualified as LMAD
    import Futhark.Pass.ExplicitAllocations

    instance SizeSubst (SegOp lvl rep)

    allocInKernelBody ::
      (Allocable fromrep torep inner) =>
      KernelBody fromrep ->
      AllocM fromrep torep (KernelBody torep)
    allocInKernelBody :: forall fromrep torep (inner :: * -> *).
    Allocable fromrep torep inner =>
    KernelBody fromrep -> AllocM fromrep torep (KernelBody torep)
    allocInKernelBody (KernelBody () Stms fromrep
    stms [KernelResult]
    res) =
      ([KernelResult] -> Stms torep -> KernelBody torep)
    -> ([KernelResult], Stms torep) -> KernelBody torep
    forall a b c. (a -> b -> c) -> (a, b) -> c
    uncurry ((Stms torep -> [KernelResult] -> KernelBody torep)
    -> [KernelResult] -> Stms torep -> KernelBody torep
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip (BodyDec torep -> Stms torep -> [KernelResult] -> KernelBody torep
    forall rep.
    BodyDec rep -> Stms rep -> [KernelResult] -> KernelBody rep
    KernelBody ()))
        (([KernelResult], Stms torep) -> KernelBody torep)
    -> AllocM fromrep torep ([KernelResult], Stms torep)
    -> AllocM fromrep torep (KernelBody torep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> AllocM fromrep torep [KernelResult]
    -> AllocM
         fromrep torep ([KernelResult], Stms (Rep (AllocM fromrep torep)))
    forall a.
    AllocM fromrep torep a
    -> AllocM fromrep torep (a, Stms (Rep (AllocM fromrep torep)))
    forall (m :: * -> *) a.
    MonadBuilder m =>
    m a -> m (a, Stms (Rep m))
    collectStms (Stms fromrep
    -> AllocM fromrep torep [KernelResult]
    -> AllocM fromrep torep [KernelResult]
    forall fromrep torep (inner :: * -> *) a.
    Allocable fromrep torep inner =>
    Stms fromrep -> AllocM fromrep torep a -> AllocM fromrep torep a
    allocInStms Stms fromrep
    stms ([KernelResult] -> AllocM fromrep torep [KernelResult]
    forall a. a -> AllocM fromrep torep a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [KernelResult]
    res))

    allocInLambda ::
      (Allocable fromrep torep inner) =>
      [LParam torep] ->
      Body fromrep ->
      AllocM fromrep torep (Lambda torep)
    allocInLambda :: forall fromrep torep (inner :: * -> *).
    Allocable fromrep torep inner =>
    [LParam torep]
    -> Body fromrep -> AllocM fromrep torep (Lambda torep)
    allocInLambda [LParam torep]
    params Body fromrep
    body =
      [LParam (Rep (AllocM fromrep torep))]
    -> AllocM fromrep torep Result
    -> AllocM fromrep torep (Lambda (Rep (AllocM fromrep torep)))
    forall (m :: * -> *).
    MonadBuilder m =>
    [LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
    mkLambda [LParam torep]
    [LParam (Rep (AllocM fromrep torep))]
    params (AllocM fromrep torep Result
     -> AllocM fromrep torep (Lambda torep))
    -> (AllocM fromrep torep Result -> AllocM fromrep torep Result)
    -> AllocM fromrep torep Result
    -> AllocM fromrep torep (Lambda torep)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Stms fromrep
    -> AllocM fromrep torep Result -> AllocM fromrep torep Result
    forall fromrep torep (inner :: * -> *) a.
    Allocable fromrep torep inner =>
    Stms fromrep -> AllocM fromrep torep a -> AllocM fromrep torep a
    allocInStms (Body fromrep -> Stms fromrep
    forall rep. Body rep -> Stms rep
    bodyStms Body fromrep
    body) (AllocM fromrep torep Result
     -> AllocM fromrep torep (Lambda torep))
    -> AllocM fromrep torep Result
    -> AllocM fromrep torep (Lambda torep)
    forall a b. (a -> b) -> a -> b
    $
        Result -> AllocM fromrep torep Result
    forall a. a -> AllocM fromrep torep a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Result -> AllocM fromrep torep Result)
    -> Result -> AllocM fromrep torep Result
    forall a b. (a -> b) -> a -> b
    $
          Body fromrep -> Result
    forall rep. Body rep -> Result
    bodyResult Body fromrep
    body

    allocInBinOpParams ::
      (Allocable fromrep torep inner) =>
      SubExp ->
      TPrimExp Int64 VName ->
      TPrimExp Int64 VName ->
      [LParam fromrep] ->
      [LParam fromrep] ->
      AllocM fromrep torep ([LParam torep], [LParam torep])
    allocInBinOpParams :: forall fromrep torep (inner :: * -> *).
    Allocable fromrep torep inner =>
    SubExp
    -> TPrimExp Int64 VName
    -> TPrimExp Int64 VName
    -> [LParam fromrep]
    -> [LParam fromrep]
    -> AllocM fromrep torep ([LParam torep], [LParam torep])
    allocInBinOpParams SubExp
    num_threads TPrimExp Int64 VName
    my_id TPrimExp Int64 VName
    other_id [LParam fromrep]
    xs [LParam fromrep]
    ys = [(Param LParamMem, Param LParamMem)]
    -> ([Param LParamMem], [Param LParamMem])
    forall a b. [(a, b)] -> ([a], [b])
    unzip ([(Param LParamMem, Param LParamMem)]
     -> ([Param LParamMem], [Param LParamMem]))
    -> AllocM fromrep torep [(Param LParamMem, Param LParamMem)]
    -> AllocM fromrep torep ([Param LParamMem], [Param LParamMem])
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Param Type
     -> Param Type
     -> AllocM fromrep torep (Param LParamMem, Param LParamMem))
    -> [Param Type]
    -> [Param Type]
    -> AllocM fromrep torep [(Param LParamMem, Param LParamMem)]
    forall (m :: * -> *) a b c.
    Applicative m =>
    (a -> b -> m c) -> [a] -> [b] -> m [c]
    zipWithM Param Type
    -> Param Type
    -> AllocM fromrep torep (Param LParamMem, Param LParamMem)
    alloc [Param Type]
    [LParam fromrep]
    xs [Param Type]
    [LParam fromrep]
    ys
      where
        alloc :: Param Type
    -> Param Type
    -> AllocM fromrep torep (Param LParamMem, Param LParamMem)
    alloc Param Type
    x Param Type
    y =
          case Param Type -> Type
    forall dec. Typed dec => Param dec -> Type
    paramType Param Type
    x of
            Array PrimType
    pt ShapeBase SubExp
    shape NoUniqueness
    u -> do
              let name :: String
    name = String -> (VName -> String) -> Maybe VName -> String
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe String
    "num_threads" VName -> String
    baseString (SubExp -> Maybe VName
    subExpVar SubExp
    num_threads)
              twice_num_threads <-
                String
    -> Exp (Rep (AllocM fromrep torep)) -> AllocM fromrep torep SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m SubExp
    letSubExp (String
    "twice_" String -> String -> String
    forall a. Semigroup a => a -> a -> a
    <> String
    name) (Exp torep -> AllocM fromrep torep SubExp)
    -> (BasicOp -> Exp torep) -> BasicOp -> AllocM fromrep torep SubExp
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . BasicOp -> Exp torep
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> AllocM fromrep torep SubExp)
    -> BasicOp -> AllocM fromrep torep SubExp
    forall a b. (a -> b) -> a -> b
    $
                  BinOp -> SubExp -> SubExp -> BasicOp
    BinOp (IntType -> Overflow -> BinOp
    Mul IntType
    Int64 Overflow
    OverflowUndef) SubExp
    num_threads (IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    2)
              let t = Param Type -> Type
    forall dec. Typed dec => Param dec -> Type
    paramType Param Type
    x Type -> SubExp -> Type
    forall d.
    ArrayShape (ShapeBase d) =>
    TypeBase (ShapeBase d) NoUniqueness
    -> d -> TypeBase (ShapeBase d) NoUniqueness
    `arrayOfRow` SubExp
    twice_num_threads
              mem <- allocForArray t =<< askDefaultSpace
              -- XXX: this iota lmad is a bit inefficient; leading to
              -- uncoalesced access.
              let base_dims = (SubExp -> TPrimExp Int64 VName)
    -> [SubExp] -> [TPrimExp Int64 VName]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> TPrimExp Int64 VName
    pe64 ([SubExp] -> [TPrimExp Int64 VName])
    -> [SubExp] -> [TPrimExp Int64 VName]
    forall a b. (a -> b) -> a -> b
    $ Type -> [SubExp]
    forall u. TypeBase (ShapeBase SubExp) u -> [SubExp]
    arrayDims Type
    t
                  lmad_base = TPrimExp Int64 VName
    -> [TPrimExp Int64 VName] -> LMAD (TPrimExp Int64 VName)
    forall num. IntegralExp num => num -> [num] -> LMAD num
    LMAD.iota TPrimExp Int64 VName
    0 [TPrimExp Int64 VName]
    base_dims
                  lmad_x =
                    LMAD (TPrimExp Int64 VName)
    -> Slice (TPrimExp Int64 VName) -> LMAD (TPrimExp Int64 VName)
    forall num.
    (Eq num, IntegralExp num) =>
    LMAD num -> Slice num -> LMAD num
    LMAD.slice LMAD (TPrimExp Int64 VName)
    lmad_base (Slice (TPrimExp Int64 VName) -> LMAD (TPrimExp Int64 VName))
    -> Slice (TPrimExp Int64 VName) -> LMAD (TPrimExp Int64 VName)
    forall a b. (a -> b) -> a -> b
    $
                      [TPrimExp Int64 VName]
    -> [DimIndex (TPrimExp Int64 VName)]
    -> Slice (TPrimExp Int64 VName)
    forall d. Num d => [d] -> [DimIndex d] -> Slice d
    fullSliceNum [TPrimExp Int64 VName]
    base_dims [TPrimExp Int64 VName -> DimIndex (TPrimExp Int64 VName)
    forall d. d -> DimIndex d
    DimFix TPrimExp Int64 VName
    my_id]
                  lmad_y =
                    LMAD (TPrimExp Int64 VName)
    -> Slice (TPrimExp Int64 VName) -> LMAD (TPrimExp Int64 VName)
    forall num.
    (Eq num, IntegralExp num) =>
    LMAD num -> Slice num -> LMAD num
    LMAD.slice LMAD (TPrimExp Int64 VName)
    lmad_base (Slice (TPrimExp Int64 VName) -> LMAD (TPrimExp Int64 VName))
    -> Slice (TPrimExp Int64 VName) -> LMAD (TPrimExp Int64 VName)
    forall a b. (a -> b) -> a -> b
    $
                      [TPrimExp Int64 VName]
    -> [DimIndex (TPrimExp Int64 VName)]
    -> Slice (TPrimExp Int64 VName)
    forall d. Num d => [d] -> [DimIndex d] -> Slice d
    fullSliceNum [TPrimExp Int64 VName]
    base_dims [TPrimExp Int64 VName -> DimIndex (TPrimExp Int64 VName)
    forall d. d -> DimIndex d
    DimFix TPrimExp Int64 VName
    other_id]
              pure
                ( x {paramDec = MemArray pt shape u $ ArrayIn mem lmad_x},
                  y {paramDec = MemArray pt shape u $ ArrayIn mem lmad_y}
                )
            Prim PrimType
    bt ->
              (Param LParamMem, Param LParamMem)
    -> AllocM fromrep torep (Param LParamMem, Param LParamMem)
    forall a. a -> AllocM fromrep torep a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
                ( Param Type
    x {paramDec = MemPrim bt},
                  Param Type
    y {paramDec = MemPrim bt}
                )
            Mem Space
    space ->
              (Param LParamMem, Param LParamMem)
    -> AllocM fromrep torep (Param LParamMem, Param LParamMem)
    forall a. a -> AllocM fromrep torep a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
                ( Param Type
    x {paramDec = MemMem space},
                  Param Type
    y {paramDec = MemMem space}
                )
            -- This next case will never happen.
            Acc VName
    acc ShapeBase SubExp
    ispace [Type]
    ts NoUniqueness
    u ->
              (Param LParamMem, Param LParamMem)
    -> AllocM fromrep torep (Param LParamMem, Param LParamMem)
    forall a. a -> AllocM fromrep torep a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
                ( Param Type
    x {paramDec = MemAcc acc ispace ts u},
                  Param Type
    y {paramDec = MemAcc acc ispace ts u}
                )

    allocInBinOpLambda ::
      (Allocable fromrep torep inner) =>
      SubExp ->
      SegSpace ->
      Lambda fromrep ->
      AllocM fromrep torep (Lambda torep)
    allocInBinOpLambda :: forall fromrep torep (inner :: * -> *).
    Allocable fromrep torep inner =>
    SubExp
    -> SegSpace
    -> Lambda fromrep
    -> AllocM fromrep torep (Lambda torep)
    allocInBinOpLambda SubExp
    num_threads (SegSpace VName
    flat [(VName, SubExp)]
    _) Lambda fromrep
    lam = do
      let ([Param Type]
    acc_params, [Param Type]
    arr_params) =
            Int -> [Param Type] -> ([Param Type], [Param Type])
    forall a. Int -> [a] -> ([a], [a])
    splitAt ([Param Type] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length (Lambda fromrep -> [LParam fromrep]
    forall rep. Lambda rep -> [LParam rep]
    lambdaParams Lambda fromrep
    lam) Int -> Int -> Int
    forall a. Integral a => a -> a -> a
    `div` Int
    2) ([Param Type] -> ([Param Type], [Param Type]))
    -> [Param Type] -> ([Param Type], [Param Type])
    forall a b. (a -> b) -> a -> b
    $ Lambda fromrep -> [LParam fromrep]
    forall rep. Lambda rep -> [LParam rep]
    lambdaParams Lambda fromrep
    lam
          index_x :: TPrimExp Int64 VName
    index_x = PrimExp VName -> TPrimExp Int64 VName
    forall {k} (t :: k) v. PrimExp v -> TPrimExp t v
    TPrimExp (PrimExp VName -> TPrimExp Int64 VName)
    -> PrimExp VName -> TPrimExp Int64 VName
    forall a b. (a -> b) -> a -> b
    $ VName -> PrimType -> PrimExp VName
    forall v. v -> PrimType -> PrimExp v
    LeafExp VName
    flat PrimType
    int64
          index_y :: TPrimExp Int64 VName
    index_y = TPrimExp Int64 VName
    index_x TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    + SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    num_threads
      (acc_params', arr_params') <-
        SubExp
    -> TPrimExp Int64 VName
    -> TPrimExp Int64 VName
    -> [LParam fromrep]
    -> [LParam fromrep]
    -> AllocM fromrep torep ([LParam torep], [LParam torep])
    forall fromrep torep (inner :: * -> *).
    Allocable fromrep torep inner =>
    SubExp
    -> TPrimExp Int64 VName
    -> TPrimExp Int64 VName
    -> [LParam fromrep]
    -> [LParam fromrep]
    -> AllocM fromrep torep ([LParam torep], [LParam torep])
    allocInBinOpParams SubExp
    num_threads TPrimExp Int64 VName
    index_x TPrimExp Int64 VName
    index_y [Param Type]
    [LParam fromrep]
    acc_params [Param Type]
    [LParam fromrep]
    arr_params

      allocInLambda (acc_params' ++ arr_params') (lambdaBody lam)