{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
-- |
-- Module      : Data.Array.Accelerate.LLVM.Native.CodeGen.Transform
-- Copyright   : [2014..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--

module Data.Array.Accelerate.LLVM.Native.CodeGen.Transform
  where

-- accelerate
import Data.Array.Accelerate.Representation.Array
import Data.Array.Accelerate.Type

import Data.Array.Accelerate.LLVM.CodeGen.Array
import Data.Array.Accelerate.LLVM.CodeGen.Base
import Data.Array.Accelerate.LLVM.CodeGen.Environment
import Data.Array.Accelerate.LLVM.CodeGen.Exp
import Data.Array.Accelerate.LLVM.CodeGen.Monad
import Data.Array.Accelerate.LLVM.CodeGen.Sugar
import Data.Array.Accelerate.LLVM.Compile.Cache

import Data.Array.Accelerate.LLVM.Native.Target                     ( Native )
import Data.Array.Accelerate.LLVM.Native.CodeGen.Base
import Data.Array.Accelerate.LLVM.Native.CodeGen.Loop


-- Hybrid map/backpermute operation
--
mkTransform
    :: UID
    -> Gamma aenv
    -> ArrayR (Array sh  a)
    -> ArrayR (Array sh' b)
    -> IRFun1  Native aenv (sh' -> sh)
    -> IRFun1  Native aenv (a -> b)
    -> CodeGen Native      (IROpenAcc Native aenv (Array sh' b))
mkTransform :: UID
-> Gamma aenv
-> ArrayR (Array sh a)
-> ArrayR (Array sh' b)
-> IRFun1 Native aenv (sh' -> sh)
-> IRFun1 Native aenv (a -> b)
-> CodeGen Native (IROpenAcc Native aenv (Array sh' b))
mkTransform UID
uid Gamma aenv
aenv ArrayR (Array sh a)
reprIn ArrayR (Array sh' b)
reprOut IRFun1 Native aenv (sh' -> sh)
p IRFun1 Native aenv (a -> b)
f =
  let
      (Operands sh'
start, Operands sh'
end, [Parameter]
paramGang)   = ShapeR sh' -> (Operands sh', Operands sh', [Parameter])
forall sh. ShapeR sh -> (Operands sh, Operands sh, [Parameter])
gangParam (ArrayR (Array sh' b) -> ShapeR sh'
forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape ArrayR (Array sh' b)
reprOut)
      (IRArray (Array sh a)
arrIn,  [Parameter]
paramIn)         = ArrayR (Array sh a)
-> Name (Array sh a) -> (IRArray (Array sh a), [Parameter])
forall sh e.
ArrayR (Array sh e)
-> Name (Array sh e) -> (IRArray (Array sh e), [Parameter])
mutableArray ArrayR (Array sh a)
reprIn  Name (Array sh a)
"in"
      (IRArray (Array sh' b)
arrOut, [Parameter]
paramOut)        = ArrayR (Array sh' b)
-> Name (Array sh' b) -> (IRArray (Array sh' b), [Parameter])
forall sh e.
ArrayR (Array sh e)
-> Name (Array sh e) -> (IRArray (Array sh e), [Parameter])
mutableArray ArrayR (Array sh' b)
reprOut Name (Array sh' b)
"out"
      paramEnv :: [Parameter]
paramEnv                  = Gamma aenv -> [Parameter]
forall aenv. Gamma aenv -> [Parameter]
envParam Gamma aenv
aenv
      shIn :: Operands sh
shIn                      = IRArray (Array sh a) -> Operands sh
forall sh e. IRArray (Array sh e) -> Operands sh
irArrayShape IRArray (Array sh a)
arrIn
      shOut :: Operands sh'
shOut                     = IRArray (Array sh' b) -> Operands sh'
forall sh e. IRArray (Array sh e) -> Operands sh
irArrayShape IRArray (Array sh' b)
arrOut
  in
  UID
-> Label
-> [Parameter]
-> CodeGen Native ()
-> CodeGen Native (IROpenAcc Native aenv (Array sh' b))
forall aenv a.
UID
-> Label
-> [Parameter]
-> CodeGen Native ()
-> CodeGen Native (IROpenAcc Native aenv a)
makeOpenAcc UID
uid Label
"transform" ([Parameter]
paramGang [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramOut [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramIn [Parameter] -> [Parameter] -> [Parameter]
forall a. [a] -> [a] -> [a]
++ [Parameter]
paramEnv) (CodeGen Native ()
 -> CodeGen Native (IROpenAcc Native aenv (Array sh' b)))
-> CodeGen Native ()
-> CodeGen Native (IROpenAcc Native aenv (Array sh' b))
forall a b. (a -> b) -> a -> b
$ do

    ShapeR sh'
-> Operands sh'
-> Operands sh'
-> Operands sh'
-> (Operands sh' -> Operands Int -> CodeGen Native ())
-> CodeGen Native ()
forall sh.
ShapeR sh
-> Operands sh
-> Operands sh
-> Operands sh
-> (Operands sh -> Operands Int -> CodeGen Native ())
-> CodeGen Native ()
imapNestFromTo (ArrayR (Array sh' b) -> ShapeR sh'
forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape ArrayR (Array sh' b)
reprOut) Operands sh'
start Operands sh'
end Operands sh'
shOut ((Operands sh' -> Operands Int -> CodeGen Native ())
 -> CodeGen Native ())
-> (Operands sh' -> Operands Int -> CodeGen Native ())
-> CodeGen Native ()
forall a b. (a -> b) -> a -> b
$ \Operands sh'
ix' Operands Int
i' -> do
      Operands sh
ix  <- IRFun1 Native aenv (sh' -> sh)
-> Operands sh' -> IROpenExp Native ((), sh') aenv sh
forall arch env aenv a b.
IROpenFun1 arch env aenv (a -> b)
-> Operands a -> IROpenExp arch (env, a) aenv b
app1 IRFun1 Native aenv (sh' -> sh)
p Operands sh'
ix'
      Operands Int
i   <- ShapeR sh
-> Operands sh -> Operands sh -> CodeGen Native (Operands Int)
forall sh arch.
ShapeR sh
-> Operands sh -> Operands sh -> CodeGen arch (Operands Int)
intOfIndex (ArrayR (Array sh a) -> ShapeR sh
forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape ArrayR (Array sh a)
reprIn) Operands sh
shIn Operands sh
ix
      Operands a
a   <- IntegralType Int
-> IRArray (Array sh a)
-> Operands Int
-> CodeGen Native (Operands a)
forall int sh e arch.
IntegralType int
-> IRArray (Array sh e)
-> Operands int
-> CodeGen arch (Operands e)
readArray IntegralType Int
TypeInt IRArray (Array sh a)
arrIn Operands Int
i
      Operands b
b   <- IRFun1 Native aenv (a -> b)
-> Operands a -> IROpenExp Native ((), a) aenv b
forall arch env aenv a b.
IROpenFun1 arch env aenv (a -> b)
-> Operands a -> IROpenExp arch (env, a) aenv b
app1 IRFun1 Native aenv (a -> b)
f Operands a
a
      IntegralType Int
-> IRArray (Array sh' b)
-> Operands Int
-> Operands b
-> CodeGen Native ()
forall int sh e arch.
IntegralType int
-> IRArray (Array sh e)
-> Operands int
-> Operands e
-> CodeGen arch ()
writeArray IntegralType Int
TypeInt IRArray (Array sh' b)
arrOut Operands Int
i' Operands b
b

    CodeGen Native ()
forall arch. HasCallStack => CodeGen arch ()
return_