-- | A generic Python code generator which is polymorphic in the type
-- of the operations.  Concretely, we use this to handle both
-- sequential and PyOpenCL Python code.
module Futhark.CodeGen.Backends.GenericPython
  ( compileProg,
    CompilerMode,
    Constructor (..),
    emptyConstructor,
    compileName,
    compileVar,
    compileDim,
    compileExp,
    compilePrimExp,
    compileCode,
    compilePrimValue,
    compilePrimType,
    compilePrimTypeExt,
    compilePrimToNp,
    compilePrimToExtNp,
    fromStorage,
    toStorage,
    Operations (..),
    defaultOperations,
    unpackDim,
    CompilerM (..),
    OpCompiler,
    WriteScalar,
    ReadScalar,
    Allocate,
    Copy,
    EntryOutput,
    EntryInput,
    CompilerEnv (..),
    CompilerState (..),
    stm,
    atInit,
    collect',
    collect,
    simpleCall,
    copyMemoryDefaultSpace,
  )
where

import Control.Monad.Identity
import Control.Monad.RWS
import Data.Char (isAlpha, isAlphaNum)
import Data.Map qualified as M
import Data.Maybe
import Data.Text qualified as T
import Futhark.CodeGen.Backends.GenericPython.AST
import Futhark.CodeGen.Backends.GenericPython.Options
import Futhark.CodeGen.ImpCode qualified as Imp
import Futhark.CodeGen.RTS.Python
import Futhark.Compiler.Config (CompilerMode (..))
import Futhark.IR.Prop (isBuiltInFunction, subExpVars)
import Futhark.IR.Syntax.Core (Space (..))
import Futhark.MonadFreshNames
import Futhark.Util (zEncodeText)
import Futhark.Util.Pretty (prettyString, prettyText)
import Language.Futhark.Primitive hiding (Bool)

-- | A substitute expression compiler, tried before the main
-- compilation function.
type OpCompiler op s = op -> CompilerM op s ()

-- | Write a scalar to the given memory block with the given index and
-- in the given memory space.
type WriteScalar op s =
  PyExp ->
  PyExp ->
  PrimType ->
  Imp.SpaceId ->
  PyExp ->
  CompilerM op s ()

-- | Read a scalar from the given memory block with the given index and
-- in the given memory space.
type ReadScalar op s =
  PyExp ->
  PyExp ->
  PrimType ->
  Imp.SpaceId ->
  CompilerM op s PyExp

-- | Allocate a memory block of the given size in the given memory
-- space, saving a reference in the given variable name.
type Allocate op s =
  PyExp ->
  PyExp ->
  Imp.SpaceId ->
  CompilerM op s ()

-- | Copy from one memory block to another.
type Copy op s =
  PyExp ->
  PyExp ->
  Imp.Space ->
  PyExp ->
  PyExp ->
  Imp.Space ->
  PyExp ->
  PrimType ->
  CompilerM op s ()

-- | Construct the Python array being returned from an entry point.
type EntryOutput op s =
  VName ->
  Imp.SpaceId ->
  PrimType ->
  Imp.Signedness ->
  [Imp.DimSize] ->
  CompilerM op s PyExp

-- | Unpack the array being passed to an entry point.
type EntryInput op s =
  PyExp ->
  Imp.SpaceId ->
  PrimType ->
  Imp.Signedness ->
  [Imp.DimSize] ->
  PyExp ->
  CompilerM op s ()

data Operations op s = Operations
  { forall op s. Operations op s -> WriteScalar op s
opsWriteScalar :: WriteScalar op s,
    forall op s. Operations op s -> ReadScalar op s
opsReadScalar :: ReadScalar op s,
    forall op s. Operations op s -> Allocate op s
opsAllocate :: Allocate op s,
    forall op s. Operations op s -> Copy op s
opsCopy :: Copy op s,
    forall op s. Operations op s -> OpCompiler op s
opsCompiler :: OpCompiler op s,
    forall op s. Operations op s -> EntryOutput op s
opsEntryOutput :: EntryOutput op s,
    forall op s. Operations op s -> EntryInput op s
opsEntryInput :: EntryInput op s
  }

-- | A set of operations that fail for every operation involving
-- non-default memory spaces.  Uses plain pointers and @malloc@ for
-- memory management.
defaultOperations :: Operations op s
defaultOperations :: forall op s. Operations op s
defaultOperations =
  Operations
    { opsWriteScalar :: WriteScalar op s
opsWriteScalar = forall {p} {p} {p} {p} {p} {a}. p -> p -> p -> p -> p -> a
defWriteScalar,
      opsReadScalar :: ReadScalar op s
opsReadScalar = forall {p} {p} {p} {p} {a}. p -> p -> p -> p -> a
defReadScalar,
      opsAllocate :: Allocate op s
opsAllocate = forall {p} {p} {p} {a}. p -> p -> p -> a
defAllocate,
      opsCopy :: Copy op s
opsCopy = forall {p} {p} {p} {p} {p} {p} {p} {p} {a}.
p -> p -> p -> p -> p -> p -> p -> p -> a
defCopy,
      opsCompiler :: OpCompiler op s
opsCompiler = forall {p} {a}. p -> a
defCompiler,
      opsEntryOutput :: EntryOutput op s
opsEntryOutput = forall {p} {p} {p} {p} {a}. p -> p -> p -> p -> a
defEntryOutput,
      opsEntryInput :: EntryInput op s
opsEntryInput = forall {p} {p} {p} {p} {a}. p -> p -> p -> p -> a
defEntryInput
    }
  where
    defWriteScalar :: p -> p -> p -> p -> p -> a
defWriteScalar p
_ p
_ p
_ p
_ p
_ =
      forall a. HasCallStack => [Char] -> a
error [Char]
"Cannot write to non-default memory space because I am dumb"
    defReadScalar :: p -> p -> p -> p -> a
defReadScalar p
_ p
_ p
_ p
_ =
      forall a. HasCallStack => [Char] -> a
error [Char]
"Cannot read from non-default memory space"
    defAllocate :: p -> p -> p -> a
defAllocate p
_ p
_ p
_ =
      forall a. HasCallStack => [Char] -> a
error [Char]
"Cannot allocate in non-default memory space"
    defCopy :: p -> p -> p -> p -> p -> p -> p -> p -> a
defCopy p
_ p
_ p
_ p
_ p
_ p
_ p
_ p
_ =
      forall a. HasCallStack => [Char] -> a
error [Char]
"Cannot copy to or from non-default memory space"
    defCompiler :: p -> a
defCompiler p
_ =
      forall a. HasCallStack => [Char] -> a
error [Char]
"The default compiler cannot compile extended operations"
    defEntryOutput :: p -> p -> p -> p -> a
defEntryOutput p
_ p
_ p
_ p
_ =
      forall a. HasCallStack => [Char] -> a
error [Char]
"Cannot return array not in default memory space"
    defEntryInput :: p -> p -> p -> p -> a
defEntryInput p
_ p
_ p
_ p
_ =
      forall a. HasCallStack => [Char] -> a
error [Char]
"Cannot accept array not in default memory space"

data CompilerEnv op s = CompilerEnv
  { forall op s. CompilerEnv op s -> Operations op s
envOperations :: Operations op s,
    forall op s. CompilerEnv op s -> Map [Char] PyExp
envVarExp :: M.Map String PyExp
  }

envOpCompiler :: CompilerEnv op s -> OpCompiler op s
envOpCompiler :: forall op s. CompilerEnv op s -> OpCompiler op s
envOpCompiler = forall op s. Operations op s -> OpCompiler op s
opsCompiler forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall op s. CompilerEnv op s -> Operations op s
envOperations

envReadScalar :: CompilerEnv op s -> ReadScalar op s
envReadScalar :: forall op s. CompilerEnv op s -> ReadScalar op s
envReadScalar = forall op s. Operations op s -> ReadScalar op s
opsReadScalar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall op s. CompilerEnv op s -> Operations op s
envOperations

envWriteScalar :: CompilerEnv op s -> WriteScalar op s
envWriteScalar :: forall op s. CompilerEnv op s -> WriteScalar op s
envWriteScalar = forall op s. Operations op s -> WriteScalar op s
opsWriteScalar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall op s. CompilerEnv op s -> Operations op s
envOperations

envAllocate :: CompilerEnv op s -> Allocate op s
envAllocate :: forall op s. CompilerEnv op s -> Allocate op s
envAllocate = forall op s. Operations op s -> Allocate op s
opsAllocate forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall op s. CompilerEnv op s -> Operations op s
envOperations

envCopy :: CompilerEnv op s -> Copy op s
envCopy :: forall op s. CompilerEnv op s -> Copy op s
envCopy = forall op s. Operations op s -> Copy op s
opsCopy forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall op s. CompilerEnv op s -> Operations op s
envOperations

envEntryOutput :: CompilerEnv op s -> EntryOutput op s
envEntryOutput :: forall op s. CompilerEnv op s -> EntryOutput op s
envEntryOutput = forall op s. Operations op s -> EntryOutput op s
opsEntryOutput forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall op s. CompilerEnv op s -> Operations op s
envOperations

envEntryInput :: CompilerEnv op s -> EntryInput op s
envEntryInput :: forall op s. CompilerEnv op s -> EntryInput op s
envEntryInput = forall op s. Operations op s -> EntryInput op s
opsEntryInput forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall op s. CompilerEnv op s -> Operations op s
envOperations

newCompilerEnv :: Operations op s -> CompilerEnv op s
newCompilerEnv :: forall op s. Operations op s -> CompilerEnv op s
newCompilerEnv Operations op s
ops =
  CompilerEnv
    { envOperations :: Operations op s
envOperations = Operations op s
ops,
      envVarExp :: Map [Char] PyExp
envVarExp = forall a. Monoid a => a
mempty
    }

data CompilerState s = CompilerState
  { forall s. CompilerState s -> VNameSource
compNameSrc :: VNameSource,
    forall s. CompilerState s -> [PyStmt]
compInit :: [PyStmt],
    forall s. CompilerState s -> s
compUserState :: s
  }

newCompilerState :: VNameSource -> s -> CompilerState s
newCompilerState :: forall s. VNameSource -> s -> CompilerState s
newCompilerState VNameSource
src s
s =
  CompilerState
    { compNameSrc :: VNameSource
compNameSrc = VNameSource
src,
      compInit :: [PyStmt]
compInit = [],
      compUserState :: s
compUserState = s
s
    }

newtype CompilerM op s a = CompilerM (RWS (CompilerEnv op s) [PyStmt] (CompilerState s) a)
  deriving
    ( forall a b. a -> CompilerM op s b -> CompilerM op s a
forall a b. (a -> b) -> CompilerM op s a -> CompilerM op s b
forall op s a b. a -> CompilerM op s b -> CompilerM op s a
forall op s a b. (a -> b) -> CompilerM op s a -> CompilerM op s b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> CompilerM op s b -> CompilerM op s a
$c<$ :: forall op s a b. a -> CompilerM op s b -> CompilerM op s a
fmap :: forall a b. (a -> b) -> CompilerM op s a -> CompilerM op s b
$cfmap :: forall op s a b. (a -> b) -> CompilerM op s a -> CompilerM op s b
Functor,
      forall a. a -> CompilerM op s a
forall op s. Functor (CompilerM op s)
forall a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s a
forall a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
forall a b.
CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b
forall op s a. a -> CompilerM op s a
forall a b c.
(a -> b -> c)
-> CompilerM op s a -> CompilerM op s b -> CompilerM op s c
forall op s a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s a
forall op s a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
forall op s a b.
CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b
forall op s a b c.
(a -> b -> c)
-> CompilerM op s a -> CompilerM op s b -> CompilerM op s c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: forall a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s a
$c<* :: forall op s a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s a
*> :: forall a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
$c*> :: forall op s a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
liftA2 :: forall a b c.
(a -> b -> c)
-> CompilerM op s a -> CompilerM op s b -> CompilerM op s c
$cliftA2 :: forall op s a b c.
(a -> b -> c)
-> CompilerM op s a -> CompilerM op s b -> CompilerM op s c
<*> :: forall a b.
CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b
$c<*> :: forall op s a b.
CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b
pure :: forall a. a -> CompilerM op s a
$cpure :: forall op s a. a -> CompilerM op s a
Applicative,
      forall a. a -> CompilerM op s a
forall op s. Applicative (CompilerM op s)
forall a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
forall a b.
CompilerM op s a -> (a -> CompilerM op s b) -> CompilerM op s b
forall op s a. a -> CompilerM op s a
forall op s a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
forall op s a b.
CompilerM op s a -> (a -> CompilerM op s b) -> CompilerM op s b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: forall a. a -> CompilerM op s a
$creturn :: forall op s a. a -> CompilerM op s a
>> :: forall a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
$c>> :: forall op s a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
>>= :: forall a b.
CompilerM op s a -> (a -> CompilerM op s b) -> CompilerM op s b
$c>>= :: forall op s a b.
CompilerM op s a -> (a -> CompilerM op s b) -> CompilerM op s b
Monad,
      MonadState (CompilerState s),
      MonadReader (CompilerEnv op s),
      MonadWriter [PyStmt]
    )

instance MonadFreshNames (CompilerM op s) where
  getNameSource :: CompilerM op s VNameSource
getNameSource = forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall s. CompilerState s -> VNameSource
compNameSrc
  putNameSource :: VNameSource -> CompilerM op s ()
putNameSource VNameSource
src = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \CompilerState s
s -> CompilerState s
s {compNameSrc :: VNameSource
compNameSrc = VNameSource
src}

collect :: CompilerM op s () -> CompilerM op s [PyStmt]
collect :: forall op s. CompilerM op s () -> CompilerM op s [PyStmt]
collect CompilerM op s ()
m = forall w (m :: * -> *) a. MonadWriter w m => m (a, w -> w) -> m a
pass forall a b. (a -> b) -> a -> b
$ do
  ((), [PyStmt]
w) <- forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
listen CompilerM op s ()
m
  forall (f :: * -> *) a. Applicative f => a -> f a
pure ([PyStmt]
w, forall a b. a -> b -> a
const forall a. Monoid a => a
mempty)

collect' :: CompilerM op s a -> CompilerM op s (a, [PyStmt])
collect' :: forall op s a. CompilerM op s a -> CompilerM op s (a, [PyStmt])
collect' CompilerM op s a
m = forall w (m :: * -> *) a. MonadWriter w m => m (a, w -> w) -> m a
pass forall a b. (a -> b) -> a -> b
$ do
  (a
x, [PyStmt]
w) <- forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
listen CompilerM op s a
m
  forall (f :: * -> *) a. Applicative f => a -> f a
pure ((a
x, [PyStmt]
w), forall a b. a -> b -> a
const forall a. Monoid a => a
mempty)

atInit :: PyStmt -> CompilerM op s ()
atInit :: forall op s. PyStmt -> CompilerM op s ()
atInit PyStmt
x = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \CompilerState s
s ->
  CompilerState s
s {compInit :: [PyStmt]
compInit = forall s. CompilerState s -> [PyStmt]
compInit CompilerState s
s forall a. [a] -> [a] -> [a]
++ [PyStmt
x]}

stm :: PyStmt -> CompilerM op s ()
stm :: forall op s. PyStmt -> CompilerM op s ()
stm PyStmt
x = forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell [PyStmt
x]

futharkFun :: T.Text -> T.Text
futharkFun :: Text -> Text
futharkFun Text
s = Text
"futhark_" forall a. Semigroup a => a -> a -> a
<> Text -> Text
zEncodeText Text
s

compileOutput :: [Imp.Param] -> [PyExp]
compileOutput :: [Param] -> [PyExp]
compileOutput = forall a b. (a -> b) -> [a] -> [b]
map ([Char] -> PyExp
Var forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> [Char]
compileName forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param -> VName
Imp.paramName)

runCompilerM ::
  Operations op s ->
  VNameSource ->
  s ->
  CompilerM op s a ->
  a
runCompilerM :: forall op s a.
Operations op s -> VNameSource -> s -> CompilerM op s a -> a
runCompilerM Operations op s
ops VNameSource
src s
userstate (CompilerM RWS (CompilerEnv op s) [PyStmt] (CompilerState s) a
m) =
  forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall r w s a. RWS r w s a -> r -> s -> (a, w)
evalRWS RWS (CompilerEnv op s) [PyStmt] (CompilerState s) a
m (forall op s. Operations op s -> CompilerEnv op s
newCompilerEnv Operations op s
ops) (forall s. VNameSource -> s -> CompilerState s
newCompilerState VNameSource
src s
userstate)

standardOptions :: [Option]
standardOptions :: [Option]
standardOptions =
  [ Option
      { optionLongName :: Text
optionLongName = Text
"tuning",
        optionShortName :: Maybe Char
optionShortName = forall a. Maybe a
Nothing,
        optionArgument :: OptionArgument
optionArgument = [Char] -> OptionArgument
RequiredArgument [Char]
"open",
        optionAction :: [PyStmt]
optionAction = [PyExp -> PyStmt
Exp forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"read_tuning_file" [[Char] -> PyExp
Var [Char]
"sizes", [Char] -> PyExp
Var [Char]
"optarg"]]
      },
    -- Does not actually do anything for Python backends.
    Option
      { optionLongName :: Text
optionLongName = Text
"cache-file",
        optionShortName :: Maybe Char
optionShortName = forall a. Maybe a
Nothing,
        optionArgument :: OptionArgument
optionArgument = [Char] -> OptionArgument
RequiredArgument [Char]
"str",
        optionAction :: [PyStmt]
optionAction = [PyStmt
Pass]
      },
    Option
      { optionLongName :: Text
optionLongName = Text
"log",
        optionShortName :: Maybe Char
optionShortName = forall a. a -> Maybe a
Just Char
'L',
        optionArgument :: OptionArgument
optionArgument = OptionArgument
NoArgument,
        optionAction :: [PyStmt]
optionAction = [PyStmt
Pass]
      }
  ]

executableOptions :: [Option]
executableOptions :: [Option]
executableOptions =
  [Option]
standardOptions
    forall a. [a] -> [a] -> [a]
++ [ Option
           { optionLongName :: Text
optionLongName = Text
"write-runtime-to",
             optionShortName :: Maybe Char
optionShortName = forall a. a -> Maybe a
Just Char
't',
             optionArgument :: OptionArgument
optionArgument = [Char] -> OptionArgument
RequiredArgument [Char]
"str",
             optionAction :: [PyStmt]
optionAction =
               [ PyExp -> [PyStmt] -> [PyStmt] -> PyStmt
If
                   ([Char] -> PyExp
Var [Char]
"runtime_file")
                   [PyExp -> PyStmt
Exp forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"runtime_file.close" []]
                   [],
                 PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"runtime_file") forall a b. (a -> b) -> a -> b
$
                   [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"open" [[Char] -> PyExp
Var [Char]
"optarg", Text -> PyExp
String Text
"w"]
               ]
           },
         Option
           { optionLongName :: Text
optionLongName = Text
"runs",
             optionShortName :: Maybe Char
optionShortName = forall a. a -> Maybe a
Just Char
'r',
             optionArgument :: OptionArgument
optionArgument = [Char] -> OptionArgument
RequiredArgument [Char]
"str",
             optionAction :: [PyStmt]
optionAction =
               [ PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"num_runs") forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp
Var [Char]
"optarg",
                 PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"do_warmup_run") forall a b. (a -> b) -> a -> b
$ Bool -> PyExp
Bool Bool
True
               ]
           },
         Option
           { optionLongName :: Text
optionLongName = Text
"entry-point",
             optionShortName :: Maybe Char
optionShortName = forall a. a -> Maybe a
Just Char
'e',
             optionArgument :: OptionArgument
optionArgument = [Char] -> OptionArgument
RequiredArgument [Char]
"str",
             optionAction :: [PyStmt]
optionAction =
               [PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"entry_point") forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp
Var [Char]
"optarg"]
           },
         Option
           { optionLongName :: Text
optionLongName = Text
"binary-output",
             optionShortName :: Maybe Char
optionShortName = forall a. a -> Maybe a
Just Char
'b',
             optionArgument :: OptionArgument
optionArgument = OptionArgument
NoArgument,
             optionAction :: [PyStmt]
optionAction = [PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"binary_output") forall a b. (a -> b) -> a -> b
$ Bool -> PyExp
Bool Bool
True]
           }
       ]

functionExternalValues :: Imp.EntryPoint -> [Imp.ExternalValue]
functionExternalValues :: EntryPoint -> [ExternalValue]
functionExternalValues EntryPoint
entry =
  forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd (EntryPoint -> [(Uniqueness, ExternalValue)]
Imp.entryPointResults EntryPoint
entry) forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd (EntryPoint -> [((Name, Uniqueness), ExternalValue)]
Imp.entryPointArgs EntryPoint
entry)

-- | Is this name a valid Python identifier?  If not, it should be escaped
-- before being emitted.
isValidPyName :: T.Text -> Bool
isValidPyName :: Text -> Bool
isValidPyName = forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True (Char, Text) -> Bool
check forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Maybe (Char, Text)
T.uncons
  where
    check :: (Char, Text) -> Bool
check (Char
c, Text
cs) = Char -> Bool
isAlpha Char
c Bool -> Bool -> Bool
&& (Char -> Bool) -> Text -> Bool
T.all Char -> Bool
constituent Text
cs
    constituent :: Char -> Bool
constituent Char
c = Char -> Bool
isAlphaNum Char
c Bool -> Bool -> Bool
|| Char
c forall a. Eq a => a -> a -> Bool
== Char
'_'

-- | If the provided text is a valid identifier, then return it
-- verbatim.  Otherwise, escape it such that it becomes valid.
escapeName :: Name -> T.Text
escapeName :: Name -> Text
escapeName Name
v
  | Text -> Bool
isValidPyName Text
v' = Text
v'
  | Bool
otherwise = Text -> Text
zEncodeText Text
v'
  where
    v' :: Text
v' = Name -> Text
nameToText Name
v

opaqueDefs :: Imp.Functions a -> M.Map T.Text [PyExp]
opaqueDefs :: forall a. Functions a -> Map Text [PyExp]
opaqueDefs (Imp.Functions [(Name, Function a)]
funs) =
  forall a. Monoid a => [a] -> a
mconcat
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map ExternalValue -> Map Text [PyExp]
evd
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap EntryPoint -> [ExternalValue]
functionExternalValues
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (forall a. FunctionT a -> Maybe EntryPoint
Imp.functionEntry forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd)
    forall a b. (a -> b) -> a -> b
$ [(Name, Function a)]
funs
  where
    evd :: ExternalValue -> Map Text [PyExp]
evd Imp.TransparentValue {} = forall a. Monoid a => a
mempty
    evd (Imp.OpaqueValue Name
name [ValueDesc]
vds) = forall k a. k -> a -> Map k a
M.singleton (Name -> Text
nameToText Name
name) forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (Text -> PyExp
String forall b c a. (b -> c) -> (a -> b) -> a -> c
. ValueDesc -> Text
vd) [ValueDesc]
vds
    vd :: ValueDesc -> Text
vd (Imp.ScalarValue PrimType
pt Signedness
s VName
_) =
      PrimType -> Signedness -> Text
readTypeEnum PrimType
pt Signedness
s
    vd (Imp.ArrayValue VName
_ Space
_ PrimType
pt Signedness
s [DimSize]
dims) =
      forall a. Monoid a => [a] -> a
mconcat (forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimSize]
dims) Text
"[]") forall a. Semigroup a => a -> a -> a
<> PrimType -> Signedness -> Text
readTypeEnum PrimType
pt Signedness
s

-- | The class generated by the code generator must have a
-- constructor, although it can be vacuous.
data Constructor = Constructor [String] [PyStmt]

-- | A constructor that takes no arguments and does nothing.
emptyConstructor :: Constructor
emptyConstructor :: Constructor
emptyConstructor = [[Char]] -> [PyStmt] -> Constructor
Constructor [[Char]
"self"] [PyStmt
Pass]

constructorToFunDef :: Constructor -> [PyStmt] -> PyFunDef
constructorToFunDef :: Constructor -> [PyStmt] -> PyFunDef
constructorToFunDef (Constructor [[Char]]
params [PyStmt]
body) [PyStmt]
at_init =
  [Char] -> [[Char]] -> [PyStmt] -> PyFunDef
Def [Char]
"__init__" [[Char]]
params forall a b. (a -> b) -> a -> b
$ [PyStmt]
body forall a. Semigroup a => a -> a -> a
<> [PyStmt]
at_init

compileProg ::
  MonadFreshNames m =>
  CompilerMode ->
  String ->
  Constructor ->
  [PyStmt] ->
  [PyStmt] ->
  Operations op s ->
  s ->
  [PyStmt] ->
  [Option] ->
  Imp.Definitions op ->
  m T.Text
compileProg :: forall (m :: * -> *) op s.
MonadFreshNames m =>
CompilerMode
-> [Char]
-> Constructor
-> [PyStmt]
-> [PyStmt]
-> Operations op s
-> s
-> [PyStmt]
-> [Option]
-> Definitions op
-> m Text
compileProg CompilerMode
mode [Char]
class_name Constructor
constructor [PyStmt]
imports [PyStmt]
defines Operations op s
ops s
userstate [PyStmt]
sync [Option]
options Definitions op
prog = do
  VNameSource
src <- forall (m :: * -> *). MonadFreshNames m => m VNameSource
getNameSource
  let prog' :: [PyStmt]
prog' = forall op s a.
Operations op s -> VNameSource -> s -> CompilerM op s a -> a
runCompilerM Operations op s
ops VNameSource
src s
userstate forall {s}. CompilerM op s [PyStmt]
compileProg'
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> Text
prettyText forall b c a. (b -> c) -> (a -> b) -> a -> c
. [PyStmt] -> PyProg
PyProg forall a b. (a -> b) -> a -> b
$
    [PyStmt]
imports
      forall a. [a] -> [a] -> [a]
++ [ [Char] -> Maybe [Char] -> PyStmt
Import [Char]
"argparse" forall a. Maybe a
Nothing,
           PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"sizes") forall a b. (a -> b) -> a -> b
$ [(PyExp, PyExp)] -> PyExp
Dict []
         ]
      forall a. [a] -> [a] -> [a]
++ [PyStmt]
defines
      forall a. [a] -> [a] -> [a]
++ [ Text -> PyStmt
Escape Text
valuesPy,
           Text -> PyStmt
Escape Text
memoryPy,
           Text -> PyStmt
Escape Text
panicPy,
           Text -> PyStmt
Escape Text
tuningPy,
           Text -> PyStmt
Escape Text
scalarPy,
           Text -> PyStmt
Escape Text
serverPy
         ]
      forall a. [a] -> [a] -> [a]
++ [PyStmt]
prog'
  where
    Imp.Definitions OpaqueTypes
_types Constants op
consts (Imp.Functions [(Name, Function op)]
funs) = Definitions op
prog
    compileProg' :: CompilerM op s [PyStmt]
compileProg' = forall op s a. Constants op -> CompilerM op s a -> CompilerM op s a
withConstantSubsts Constants op
consts forall a b. (a -> b) -> a -> b
$ do
      forall op s. Constants op -> CompilerM op s ()
compileConstants Constants op
consts

      [PyFunDef]
definitions <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall op s. (Name, Function op) -> CompilerM op s PyFunDef
compileFunc [(Name, Function op)]
funs
      [PyStmt]
at_inits <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall s. CompilerState s -> [PyStmt]
compInit

      let constructor' :: PyFunDef
constructor' = Constructor -> [PyStmt] -> PyFunDef
constructorToFunDef Constructor
constructor [PyStmt]
at_inits

      case CompilerMode
mode of
        CompilerMode
ToLibrary -> do
          ([PyFunDef]
entry_points, [(PyExp, PyExp)]
entry_point_types) <-
            forall a b. [(a, b)] -> ([a], [b])
unzip forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [Maybe a] -> [a]
catMaybes forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall op s.
[PyStmt]
-> ReturnTiming
-> (Name, Function op)
-> CompilerM op s (Maybe (PyFunDef, (PyExp, PyExp)))
compileEntryFun [PyStmt]
sync ReturnTiming
DoNotReturnTiming) [(Name, Function op)]
funs
          forall (f :: * -> *) a. Applicative f => a -> f a
pure
            [ PyClassDef -> PyStmt
ClassDef forall a b. (a -> b) -> a -> b
$
                [Char] -> [PyStmt] -> PyClassDef
Class [Char]
class_name forall a b. (a -> b) -> a -> b
$
                  PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"entry_points") ([(PyExp, PyExp)] -> PyExp
Dict [(PyExp, PyExp)]
entry_point_types)
                    forall a. a -> [a] -> [a]
: PyExp -> PyExp -> PyStmt
Assign
                      ([Char] -> PyExp
Var [Char]
"opaques")
                      ([(PyExp, PyExp)] -> PyExp
Dict forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip (forall a b. (a -> b) -> [a] -> [b]
map Text -> PyExp
String [Text]
opaque_names) (forall a b. (a -> b) -> [a] -> [b]
map [PyExp] -> PyExp
Tuple [[PyExp]]
opaque_payloads))
                    forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map PyFunDef -> PyStmt
FunDef (PyFunDef
constructor' forall a. a -> [a] -> [a]
: [PyFunDef]
definitions forall a. [a] -> [a] -> [a]
++ [PyFunDef]
entry_points)
            ]
        CompilerMode
ToServer -> do
          ([PyFunDef]
entry_points, [(PyExp, PyExp)]
entry_point_types) <-
            forall a b. [(a, b)] -> ([a], [b])
unzip forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [Maybe a] -> [a]
catMaybes forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall op s.
[PyStmt]
-> ReturnTiming
-> (Name, Function op)
-> CompilerM op s (Maybe (PyFunDef, (PyExp, PyExp)))
compileEntryFun [PyStmt]
sync ReturnTiming
ReturnTiming) [(Name, Function op)]
funs
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
            [PyStmt]
parse_options_server
              forall a. [a] -> [a] -> [a]
++ [ PyClassDef -> PyStmt
ClassDef
                     ( [Char] -> [PyStmt] -> PyClassDef
Class [Char]
class_name forall a b. (a -> b) -> a -> b
$
                         PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"entry_points") ([(PyExp, PyExp)] -> PyExp
Dict [(PyExp, PyExp)]
entry_point_types)
                           forall a. a -> [a] -> [a]
: PyExp -> PyExp -> PyStmt
Assign
                             ([Char] -> PyExp
Var [Char]
"opaques")
                             ([(PyExp, PyExp)] -> PyExp
Dict forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip (forall a b. (a -> b) -> [a] -> [b]
map Text -> PyExp
String [Text]
opaque_names) (forall a b. (a -> b) -> [a] -> [b]
map [PyExp] -> PyExp
Tuple [[PyExp]]
opaque_payloads))
                           forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map PyFunDef -> PyStmt
FunDef (PyFunDef
constructor' forall a. a -> [a] -> [a]
: [PyFunDef]
definitions forall a. [a] -> [a] -> [a]
++ [PyFunDef]
entry_points)
                     ),
                   PyExp -> PyExp -> PyStmt
Assign
                     ([Char] -> PyExp
Var [Char]
"server")
                     ([Char] -> [PyExp] -> PyExp
simpleCall [Char]
"Server" [[Char] -> [PyExp] -> PyExp
simpleCall [Char]
class_name []]),
                   PyExp -> PyStmt
Exp forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"server.run" []
                 ]
        CompilerMode
ToExecutable -> do
          let classinst :: PyStmt
classinst = PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"self") forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
class_name []
          ([PyFunDef]
entry_point_defs, [Text]
entry_point_names, [PyExp]
entry_points) <-
            forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [Maybe a] -> [a]
catMaybes forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall op s.
[PyStmt]
-> (Name, Function op)
-> CompilerM op s (Maybe (PyFunDef, Text, PyExp))
callEntryFun [PyStmt]
sync) [(Name, Function op)]
funs
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
            [PyStmt]
parse_options_executable
              forall a. [a] -> [a] -> [a]
++ PyClassDef -> PyStmt
ClassDef
                ( [Char] -> [PyStmt] -> PyClassDef
Class [Char]
class_name forall a b. (a -> b) -> a -> b
$
                    forall a b. (a -> b) -> [a] -> [b]
map PyFunDef -> PyStmt
FunDef forall a b. (a -> b) -> a -> b
$
                      PyFunDef
constructor' forall a. a -> [a] -> [a]
: [PyFunDef]
definitions
                )
              forall a. a -> [a] -> [a]
: PyStmt
classinst
              forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map PyFunDef -> PyStmt
FunDef [PyFunDef]
entry_point_defs
              forall a. [a] -> [a] -> [a]
++ [Text] -> [PyExp] -> [PyStmt]
selectEntryPoint [Text]
entry_point_names [PyExp]
entry_points

    parse_options_executable :: [PyStmt]
parse_options_executable =
      PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"runtime_file") PyExp
None
        forall a. a -> [a] -> [a]
: PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"do_warmup_run") (Bool -> PyExp
Bool Bool
False)
        forall a. a -> [a] -> [a]
: PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"num_runs") (Integer -> PyExp
Integer Integer
1)
        forall a. a -> [a] -> [a]
: PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"entry_point") (Text -> PyExp
String Text
"main")
        forall a. a -> [a] -> [a]
: PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"binary_output") (Bool -> PyExp
Bool Bool
False)
        forall a. a -> [a] -> [a]
: [Option] -> [PyStmt]
generateOptionParser ([Option]
executableOptions forall a. [a] -> [a] -> [a]
++ [Option]
options)

    parse_options_server :: [PyStmt]
parse_options_server =
      [Option] -> [PyStmt]
generateOptionParser ([Option]
standardOptions forall a. [a] -> [a] -> [a]
++ [Option]
options)

    ([Text]
opaque_names, [[PyExp]]
opaque_payloads) =
      forall a b. [(a, b)] -> ([a], [b])
unzip forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
M.toList forall a b. (a -> b) -> a -> b
$ forall a. Functions a -> Map Text [PyExp]
opaqueDefs forall a b. (a -> b) -> a -> b
$ forall a. Definitions a -> Functions a
Imp.defFuns Definitions op
prog

    selectEntryPoint :: [Text] -> [PyExp] -> [PyStmt]
selectEntryPoint [Text]
entry_point_names [PyExp]
entry_points =
      [ PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"entry_points") forall a b. (a -> b) -> a -> b
$
          [(PyExp, PyExp)] -> PyExp
Dict forall a b. (a -> b) -> a -> b
$
            forall a b. [a] -> [b] -> [(a, b)]
zip (forall a b. (a -> b) -> [a] -> [b]
map Text -> PyExp
String [Text]
entry_point_names) [PyExp]
entry_points,
        PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"entry_point_fun") forall a b. (a -> b) -> a -> b
$
          [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"entry_points.get" [[Char] -> PyExp
Var [Char]
"entry_point"],
        PyExp -> [PyStmt] -> [PyStmt] -> PyStmt
If
          ([Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"==" ([Char] -> PyExp
Var [Char]
"entry_point_fun") PyExp
None)
          [ PyExp -> PyStmt
Exp forall a b. (a -> b) -> a -> b
$
              [Char] -> [PyExp] -> PyExp
simpleCall
                [Char]
"sys.exit"
                [ PyExp -> [PyArg] -> PyExp
Call
                    ( PyExp -> [Char] -> PyExp
Field
                        (Text -> PyExp
String Text
"No entry point '{}'.  Select another with --entry point.  Options are:\n{}")
                        [Char]
"format"
                    )
                    [ PyExp -> PyArg
Arg forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp
Var [Char]
"entry_point",
                      PyExp -> PyArg
Arg forall a b. (a -> b) -> a -> b
$
                        PyExp -> [PyArg] -> PyExp
Call
                          (PyExp -> [Char] -> PyExp
Field (Text -> PyExp
String Text
"\n") [Char]
"join")
                          [PyExp -> PyArg
Arg forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"entry_points.keys" []]
                    ]
                ]
          ]
          [PyExp -> PyStmt
Exp forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"entry_point_fun" []]
      ]

withConstantSubsts :: Imp.Constants op -> CompilerM op s a -> CompilerM op s a
withConstantSubsts :: forall op s a. Constants op -> CompilerM op s a -> CompilerM op s a
withConstantSubsts (Imp.Constants [Param]
ps Code op
_) =
  forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ \CompilerEnv op s
env -> CompilerEnv op s
env {envVarExp :: Map [Char] PyExp
envVarExp = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Param -> Map [Char] PyExp
constExp [Param]
ps}
  where
    constExp :: Param -> Map [Char] PyExp
constExp Param
p =
      forall k a. k -> a -> Map k a
M.singleton
        (VName -> [Char]
compileName forall a b. (a -> b) -> a -> b
$ Param -> VName
Imp.paramName Param
p)
        (PyExp -> PyIdx -> PyExp
Index ([Char] -> PyExp
Var [Char]
"self.constants") forall a b. (a -> b) -> a -> b
$ PyExp -> PyIdx
IdxExp forall a b. (a -> b) -> a -> b
$ Text -> PyExp
String forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> Text
prettyText forall a b. (a -> b) -> a -> b
$ Param -> VName
Imp.paramName Param
p)

compileConstants :: Imp.Constants op -> CompilerM op s ()
compileConstants :: forall op s. Constants op -> CompilerM op s ()
compileConstants (Imp.Constants [Param]
_ Code op
init_consts) = do
  forall op s. PyStmt -> CompilerM op s ()
atInit forall a b. (a -> b) -> a -> b
$ PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"self.constants") forall a b. (a -> b) -> a -> b
$ [(PyExp, PyExp)] -> PyExp
Dict []
  forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ forall op s. PyStmt -> CompilerM op s ()
atInit forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall op s. CompilerM op s () -> CompilerM op s [PyStmt]
collect (forall op s. Code op -> CompilerM op s ()
compileCode Code op
init_consts)

compileFunc :: (Name, Imp.Function op) -> CompilerM op s PyFunDef
compileFunc :: forall op s. (Name, Function op) -> CompilerM op s PyFunDef
compileFunc (Name
fname, Imp.Function Maybe EntryPoint
_ [Param]
outputs [Param]
inputs Code op
body) = do
  [PyStmt]
body' <- forall op s. CompilerM op s () -> CompilerM op s [PyStmt]
collect forall a b. (a -> b) -> a -> b
$ forall op s. Code op -> CompilerM op s ()
compileCode Code op
body
  let inputs' :: [[Char]]
inputs' = forall a b. (a -> b) -> [a] -> [b]
map (VName -> [Char]
compileName forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param -> VName
Imp.paramName) [Param]
inputs
  let ret :: PyStmt
ret = PyExp -> PyStmt
Return forall a b. (a -> b) -> a -> b
$ [PyExp] -> PyExp
tupleOrSingle forall a b. (a -> b) -> a -> b
$ [Param] -> [PyExp]
compileOutput [Param]
outputs
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
    [Char] -> [[Char]] -> [PyStmt] -> PyFunDef
Def (Text -> [Char]
T.unpack forall a b. (a -> b) -> a -> b
$ Text -> Text
futharkFun forall a b. (a -> b) -> a -> b
$ Name -> Text
nameToText Name
fname) ([Char]
"self" forall a. a -> [a] -> [a]
: [[Char]]
inputs') forall a b. (a -> b) -> a -> b
$
      [PyStmt]
body' forall a. [a] -> [a] -> [a]
++ [PyStmt
ret]

tupleOrSingle :: [PyExp] -> PyExp
tupleOrSingle :: [PyExp] -> PyExp
tupleOrSingle [PyExp
e] = PyExp
e
tupleOrSingle [PyExp]
es = [PyExp] -> PyExp
Tuple [PyExp]
es

-- | A 'Call' where the function is a variable and every argument is a
-- simple 'Arg'.
simpleCall :: String -> [PyExp] -> PyExp
simpleCall :: [Char] -> [PyExp] -> PyExp
simpleCall [Char]
fname = PyExp -> [PyArg] -> PyExp
Call ([Char] -> PyExp
Var [Char]
fname) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map PyExp -> PyArg
Arg

compileName :: VName -> String
compileName :: VName -> [Char]
compileName = Text -> [Char]
T.unpack forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Text
zEncodeText forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> Text
prettyText

compileDim :: Imp.DimSize -> CompilerM op s PyExp
compileDim :: forall op s. DimSize -> CompilerM op s PyExp
compileDim (Imp.Constant PrimValue
v) = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ PrimValue -> PyExp
compilePrimValue PrimValue
v
compileDim (Imp.Var VName
v) = forall op s. VName -> CompilerM op s PyExp
compileVar VName
v

unpackDim :: PyExp -> Imp.DimSize -> Int32 -> CompilerM op s ()
unpackDim :: forall op s. PyExp -> DimSize -> Int32 -> CompilerM op s ()
unpackDim PyExp
arr_name (Imp.Constant PrimValue
c) Int32
i = do
  let shape_name :: PyExp
shape_name = PyExp -> [Char] -> PyExp
Field PyExp
arr_name [Char]
"shape"
  let constant_c :: PyExp
constant_c = PrimValue -> PyExp
compilePrimValue PrimValue
c
  let constant_i :: PyExp
constant_i = Integer -> PyExp
Integer forall a b. (a -> b) -> a -> b
$ forall a. Integral a => a -> Integer
toInteger Int32
i
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$
    PyExp -> PyExp -> PyStmt
Assert ([Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"==" PyExp
constant_c (PyExp -> PyIdx -> PyExp
Index PyExp
shape_name forall a b. (a -> b) -> a -> b
$ PyExp -> PyIdx
IdxExp PyExp
constant_i)) forall a b. (a -> b) -> a -> b
$
      Text -> PyExp
String Text
"Entry point arguments have invalid sizes."
unpackDim PyExp
arr_name (Imp.Var VName
var) Int32
i = do
  let shape_name :: PyExp
shape_name = PyExp -> [Char] -> PyExp
Field PyExp
arr_name [Char]
"shape"
      src :: PyExp
src = PyExp -> PyIdx -> PyExp
Index PyExp
shape_name forall a b. (a -> b) -> a -> b
$ PyExp -> PyIdx
IdxExp forall a b. (a -> b) -> a -> b
$ Integer -> PyExp
Integer forall a b. (a -> b) -> a -> b
$ forall a. Integral a => a -> Integer
toInteger Int32
i
  PyExp
var' <- forall op s. VName -> CompilerM op s PyExp
compileVar VName
var
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$
    PyExp -> [PyStmt] -> [PyStmt] -> PyStmt
If
      ([Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"==" PyExp
var' PyExp
None)
      [PyExp -> PyExp -> PyStmt
Assign PyExp
var' forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"np.int64" [PyExp
src]]
      [ PyExp -> PyExp -> PyStmt
Assert ([Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"==" PyExp
var' PyExp
src) forall a b. (a -> b) -> a -> b
$
          Text -> PyExp
String Text
"Error: entry point arguments have invalid sizes."
      ]

entryPointOutput :: Imp.ExternalValue -> CompilerM op s PyExp
entryPointOutput :: forall op s. ExternalValue -> CompilerM op s PyExp
entryPointOutput (Imp.OpaqueValue Name
desc [ValueDesc]
vs) =
  [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"opaque" forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text -> PyExp
String (forall a. Pretty a => a -> Text
prettyText Name
desc) :)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall op s. ExternalValue -> CompilerM op s PyExp
entryPointOutput forall b c a. (b -> c) -> (a -> b) -> a -> c
. ValueDesc -> ExternalValue
Imp.TransparentValue) [ValueDesc]
vs
entryPointOutput (Imp.TransparentValue (Imp.ScalarValue PrimType
bt Signedness
ept VName
name)) = do
  PyExp
name' <- forall op s. VName -> CompilerM op s PyExp
compileVar VName
name
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
tf [PyExp
name']
  where
    tf :: [Char]
tf = PrimType -> Signedness -> [Char]
compilePrimToExtNp PrimType
bt Signedness
ept
entryPointOutput (Imp.TransparentValue (Imp.ArrayValue VName
mem (Imp.Space [Char]
sid) PrimType
bt Signedness
ept [DimSize]
dims)) = do
  EntryOutput op s
pack_output <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall op s. CompilerEnv op s -> EntryOutput op s
envEntryOutput
  EntryOutput op s
pack_output VName
mem [Char]
sid PrimType
bt Signedness
ept [DimSize]
dims
entryPointOutput (Imp.TransparentValue (Imp.ArrayValue VName
mem Space
_ PrimType
bt Signedness
ept [DimSize]
dims)) = do
  PyExp
mem' <- PyExp -> [Char] -> PyExp
Cast forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall op s. VName -> CompilerM op s PyExp
compileVar VName
mem forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure (PrimType -> Signedness -> [Char]
compilePrimTypeExt PrimType
bt Signedness
ept)
  [PyExp]
dims' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall op s. DimSize -> CompilerM op s PyExp
compileDim [DimSize]
dims
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"createArray" [PyExp
mem', [PyExp] -> PyExp
Tuple [PyExp]
dims', [Char] -> PyExp
Var forall a b. (a -> b) -> a -> b
$ PrimType -> Signedness -> [Char]
compilePrimToExtNp PrimType
bt Signedness
ept]

badInput :: Int -> PyExp -> T.Text -> PyStmt
badInput :: Int -> PyExp -> Text -> PyStmt
badInput Int
i PyExp
e Text
t =
  PyExp -> PyStmt
Raise forall a b. (a -> b) -> a -> b
$
    [Char] -> [PyExp] -> PyExp
simpleCall
      [Char]
"TypeError"
      [ PyExp -> [PyArg] -> PyExp
Call
          (PyExp -> [Char] -> PyExp
Field (Text -> PyExp
String Text
err_msg) [Char]
"format")
          [PyExp -> PyArg
Arg (Text -> PyExp
String Text
t), PyExp -> PyArg
Arg forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"type" [PyExp
e], PyExp -> PyArg
Arg PyExp
e]
      ]
  where
    err_msg :: Text
err_msg =
      [Text] -> Text
T.unlines
        [ Text
"Argument #" forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText Int
i forall a. Semigroup a => a -> a -> a
<> Text
" has invalid value",
          Text
"Futhark type: {}",
          Text
"Argument has Python type {} and value: {}"
        ]

badInputType :: Int -> PyExp -> T.Text -> PyExp -> PyExp -> PyStmt
badInputType :: Int -> PyExp -> Text -> PyExp -> PyExp -> PyStmt
badInputType Int
i PyExp
e Text
t PyExp
de PyExp
dg =
  PyExp -> PyStmt
Raise forall a b. (a -> b) -> a -> b
$
    [Char] -> [PyExp] -> PyExp
simpleCall
      [Char]
"TypeError"
      [ PyExp -> [PyArg] -> PyExp
Call
          (PyExp -> [Char] -> PyExp
Field (Text -> PyExp
String Text
err_msg) [Char]
"format")
          [PyExp -> PyArg
Arg (Text -> PyExp
String Text
t), PyExp -> PyArg
Arg forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"type" [PyExp
e], PyExp -> PyArg
Arg PyExp
e, PyExp -> PyArg
Arg PyExp
de, PyExp -> PyArg
Arg PyExp
dg]
      ]
  where
    err_msg :: Text
err_msg =
      [Text] -> Text
T.unlines
        [ Text
"Argument #" forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText Int
i forall a. Semigroup a => a -> a -> a
<> Text
" has invalid value",
          Text
"Futhark type: {}",
          Text
"Argument has Python type {} and value: {}",
          Text
"Expected array with elements of dtype: {}",
          Text
"The array given has elements of dtype: {}"
        ]

badInputDim :: Int -> PyExp -> T.Text -> Int -> PyStmt
badInputDim :: Int -> PyExp -> Text -> Int -> PyStmt
badInputDim Int
i PyExp
e Text
typ Int
dimf =
  PyExp -> PyStmt
Raise forall a b. (a -> b) -> a -> b
$
    [Char] -> [PyExp] -> PyExp
simpleCall
      [Char]
"TypeError"
      [ PyExp -> [PyArg] -> PyExp
Call
          (PyExp -> [Char] -> PyExp
Field (Text -> PyExp
String Text
err_msg) [Char]
"format")
          [PyExp -> PyArg
Arg PyExp
eft, PyExp -> PyArg
Arg PyExp
aft]
      ]
  where
    eft :: PyExp
eft = Text -> PyExp
String (forall a. Monoid a => [a] -> a
mconcat (forall a. Int -> a -> [a]
replicate Int
dimf Text
"[]") forall a. Semigroup a => a -> a -> a
<> Text
typ)
    aft :: PyExp
aft = [Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"+" ([Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"*" (Text -> PyExp
String Text
"[]") (PyExp -> [Char] -> PyExp
Field PyExp
e [Char]
"ndim")) (Text -> PyExp
String Text
typ)
    err_msg :: Text
err_msg =
      [Text] -> Text
T.unlines
        [ Text
"Argument #" forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText Int
i forall a. Semigroup a => a -> a -> a
<> Text
" has invalid value",
          Text
"Dimensionality mismatch",
          Text
"Expected Futhark type: {}",
          Text
"Bad Python value passed",
          Text
"Actual Futhark type: {}"
        ]

declEntryPointInputSizes :: [Imp.ExternalValue] -> CompilerM op s ()
declEntryPointInputSizes :: forall op s. [ExternalValue] -> CompilerM op s ()
declEntryPointInputSizes = forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ forall {op} {s}. VName -> CompilerM op s ()
onSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ExternalValue -> [VName]
sizes
  where
    sizes :: ExternalValue -> [VName]
sizes (Imp.TransparentValue ValueDesc
v) = ValueDesc -> [VName]
valueSizes ValueDesc
v
    sizes (Imp.OpaqueValue Name
_ [ValueDesc]
vs) = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ValueDesc -> [VName]
valueSizes [ValueDesc]
vs
    valueSizes :: ValueDesc -> [VName]
valueSizes (Imp.ArrayValue VName
_ Space
_ PrimType
_ Signedness
_ [DimSize]
dims) = [DimSize] -> [VName]
subExpVars [DimSize]
dims
    valueSizes Imp.ScalarValue {} = []
    onSize :: VName -> CompilerM op s ()
onSize VName
v = forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$ PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var (VName -> [Char]
compileName VName
v)) PyExp
None

entryPointInput :: (Int, Imp.ExternalValue, PyExp) -> CompilerM op s ()
entryPointInput :: forall op s. (Int, ExternalValue, PyExp) -> CompilerM op s ()
entryPointInput (Int
i, Imp.OpaqueValue Name
desc [ValueDesc]
vs, PyExp
e) = do
  let type_is_ok :: PyExp
type_is_ok =
        [Char] -> PyExp -> PyExp -> PyExp
BinOp
          [Char]
"and"
          ([Char] -> [PyExp] -> PyExp
simpleCall [Char]
"isinstance" [PyExp
e, [Char] -> PyExp
Var [Char]
"opaque"])
          ([Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"==" (PyExp -> [Char] -> PyExp
Field PyExp
e [Char]
"desc") (Text -> PyExp
String (Name -> Text
nameToText Name
desc)))
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$ PyExp -> [PyStmt] -> [PyStmt] -> PyStmt
If ([Char] -> PyExp -> PyExp
UnOp [Char]
"not" PyExp
type_is_ok) [Int -> PyExp -> Text -> PyStmt
badInput Int
i PyExp
e (Name -> Text
nameToText Name
desc)] []
  forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ forall op s. (Int, ExternalValue, PyExp) -> CompilerM op s ()
entryPointInput forall a b. (a -> b) -> a -> b
$
    forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 (forall a. a -> [a]
repeat Int
i) (forall a b. (a -> b) -> [a] -> [b]
map ValueDesc -> ExternalValue
Imp.TransparentValue [ValueDesc]
vs) forall a b. (a -> b) -> a -> b
$
      forall a b. (a -> b) -> [a] -> [b]
map (PyExp -> PyIdx -> PyExp
Index (PyExp -> [Char] -> PyExp
Field PyExp
e [Char]
"data") forall b c a. (b -> c) -> (a -> b) -> a -> c
. PyExp -> PyIdx
IdxExp forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> PyExp
Integer) [Integer
0 ..]
entryPointInput (Int
i, Imp.TransparentValue (Imp.ScalarValue PrimType
bt Signedness
s VName
name), PyExp
e) = do
  PyExp
vname' <- forall op s. VName -> CompilerM op s PyExp
compileVar VName
name
  let -- HACK: A Numpy int64 will signal an OverflowError if we pass
      -- it a number bigger than 2**63.  This does not happen if we
      -- pass e.g. int8 a number bigger than 2**7.  As a workaround,
      -- we first go through the corresponding ctypes type, which does
      -- not have this problem.
      ctobject :: [Char]
ctobject = PrimType -> [Char]
compilePrimType PrimType
bt
      npobject :: [Char]
npobject = PrimType -> [Char]
compilePrimToNp PrimType
bt
      npcall :: PyExp
npcall =
        [Char] -> [PyExp] -> PyExp
simpleCall
          [Char]
npobject
          [ case PrimType
bt of
              IntType IntType
Int64 -> [Char] -> [PyExp] -> PyExp
simpleCall [Char]
ctobject [PyExp
e]
              PrimType
_ -> PyExp
e
          ]
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$
    [PyStmt] -> [PyExcept] -> PyStmt
Try
      [PyExp -> PyExp -> PyStmt
Assign PyExp
vname' PyExp
npcall]
      [ PyExp -> [PyStmt] -> PyExcept
Catch
          ([PyExp] -> PyExp
Tuple [[Char] -> PyExp
Var [Char]
"TypeError", [Char] -> PyExp
Var [Char]
"AssertionError"])
          [Int -> PyExp -> Text -> PyStmt
badInput Int
i PyExp
e forall a b. (a -> b) -> a -> b
$ Bool -> PrimType -> Text
prettySigned (Signedness
s forall a. Eq a => a -> a -> Bool
== Signedness
Imp.Unsigned) PrimType
bt]
      ]
entryPointInput (Int
i, Imp.TransparentValue (Imp.ArrayValue VName
mem (Imp.Space [Char]
sid) PrimType
bt Signedness
ept [DimSize]
dims), PyExp
e) = do
  EntryInput op s
unpack_input <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall op s. CompilerEnv op s -> EntryInput op s
envEntryInput
  PyExp
mem' <- forall op s. VName -> CompilerM op s PyExp
compileVar VName
mem
  [PyStmt]
unpack <- forall op s. CompilerM op s () -> CompilerM op s [PyStmt]
collect forall a b. (a -> b) -> a -> b
$ EntryInput op s
unpack_input PyExp
mem' [Char]
sid PrimType
bt Signedness
ept [DimSize]
dims PyExp
e
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$
    [PyStmt] -> [PyExcept] -> PyStmt
Try
      [PyStmt]
unpack
      [ PyExp -> [PyStmt] -> PyExcept
Catch
          ([PyExp] -> PyExp
Tuple [[Char] -> PyExp
Var [Char]
"TypeError", [Char] -> PyExp
Var [Char]
"AssertionError"])
          [ Int -> PyExp -> Text -> PyStmt
badInput Int
i PyExp
e forall a b. (a -> b) -> a -> b
$
              forall a. Monoid a => [a] -> a
mconcat (forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimSize]
dims) Text
"[]")
                forall a. Semigroup a => a -> a -> a
<> Bool -> PrimType -> Text
prettySigned (Signedness
ept forall a. Eq a => a -> a -> Bool
== Signedness
Imp.Unsigned) PrimType
bt
          ]
      ]
entryPointInput (Int
i, Imp.TransparentValue (Imp.ArrayValue VName
mem Space
_ PrimType
t Signedness
s [DimSize]
dims), PyExp
e) = do
  let type_is_wrong :: PyExp
type_is_wrong = [Char] -> PyExp -> PyExp
UnOp [Char]
"not" forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"in" ([Char] -> [PyExp] -> PyExp
simpleCall [Char]
"type" [PyExp
e]) forall a b. (a -> b) -> a -> b
$ [PyExp] -> PyExp
List [[Char] -> PyExp
Var [Char]
"np.ndarray"]
  let dtype_is_wrong :: PyExp
dtype_is_wrong = [Char] -> PyExp -> PyExp
UnOp [Char]
"not" forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"==" (PyExp -> [Char] -> PyExp
Field PyExp
e [Char]
"dtype") forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp
Var forall a b. (a -> b) -> a -> b
$ PrimType -> Signedness -> [Char]
compilePrimToExtNp PrimType
t Signedness
s
  let dim_is_wrong :: PyExp
dim_is_wrong = [Char] -> PyExp -> PyExp
UnOp [Char]
"not" forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"==" (PyExp -> [Char] -> PyExp
Field PyExp
e [Char]
"ndim") forall a b. (a -> b) -> a -> b
$ Integer -> PyExp
Integer forall a b. (a -> b) -> a -> b
$ forall a. Integral a => a -> Integer
toInteger forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimSize]
dims
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$
    PyExp -> [PyStmt] -> [PyStmt] -> PyStmt
If
      PyExp
type_is_wrong
      [ Int -> PyExp -> Text -> PyStmt
badInput Int
i PyExp
e forall a b. (a -> b) -> a -> b
$
          forall a. Monoid a => [a] -> a
mconcat (forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimSize]
dims) Text
"[]")
            forall a. Semigroup a => a -> a -> a
<> Bool -> PrimType -> Text
prettySigned (Signedness
s forall a. Eq a => a -> a -> Bool
== Signedness
Imp.Unsigned) PrimType
t
      ]
      []
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$
    PyExp -> [PyStmt] -> [PyStmt] -> PyStmt
If
      PyExp
dtype_is_wrong
      [ Int -> PyExp -> Text -> PyExp -> PyExp -> PyStmt
badInputType
          Int
i
          PyExp
e
          (forall a. Monoid a => [a] -> a
mconcat (forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimSize]
dims) Text
"[]") forall a. Semigroup a => a -> a -> a
<> Bool -> PrimType -> Text
prettySigned (Signedness
s forall a. Eq a => a -> a -> Bool
== Signedness
Imp.Unsigned) PrimType
t)
          ([Char] -> [PyExp] -> PyExp
simpleCall [Char]
"np.dtype" [[Char] -> PyExp
Var (PrimType -> Signedness -> [Char]
compilePrimToExtNp PrimType
t Signedness
s)])
          (PyExp -> [Char] -> PyExp
Field PyExp
e [Char]
"dtype")
      ]
      []
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$
    PyExp -> [PyStmt] -> [PyStmt] -> PyStmt
If
      PyExp
dim_is_wrong
      [Int -> PyExp -> Text -> Int -> PyStmt
badInputDim Int
i PyExp
e (Bool -> PrimType -> Text
prettySigned (Signedness
s forall a. Eq a => a -> a -> Bool
== Signedness
Imp.Unsigned) PrimType
t) (forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimSize]
dims)]
      []

  forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ (forall op s. PyExp -> DimSize -> Int32 -> CompilerM op s ()
unpackDim PyExp
e) [DimSize]
dims [Int32
0 ..]
  PyExp
dest <- forall op s. VName -> CompilerM op s PyExp
compileVar VName
mem
  let unwrap_call :: PyExp
unwrap_call = [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"unwrapArray" [PyExp
e]

  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$ PyExp -> PyExp -> PyStmt
Assign PyExp
dest PyExp
unwrap_call

extValueDescName :: Imp.ExternalValue -> T.Text
extValueDescName :: ExternalValue -> Text
extValueDescName (Imp.TransparentValue ValueDesc
v) = Text -> Text
extName forall a b. (a -> b) -> a -> b
$ [Char] -> Text
T.pack forall a b. (a -> b) -> a -> b
$ VName -> [Char]
compileName forall a b. (a -> b) -> a -> b
$ ValueDesc -> VName
valueDescVName ValueDesc
v
extValueDescName (Imp.OpaqueValue Name
desc []) = Text -> Text
extName forall a b. (a -> b) -> a -> b
$ Text -> Text
zEncodeText forall a b. (a -> b) -> a -> b
$ Name -> Text
nameToText Name
desc
extValueDescName (Imp.OpaqueValue Name
desc (ValueDesc
v : [ValueDesc]
_)) =
  Text -> Text
extName forall a b. (a -> b) -> a -> b
$ Text -> Text
zEncodeText (Name -> Text
nameToText Name
desc) forall a. Semigroup a => a -> a -> a
<> Text
"_" forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText (VName -> Int
baseTag (ValueDesc -> VName
valueDescVName ValueDesc
v))

extName :: T.Text -> T.Text
extName :: Text -> Text
extName = (forall a. Semigroup a => a -> a -> a
<> Text
"_ext")

valueDescVName :: Imp.ValueDesc -> VName
valueDescVName :: ValueDesc -> VName
valueDescVName (Imp.ScalarValue PrimType
_ Signedness
_ VName
vname) = VName
vname
valueDescVName (Imp.ArrayValue VName
vname Space
_ PrimType
_ Signedness
_ [DimSize]
_) = VName
vname

-- Key into the FUTHARK_PRIMTYPES dict.
readTypeEnum :: PrimType -> Imp.Signedness -> T.Text
readTypeEnum :: PrimType -> Signedness -> Text
readTypeEnum (IntType IntType
Int8) Signedness
Imp.Unsigned = Text
"u8"
readTypeEnum (IntType IntType
Int16) Signedness
Imp.Unsigned = Text
"u16"
readTypeEnum (IntType IntType
Int32) Signedness
Imp.Unsigned = Text
"u32"
readTypeEnum (IntType IntType
Int64) Signedness
Imp.Unsigned = Text
"u64"
readTypeEnum (IntType IntType
Int8) Signedness
Imp.Signed = Text
"i8"
readTypeEnum (IntType IntType
Int16) Signedness
Imp.Signed = Text
"i16"
readTypeEnum (IntType IntType
Int32) Signedness
Imp.Signed = Text
"i32"
readTypeEnum (IntType IntType
Int64) Signedness
Imp.Signed = Text
"i64"
readTypeEnum (FloatType FloatType
Float16) Signedness
_ = Text
"f16"
readTypeEnum (FloatType FloatType
Float32) Signedness
_ = Text
"f32"
readTypeEnum (FloatType FloatType
Float64) Signedness
_ = Text
"f64"
readTypeEnum PrimType
Imp.Bool Signedness
_ = Text
"bool"
readTypeEnum PrimType
Unit Signedness
_ = Text
"bool"

readInput :: Imp.ExternalValue -> PyStmt
readInput :: ExternalValue -> PyStmt
readInput (Imp.OpaqueValue Name
desc [ValueDesc]
_) =
  PyExp -> PyStmt
Raise forall a b. (a -> b) -> a -> b
$
    [Char] -> [PyExp] -> PyExp
simpleCall
      [Char]
"Exception"
      [Text -> PyExp
String forall a b. (a -> b) -> a -> b
$ Text
"Cannot read argument of type " forall a. Semigroup a => a -> a -> a
<> Name -> Text
nameToText Name
desc forall a. Semigroup a => a -> a -> a
<> Text
"."]
readInput decl :: ExternalValue
decl@(Imp.TransparentValue (Imp.ScalarValue PrimType
bt Signedness
ept VName
_)) =
  let type_name :: Text
type_name = PrimType -> Signedness -> Text
readTypeEnum PrimType
bt Signedness
ept
   in PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var forall a b. (a -> b) -> a -> b
$ Text -> [Char]
T.unpack forall a b. (a -> b) -> a -> b
$ ExternalValue -> Text
extValueDescName ExternalValue
decl) forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"read_value" [Text -> PyExp
String Text
type_name]
readInput decl :: ExternalValue
decl@(Imp.TransparentValue (Imp.ArrayValue VName
_ Space
_ PrimType
bt Signedness
ept [DimSize]
dims)) =
  let type_name :: Text
type_name = PrimType -> Signedness -> Text
readTypeEnum PrimType
bt Signedness
ept
   in PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var forall a b. (a -> b) -> a -> b
$ Text -> [Char]
T.unpack forall a b. (a -> b) -> a -> b
$ ExternalValue -> Text
extValueDescName ExternalValue
decl) forall a b. (a -> b) -> a -> b
$
        [Char] -> [PyExp] -> PyExp
simpleCall
          [Char]
"read_value"
          [Text -> PyExp
String forall a b. (a -> b) -> a -> b
$ forall a. Monoid a => [a] -> a
mconcat (forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimSize]
dims) Text
"[]") forall a. Semigroup a => a -> a -> a
<> Text
type_name]

printValue :: [(Imp.ExternalValue, PyExp)] -> CompilerM op s [PyStmt]
printValue :: forall op s. [(ExternalValue, PyExp)] -> CompilerM op s [PyStmt]
printValue = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall {f :: * -> *}.
Applicative f =>
ExternalValue -> PyExp -> f [PyStmt]
printValue')
  where
    -- We copy non-host arrays to the host before printing.  This is
    -- done in a hacky way - we assume the value has a .get()-method
    -- that returns an equivalent Numpy array.  This works for PyOpenCL,
    -- but we will probably need yet another plugin mechanism here in
    -- the future.
    printValue' :: ExternalValue -> PyExp -> f [PyStmt]
printValue' (Imp.OpaqueValue Name
desc [ValueDesc]
_) PyExp
_ =
      forall (f :: * -> *) a. Applicative f => a -> f a
pure
        [ PyExp -> PyStmt
Exp forall a b. (a -> b) -> a -> b
$
            [Char] -> [PyExp] -> PyExp
simpleCall
              [Char]
"sys.stdout.write"
              [Text -> PyExp
String forall a b. (a -> b) -> a -> b
$ Text
"#<opaque " forall a. Semigroup a => a -> a -> a
<> Name -> Text
nameToText Name
desc forall a. Semigroup a => a -> a -> a
<> Text
">"]
        ]
    printValue' (Imp.TransparentValue (Imp.ArrayValue VName
mem (Space [Char]
_) PrimType
bt Signedness
ept [DimSize]
shape)) PyExp
e =
      ExternalValue -> PyExp -> f [PyStmt]
printValue' (ValueDesc -> ExternalValue
Imp.TransparentValue (VName -> Space -> PrimType -> Signedness -> [DimSize] -> ValueDesc
Imp.ArrayValue VName
mem Space
DefaultSpace PrimType
bt Signedness
ept [DimSize]
shape)) forall a b. (a -> b) -> a -> b
$
        [Char] -> [PyExp] -> PyExp
simpleCall (forall a. Pretty a => a -> [Char]
prettyString PyExp
e forall a. [a] -> [a] -> [a]
++ [Char]
".get") []
    printValue' (Imp.TransparentValue ValueDesc
_) PyExp
e =
      forall (f :: * -> *) a. Applicative f => a -> f a
pure
        [ PyExp -> PyStmt
Exp forall a b. (a -> b) -> a -> b
$
            PyExp -> [PyArg] -> PyExp
Call
              ([Char] -> PyExp
Var [Char]
"write_value")
              [ PyExp -> PyArg
Arg PyExp
e,
                [Char] -> PyExp -> PyArg
ArgKeyword [Char]
"binary" ([Char] -> PyExp
Var [Char]
"binary_output")
              ],
          PyExp -> PyStmt
Exp forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"sys.stdout.write" [Text -> PyExp
String Text
"\n"]
        ]

prepareEntry ::
  Imp.EntryPoint ->
  (Name, Imp.Function op) ->
  CompilerM
    op
    s
    ( [String],
      [PyStmt],
      [PyStmt],
      [PyStmt],
      [(Imp.ExternalValue, PyExp)]
    )
prepareEntry :: forall op s.
EntryPoint
-> (Name, Function op)
-> CompilerM
     op
     s
     ([[Char]], [PyStmt], [PyStmt], [PyStmt], [(ExternalValue, PyExp)])
prepareEntry (Imp.EntryPoint Name
_ [(Uniqueness, ExternalValue)]
results [((Name, Uniqueness), ExternalValue)]
args) (Name
fname, Imp.Function Maybe EntryPoint
_ [Param]
outputs [Param]
inputs Code op
_) = do
  let output_paramNames :: [[Char]]
output_paramNames = forall a b. (a -> b) -> [a] -> [b]
map (VName -> [Char]
compileName forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param -> VName
Imp.paramName) [Param]
outputs
      funTuple :: PyExp
funTuple = [PyExp] -> PyExp
tupleOrSingle forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Char] -> PyExp
Var [[Char]]
output_paramNames

  [PyStmt]
prepareIn <- forall op s. CompilerM op s () -> CompilerM op s [PyStmt]
collect forall a b. (a -> b) -> a -> b
$ do
    forall op s. [ExternalValue] -> CompilerM op s ()
declEntryPointInputSizes forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [((Name, Uniqueness), ExternalValue)]
args
    forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ forall op s. (Int, ExternalValue, PyExp) -> CompilerM op s ()
entryPointInput forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Int
0 ..] (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [((Name, Uniqueness), ExternalValue)]
args) forall a b. (a -> b) -> a -> b
$
      forall a b. (a -> b) -> [a] -> [b]
map ([Char] -> PyExp
Var forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> [Char]
T.unpack forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExternalValue -> Text
extValueDescName forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [((Name, Uniqueness), ExternalValue)]
args
  ([PyExp]
res, [PyStmt]
prepareOut) <- forall op s a. CompilerM op s a -> CompilerM op s (a, [PyStmt])
collect' forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall op s. ExternalValue -> CompilerM op s PyExp
entryPointOutput forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(Uniqueness, ExternalValue)]
results

  let argexps_lib :: [[Char]]
argexps_lib = forall a b. (a -> b) -> [a] -> [b]
map (VName -> [Char]
compileName forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param -> VName
Imp.paramName) [Param]
inputs
      fname' :: Text
fname' = Text
"self." forall a. Semigroup a => a -> a -> a
<> Text -> Text
futharkFun (Name -> Text
nameToText Name
fname)

      -- We ignore overflow errors and the like for executable entry
      -- points.  These are (somewhat) well-defined in Futhark.
      ignore :: [Char] -> PyArg
ignore [Char]
s = [Char] -> PyExp -> PyArg
ArgKeyword [Char]
s forall a b. (a -> b) -> a -> b
$ Text -> PyExp
String Text
"ignore"
      errstate :: PyExp
errstate = PyExp -> [PyArg] -> PyExp
Call ([Char] -> PyExp
Var [Char]
"np.errstate") forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map [Char] -> PyArg
ignore [[Char]
"divide", [Char]
"over", [Char]
"under", [Char]
"invalid"]

      call :: [[Char]] -> [PyStmt]
call [[Char]]
argexps =
        [ PyExp -> [PyStmt] -> PyStmt
With
            PyExp
errstate
            [PyExp -> PyExp -> PyStmt
Assign PyExp
funTuple forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall (Text -> [Char]
T.unpack Text
fname') (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Char] -> PyExp
Var [[Char]]
argexps)]
        ]

  forall (f :: * -> *) a. Applicative f => a -> f a
pure
    ( forall a b. (a -> b) -> [a] -> [b]
map (Text -> [Char]
T.unpack forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExternalValue -> Text
extValueDescName forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [((Name, Uniqueness), ExternalValue)]
args,
      [PyStmt]
prepareIn,
      [[Char]] -> [PyStmt]
call [[Char]]
argexps_lib,
      [PyStmt]
prepareOut,
      forall a b. [a] -> [b] -> [(a, b)]
zip (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(Uniqueness, ExternalValue)]
results) [PyExp]
res
    )

copyMemoryDefaultSpace ::
  PyExp ->
  PyExp ->
  PyExp ->
  PyExp ->
  PyExp ->
  CompilerM op s ()
copyMemoryDefaultSpace :: forall op s.
PyExp -> PyExp -> PyExp -> PyExp -> PyExp -> CompilerM op s ()
copyMemoryDefaultSpace PyExp
destmem PyExp
destidx PyExp
srcmem PyExp
srcidx PyExp
nbytes = do
  let offset_call1 :: PyExp
offset_call1 =
        [Char] -> [PyExp] -> PyExp
simpleCall
          [Char]
"addressOffset"
          [PyExp
destmem, PyExp
destidx, [Char] -> PyExp
Var [Char]
"ct.c_byte"]
  let offset_call2 :: PyExp
offset_call2 =
        [Char] -> [PyExp] -> PyExp
simpleCall
          [Char]
"addressOffset"
          [PyExp
srcmem, PyExp
srcidx, [Char] -> PyExp
Var [Char]
"ct.c_byte"]
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$ PyExp -> PyStmt
Exp forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"ct.memmove" [PyExp
offset_call1, PyExp
offset_call2, PyExp
nbytes]

data ReturnTiming = ReturnTiming | DoNotReturnTiming

compileEntryFun ::
  [PyStmt] ->
  ReturnTiming ->
  (Name, Imp.Function op) ->
  CompilerM op s (Maybe (PyFunDef, (PyExp, PyExp)))
compileEntryFun :: forall op s.
[PyStmt]
-> ReturnTiming
-> (Name, Function op)
-> CompilerM op s (Maybe (PyFunDef, (PyExp, PyExp)))
compileEntryFun [PyStmt]
sync ReturnTiming
timing (Name, Function op)
fun
  | Just EntryPoint
entry <- forall a. FunctionT a -> Maybe EntryPoint
Imp.functionEntry forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> b
snd (Name, Function op)
fun = do
      let ename :: Name
ename = EntryPoint -> Name
Imp.entryPointName EntryPoint
entry
      ([[Char]]
params, [PyStmt]
prepareIn, [PyStmt]
body_lib, [PyStmt]
prepareOut, [(ExternalValue, PyExp)]
res) <- forall op s.
EntryPoint
-> (Name, Function op)
-> CompilerM
     op
     s
     ([[Char]], [PyStmt], [PyStmt], [PyStmt], [(ExternalValue, PyExp)])
prepareEntry EntryPoint
entry (Name, Function op)
fun
      let ([PyStmt]
maybe_sync, PyStmt
ret) =
            case ReturnTiming
timing of
              ReturnTiming
DoNotReturnTiming ->
                ( [],
                  PyExp -> PyStmt
Return forall a b. (a -> b) -> a -> b
$ [PyExp] -> PyExp
tupleOrSingle forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(ExternalValue, PyExp)]
res
                )
              ReturnTiming
ReturnTiming ->
                ( [PyStmt]
sync,
                  PyExp -> PyStmt
Return forall a b. (a -> b) -> a -> b
$
                    [PyExp] -> PyExp
Tuple
                      [ [Char] -> PyExp
Var [Char]
"runtime",
                        [PyExp] -> PyExp
tupleOrSingle forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(ExternalValue, PyExp)]
res
                      ]
                )
          ([Text]
pts, [Text]
rts) = EntryPoint -> ([Text], [Text])
entryTypes EntryPoint
entry

          do_run :: [PyStmt]
do_run =
            PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"time_start") ([Char] -> [PyExp] -> PyExp
simpleCall [Char]
"time.time" [])
              forall a. a -> [a] -> [a]
: [PyStmt]
body_lib
              forall a. [a] -> [a] -> [a]
++ [PyStmt]
maybe_sync
              forall a. [a] -> [a] -> [a]
++ [ PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"runtime") forall a b. (a -> b) -> a -> b
$
                     [Char] -> PyExp -> PyExp -> PyExp
BinOp
                       [Char]
"-"
                       (PyExp -> PyExp
toMicroseconds ([Char] -> [PyExp] -> PyExp
simpleCall [Char]
"time.time" []))
                       (PyExp -> PyExp
toMicroseconds ([Char] -> PyExp
Var [Char]
"time_start"))
                 ]

      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
        forall a. a -> Maybe a
Just
          ( [Char] -> [[Char]] -> [PyStmt] -> PyFunDef
Def (Text -> [Char]
T.unpack (Name -> Text
escapeName Name
ename)) ([Char]
"self" forall a. a -> [a] -> [a]
: [[Char]]
params) forall a b. (a -> b) -> a -> b
$
              [PyStmt]
prepareIn forall a. [a] -> [a] -> [a]
++ [PyStmt]
do_run forall a. [a] -> [a] -> [a]
++ [PyStmt]
prepareOut forall a. [a] -> [a] -> [a]
++ [PyStmt]
sync forall a. [a] -> [a] -> [a]
++ [PyStmt
ret],
            ( Text -> PyExp
String (Name -> Text
nameToText Name
ename),
              [PyExp] -> PyExp
Tuple
                [ Text -> PyExp
String (Name -> Text
escapeName Name
ename),
                  [PyExp] -> PyExp
List (forall a b. (a -> b) -> [a] -> [b]
map Text -> PyExp
String [Text]
pts),
                  [PyExp] -> PyExp
List (forall a b. (a -> b) -> [a] -> [b]
map Text -> PyExp
String [Text]
rts)
                ]
            )
          )
  | Bool
otherwise = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing

entryTypes :: Imp.EntryPoint -> ([T.Text], [T.Text])
entryTypes :: EntryPoint -> ([Text], [Text])
entryTypes (Imp.EntryPoint Name
_ [(Uniqueness, ExternalValue)]
res [((Name, Uniqueness), ExternalValue)]
args) =
  (forall a b. (a -> b) -> [a] -> [b]
map forall {a} {a}. Pretty a => ((a, a), ExternalValue) -> Text
descArg [((Name, Uniqueness), ExternalValue)]
args, forall a b. (a -> b) -> [a] -> [b]
map forall {a}. Pretty a => (a, ExternalValue) -> Text
desc [(Uniqueness, ExternalValue)]
res)
  where
    descArg :: ((a, a), ExternalValue) -> Text
descArg ((a
_, a
u), ExternalValue
d) = forall {a}. Pretty a => (a, ExternalValue) -> Text
desc (a
u, ExternalValue
d)
    desc :: (a, ExternalValue) -> Text
desc (a
u, Imp.OpaqueValue Name
d [ValueDesc]
_) = forall a. Pretty a => a -> Text
prettyText a
u forall a. Semigroup a => a -> a -> a
<> Name -> Text
nameToText Name
d
    desc (a
u, Imp.TransparentValue (Imp.ScalarValue PrimType
pt Signedness
s VName
_)) = forall a. Pretty a => a -> Text
prettyText a
u forall a. Semigroup a => a -> a -> a
<> PrimType -> Signedness -> Text
readTypeEnum PrimType
pt Signedness
s
    desc (a
u, Imp.TransparentValue (Imp.ArrayValue VName
_ Space
_ PrimType
pt Signedness
s [DimSize]
dims)) =
      forall a. Pretty a => a -> Text
prettyText a
u forall a. Semigroup a => a -> a -> a
<> forall a. Monoid a => [a] -> a
mconcat (forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimSize]
dims) Text
"[]") forall a. Semigroup a => a -> a -> a
<> PrimType -> Signedness -> Text
readTypeEnum PrimType
pt Signedness
s

callEntryFun ::
  [PyStmt] ->
  (Name, Imp.Function op) ->
  CompilerM op s (Maybe (PyFunDef, T.Text, PyExp))
callEntryFun :: forall op s.
[PyStmt]
-> (Name, Function op)
-> CompilerM op s (Maybe (PyFunDef, Text, PyExp))
callEntryFun [PyStmt]
_ (Name
_, Imp.Function Maybe EntryPoint
Nothing [Param]
_ [Param]
_ Code op
_) = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
callEntryFun [PyStmt]
pre_timing fun :: (Name, Function op)
fun@(Name
fname, Imp.Function (Just EntryPoint
entry) [Param]
_ [Param]
_ Code op
_) = do
  let Imp.EntryPoint Name
ename [(Uniqueness, ExternalValue)]
_ [((Name, Uniqueness), ExternalValue)]
decl_args = EntryPoint
entry
  ([[Char]]
_, [PyStmt]
prepare_in, [PyStmt]
body_bin, [PyStmt]
_, [(ExternalValue, PyExp)]
res) <- forall op s.
EntryPoint
-> (Name, Function op)
-> CompilerM
     op
     s
     ([[Char]], [PyStmt], [PyStmt], [PyStmt], [(ExternalValue, PyExp)])
prepareEntry EntryPoint
entry (Name, Function op)
fun

  let str_input :: [PyStmt]
str_input = forall a b. (a -> b) -> [a] -> [b]
map (ExternalValue -> PyStmt
readInput forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [((Name, Uniqueness), ExternalValue)]
decl_args
      end_of_input :: [PyStmt]
end_of_input = [PyExp -> PyStmt
Exp forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"end_of_input" [Text -> PyExp
String forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> Text
prettyText Name
fname]]

      exitcall :: [PyStmt]
exitcall = [PyExp -> PyStmt
Exp forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"sys.exit" [PyExp -> [Char] -> PyExp
Field (Text -> PyExp
String Text
"Assertion.{} failed") [Char]
"format(e)"]]
      except' :: PyExcept
except' = PyExp -> [PyStmt] -> PyExcept
Catch ([Char] -> PyExp
Var [Char]
"AssertionError") [PyStmt]
exitcall
      do_run :: [PyStmt]
do_run = [PyStmt]
body_bin forall a. [a] -> [a] -> [a]
++ [PyStmt]
pre_timing
      ([PyStmt]
do_run_with_timing, PyStmt
close_runtime_file) = [PyStmt] -> ([PyStmt], PyStmt)
addTiming [PyStmt]
do_run

      do_warmup_run :: PyStmt
do_warmup_run =
        PyExp -> [PyStmt] -> [PyStmt] -> PyStmt
If ([Char] -> PyExp
Var [Char]
"do_warmup_run") [PyStmt]
do_run []

      do_num_runs :: PyStmt
do_num_runs =
        [Char] -> PyExp -> [PyStmt] -> PyStmt
For
          [Char]
"i"
          ([Char] -> [PyExp] -> PyExp
simpleCall [Char]
"range" [[Char] -> [PyExp] -> PyExp
simpleCall [Char]
"int" [[Char] -> PyExp
Var [Char]
"num_runs"]])
          [PyStmt]
do_run_with_timing

  [PyStmt]
str_output <- forall op s. [(ExternalValue, PyExp)] -> CompilerM op s [PyStmt]
printValue [(ExternalValue, PyExp)]
res

  let fname' :: [Char]
fname' = [Char]
"entry_" forall a. [a] -> [a] -> [a]
++ Text -> [Char]
T.unpack (Name -> Text
escapeName Name
fname)

  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
    forall a. a -> Maybe a
Just
      ( [Char] -> [[Char]] -> [PyStmt] -> PyFunDef
Def [Char]
fname' [] forall a b. (a -> b) -> a -> b
$
          [PyStmt]
str_input
            forall a. [a] -> [a] -> [a]
++ [PyStmt]
end_of_input
            forall a. [a] -> [a] -> [a]
++ [PyStmt]
prepare_in
            forall a. [a] -> [a] -> [a]
++ [[PyStmt] -> [PyExcept] -> PyStmt
Try [PyStmt
do_warmup_run, PyStmt
do_num_runs] [PyExcept
except']]
            forall a. [a] -> [a] -> [a]
++ [PyStmt
close_runtime_file]
            forall a. [a] -> [a] -> [a]
++ [PyStmt]
str_output,
        Name -> Text
nameToText Name
ename,
        [Char] -> PyExp
Var [Char]
fname'
      )

addTiming :: [PyStmt] -> ([PyStmt], PyStmt)
addTiming :: [PyStmt] -> ([PyStmt], PyStmt)
addTiming [PyStmt]
statements =
  ( [PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"time_start") forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"time.time" []]
      forall a. [a] -> [a] -> [a]
++ [PyStmt]
statements
      forall a. [a] -> [a] -> [a]
++ [ PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"time_end") forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"time.time" [],
           PyExp -> [PyStmt] -> [PyStmt] -> PyStmt
If ([Char] -> PyExp
Var [Char]
"runtime_file") [PyStmt]
print_runtime []
         ],
    PyExp -> [PyStmt] -> [PyStmt] -> PyStmt
If ([Char] -> PyExp
Var [Char]
"runtime_file") [PyExp -> PyStmt
Exp forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"runtime_file.close" []] []
  )
  where
    print_runtime :: [PyStmt]
print_runtime =
      [ PyExp -> PyStmt
Exp forall a b. (a -> b) -> a -> b
$
          [Char] -> [PyExp] -> PyExp
simpleCall
            [Char]
"runtime_file.write"
            [ [Char] -> [PyExp] -> PyExp
simpleCall
                [Char]
"str"
                [ [Char] -> PyExp -> PyExp -> PyExp
BinOp
                    [Char]
"-"
                    (PyExp -> PyExp
toMicroseconds ([Char] -> PyExp
Var [Char]
"time_end"))
                    (PyExp -> PyExp
toMicroseconds ([Char] -> PyExp
Var [Char]
"time_start"))
                ]
            ],
        PyExp -> PyStmt
Exp forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"runtime_file.write" [Text -> PyExp
String Text
"\n"],
        PyExp -> PyStmt
Exp forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"runtime_file.flush" []
      ]

toMicroseconds :: PyExp -> PyExp
toMicroseconds :: PyExp -> PyExp
toMicroseconds PyExp
x =
  [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"int" [[Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"*" PyExp
x forall a b. (a -> b) -> a -> b
$ Integer -> PyExp
Integer Integer
1000000]

compileUnOp :: Imp.UnOp -> String
compileUnOp :: UnOp -> [Char]
compileUnOp UnOp
op =
  case UnOp
op of
    UnOp
Not -> [Char]
"not"
    Complement {} -> [Char]
"~"
    Abs {} -> [Char]
"abs"
    FAbs {} -> [Char]
"abs"
    SSignum {} -> [Char]
"ssignum"
    USignum {} -> [Char]
"usignum"
    FSignum {} -> [Char]
"np.sign"

compileBinOpLike ::
  Monad m =>
  (v -> m PyExp) ->
  Imp.PrimExp v ->
  Imp.PrimExp v ->
  m (PyExp, PyExp, String -> m PyExp)
compileBinOpLike :: forall (m :: * -> *) v.
Monad m =>
(v -> m PyExp)
-> PrimExp v -> PrimExp v -> m (PyExp, PyExp, [Char] -> m PyExp)
compileBinOpLike v -> m PyExp
f PrimExp v
x PrimExp v
y = do
  PyExp
x' <- forall (m :: * -> *) v.
Monad m =>
(v -> m PyExp) -> PrimExp v -> m PyExp
compilePrimExp v -> m PyExp
f PrimExp v
x
  PyExp
y' <- forall (m :: * -> *) v.
Monad m =>
(v -> m PyExp) -> PrimExp v -> m PyExp
compilePrimExp v -> m PyExp
f PrimExp v
y
  let simple :: [Char] -> f PyExp
simple [Char]
s = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
s PyExp
x' PyExp
y'
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (PyExp
x', PyExp
y', forall {f :: * -> *}. Applicative f => [Char] -> f PyExp
simple)

-- | The ctypes type corresponding to a 'PrimType'.
compilePrimType :: PrimType -> String
compilePrimType :: PrimType -> [Char]
compilePrimType PrimType
t =
  case PrimType
t of
    IntType IntType
Int8 -> [Char]
"ct.c_int8"
    IntType IntType
Int16 -> [Char]
"ct.c_int16"
    IntType IntType
Int32 -> [Char]
"ct.c_int32"
    IntType IntType
Int64 -> [Char]
"ct.c_int64"
    FloatType FloatType
Float16 -> [Char]
"ct.c_uint16"
    FloatType FloatType
Float32 -> [Char]
"ct.c_float"
    FloatType FloatType
Float64 -> [Char]
"ct.c_double"
    PrimType
Imp.Bool -> [Char]
"ct.c_bool"
    PrimType
Unit -> [Char]
"ct.c_bool"

-- | The ctypes type corresponding to a 'PrimType', taking sign into account.
compilePrimTypeExt :: PrimType -> Imp.Signedness -> String
compilePrimTypeExt :: PrimType -> Signedness -> [Char]
compilePrimTypeExt PrimType
t Signedness
ept =
  case (PrimType
t, Signedness
ept) of
    (IntType IntType
Int8, Signedness
Imp.Unsigned) -> [Char]
"ct.c_uint8"
    (IntType IntType
Int16, Signedness
Imp.Unsigned) -> [Char]
"ct.c_uint16"
    (IntType IntType
Int32, Signedness
Imp.Unsigned) -> [Char]
"ct.c_uint32"
    (IntType IntType
Int64, Signedness
Imp.Unsigned) -> [Char]
"ct.c_uint64"
    (IntType IntType
Int8, Signedness
_) -> [Char]
"ct.c_int8"
    (IntType IntType
Int16, Signedness
_) -> [Char]
"ct.c_int16"
    (IntType IntType
Int32, Signedness
_) -> [Char]
"ct.c_int32"
    (IntType IntType
Int64, Signedness
_) -> [Char]
"ct.c_int64"
    (FloatType FloatType
Float16, Signedness
_) -> [Char]
"ct.c_uint16"
    (FloatType FloatType
Float32, Signedness
_) -> [Char]
"ct.c_float"
    (FloatType FloatType
Float64, Signedness
_) -> [Char]
"ct.c_double"
    (PrimType
Imp.Bool, Signedness
_) -> [Char]
"ct.c_bool"
    (PrimType
Unit, Signedness
_) -> [Char]
"ct.c_byte"

-- | The Numpy type corresponding to a 'PrimType'.
compilePrimToNp :: Imp.PrimType -> String
compilePrimToNp :: PrimType -> [Char]
compilePrimToNp PrimType
bt =
  case PrimType
bt of
    IntType IntType
Int8 -> [Char]
"np.int8"
    IntType IntType
Int16 -> [Char]
"np.int16"
    IntType IntType
Int32 -> [Char]
"np.int32"
    IntType IntType
Int64 -> [Char]
"np.int64"
    FloatType FloatType
Float16 -> [Char]
"np.float16"
    FloatType FloatType
Float32 -> [Char]
"np.float32"
    FloatType FloatType
Float64 -> [Char]
"np.float64"
    PrimType
Imp.Bool -> [Char]
"np.byte"
    PrimType
Unit -> [Char]
"np.byte"

-- | The Numpy type corresponding to a 'PrimType', taking sign into account.
compilePrimToExtNp :: Imp.PrimType -> Imp.Signedness -> String
compilePrimToExtNp :: PrimType -> Signedness -> [Char]
compilePrimToExtNp PrimType
bt Signedness
ept =
  case (PrimType
bt, Signedness
ept) of
    (IntType IntType
Int8, Signedness
Imp.Unsigned) -> [Char]
"np.uint8"
    (IntType IntType
Int16, Signedness
Imp.Unsigned) -> [Char]
"np.uint16"
    (IntType IntType
Int32, Signedness
Imp.Unsigned) -> [Char]
"np.uint32"
    (IntType IntType
Int64, Signedness
Imp.Unsigned) -> [Char]
"np.uint64"
    (IntType IntType
Int8, Signedness
_) -> [Char]
"np.int8"
    (IntType IntType
Int16, Signedness
_) -> [Char]
"np.int16"
    (IntType IntType
Int32, Signedness
_) -> [Char]
"np.int32"
    (IntType IntType
Int64, Signedness
_) -> [Char]
"np.int64"
    (FloatType FloatType
Float16, Signedness
_) -> [Char]
"np.float16"
    (FloatType FloatType
Float32, Signedness
_) -> [Char]
"np.float32"
    (FloatType FloatType
Float64, Signedness
_) -> [Char]
"np.float64"
    (PrimType
Imp.Bool, Signedness
_) -> [Char]
"np.bool_"
    (PrimType
Unit, Signedness
_) -> [Char]
"np.byte"

-- | Convert from scalar to storage representation for the given type.
toStorage :: PrimType -> PyExp -> PyExp
toStorage :: PrimType -> PyExp -> PyExp
toStorage (FloatType FloatType
Float16) PyExp
e =
  [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"ct.c_int16" [[Char] -> [PyExp] -> PyExp
simpleCall [Char]
"futhark_to_bits16" [PyExp
e]]
toStorage PrimType
t PyExp
e = [Char] -> [PyExp] -> PyExp
simpleCall (PrimType -> [Char]
compilePrimType PrimType
t) [PyExp
e]

-- | Convert from storage to scalar representation for the given type.
fromStorage :: PrimType -> PyExp -> PyExp
fromStorage :: PrimType -> PyExp -> PyExp
fromStorage (FloatType FloatType
Float16) PyExp
e =
  [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"futhark_from_bits16" [[Char] -> [PyExp] -> PyExp
simpleCall [Char]
"np.int16" [PyExp
e]]
fromStorage PrimType
t PyExp
e = [Char] -> [PyExp] -> PyExp
simpleCall (PrimType -> [Char]
compilePrimToNp PrimType
t) [PyExp
e]

compilePrimValue :: Imp.PrimValue -> PyExp
compilePrimValue :: PrimValue -> PyExp
compilePrimValue (IntValue (Int8Value Int8
v)) =
  [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"np.int8" [Integer -> PyExp
Integer forall a b. (a -> b) -> a -> b
$ forall a. Integral a => a -> Integer
toInteger Int8
v]
compilePrimValue (IntValue (Int16Value Int16
v)) =
  [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"np.int16" [Integer -> PyExp
Integer forall a b. (a -> b) -> a -> b
$ forall a. Integral a => a -> Integer
toInteger Int16
v]
compilePrimValue (IntValue (Int32Value Int32
v)) =
  [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"np.int32" [Integer -> PyExp
Integer forall a b. (a -> b) -> a -> b
$ forall a. Integral a => a -> Integer
toInteger Int32
v]
compilePrimValue (IntValue (Int64Value Int64
v)) =
  [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"np.int64" [Integer -> PyExp
Integer forall a b. (a -> b) -> a -> b
$ forall a. Integral a => a -> Integer
toInteger Int64
v]
compilePrimValue (FloatValue (Float16Value Half
v))
  | forall a. RealFloat a => a -> Bool
isInfinite Half
v =
      if Half
v forall a. Ord a => a -> a -> Bool
> Half
0 then [Char] -> PyExp
Var [Char]
"np.inf" else [Char] -> PyExp
Var [Char]
"-np.inf"
  | forall a. RealFloat a => a -> Bool
isNaN Half
v =
      [Char] -> PyExp
Var [Char]
"np.nan"
  | Bool
otherwise = [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"np.float16" [Double -> PyExp
Float forall a b. (a -> b) -> a -> b
$ forall a. Fractional a => Rational -> a
fromRational forall a b. (a -> b) -> a -> b
$ forall a. Real a => a -> Rational
toRational Half
v]
compilePrimValue (FloatValue (Float32Value Float
v))
  | forall a. RealFloat a => a -> Bool
isInfinite Float
v =
      if Float
v forall a. Ord a => a -> a -> Bool
> Float
0 then [Char] -> PyExp
Var [Char]
"np.inf" else [Char] -> PyExp
Var [Char]
"-np.inf"
  | forall a. RealFloat a => a -> Bool
isNaN Float
v =
      [Char] -> PyExp
Var [Char]
"np.nan"
  | Bool
otherwise = [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"np.float32" [Double -> PyExp
Float forall a b. (a -> b) -> a -> b
$ forall a. Fractional a => Rational -> a
fromRational forall a b. (a -> b) -> a -> b
$ forall a. Real a => a -> Rational
toRational Float
v]
compilePrimValue (FloatValue (Float64Value Double
v))
  | forall a. RealFloat a => a -> Bool
isInfinite Double
v =
      if Double
v forall a. Ord a => a -> a -> Bool
> Double
0 then [Char] -> PyExp
Var [Char]
"np.inf" else [Char] -> PyExp
Var [Char]
"-np.inf"
  | forall a. RealFloat a => a -> Bool
isNaN Double
v =
      [Char] -> PyExp
Var [Char]
"np.nan"
  | Bool
otherwise = [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"np.float64" [Double -> PyExp
Float forall a b. (a -> b) -> a -> b
$ forall a. Fractional a => Rational -> a
fromRational forall a b. (a -> b) -> a -> b
$ forall a. Real a => a -> Rational
toRational Double
v]
compilePrimValue (BoolValue Bool
v) = Bool -> PyExp
Bool Bool
v
compilePrimValue PrimValue
UnitValue = [Char] -> PyExp
Var [Char]
"None"

compileVar :: VName -> CompilerM op s PyExp
compileVar :: forall op s. VName -> CompilerM op s PyExp
compileVar VName
v = forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a -> a
fromMaybe ([Char] -> PyExp
Var [Char]
v') forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup [Char]
v' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall op s. CompilerEnv op s -> Map [Char] PyExp
envVarExp
  where
    v' :: [Char]
v' = VName -> [Char]
compileName VName
v

-- | Tell me how to compile a @v@, and I'll Compile any @PrimExp v@ for you.
compilePrimExp :: Monad m => (v -> m PyExp) -> Imp.PrimExp v -> m PyExp
compilePrimExp :: forall (m :: * -> *) v.
Monad m =>
(v -> m PyExp) -> PrimExp v -> m PyExp
compilePrimExp v -> m PyExp
_ (Imp.ValueExp PrimValue
v) = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ PrimValue -> PyExp
compilePrimValue PrimValue
v
compilePrimExp v -> m PyExp
f (Imp.LeafExp v
v PrimType
_) = v -> m PyExp
f v
v
compilePrimExp v -> m PyExp
f (Imp.BinOpExp BinOp
op PrimExp v
x PrimExp v
y) = do
  (PyExp
x', PyExp
y', [Char] -> m PyExp
simple) <- forall (m :: * -> *) v.
Monad m =>
(v -> m PyExp)
-> PrimExp v -> PrimExp v -> m (PyExp, PyExp, [Char] -> m PyExp)
compileBinOpLike v -> m PyExp
f PrimExp v
x PrimExp v
y
  case BinOp
op of
    Add {} -> [Char] -> m PyExp
simple [Char]
"+"
    Sub {} -> [Char] -> m PyExp
simple [Char]
"-"
    Mul {} -> [Char] -> m PyExp
simple [Char]
"*"
    FAdd {} -> [Char] -> m PyExp
simple [Char]
"+"
    FSub {} -> [Char] -> m PyExp
simple [Char]
"-"
    FMul {} -> [Char] -> m PyExp
simple [Char]
"*"
    FDiv {} -> [Char] -> m PyExp
simple [Char]
"/"
    FMod {} -> [Char] -> m PyExp
simple [Char]
"%"
    Xor {} -> [Char] -> m PyExp
simple [Char]
"^"
    And {} -> [Char] -> m PyExp
simple [Char]
"&"
    Or {} -> [Char] -> m PyExp
simple [Char]
"|"
    Shl {} -> [Char] -> m PyExp
simple [Char]
"<<"
    LogAnd {} -> [Char] -> m PyExp
simple [Char]
"and"
    LogOr {} -> [Char] -> m PyExp
simple [Char]
"or"
    BinOp
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall (forall a. Pretty a => a -> [Char]
prettyString BinOp
op) [PyExp
x', PyExp
y']
compilePrimExp v -> m PyExp
f (Imp.ConvOpExp ConvOp
conv PrimExp v
x) = do
  PyExp
x' <- forall (m :: * -> *) v.
Monad m =>
(v -> m PyExp) -> PrimExp v -> m PyExp
compilePrimExp v -> m PyExp
f PrimExp v
x
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall (forall a. Pretty a => a -> [Char]
prettyString ConvOp
conv) [PyExp
x']
compilePrimExp v -> m PyExp
f (Imp.CmpOpExp CmpOp
cmp PrimExp v
x PrimExp v
y) = do
  (PyExp
x', PyExp
y', [Char] -> m PyExp
simple) <- forall (m :: * -> *) v.
Monad m =>
(v -> m PyExp)
-> PrimExp v -> PrimExp v -> m (PyExp, PyExp, [Char] -> m PyExp)
compileBinOpLike v -> m PyExp
f PrimExp v
x PrimExp v
y
  case CmpOp
cmp of
    CmpEq {} -> [Char] -> m PyExp
simple [Char]
"=="
    FCmpLt {} -> [Char] -> m PyExp
simple [Char]
"<"
    FCmpLe {} -> [Char] -> m PyExp
simple [Char]
"<="
    CmpOp
CmpLlt -> [Char] -> m PyExp
simple [Char]
"<"
    CmpOp
CmpLle -> [Char] -> m PyExp
simple [Char]
"<="
    CmpOp
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall (forall a. Pretty a => a -> [Char]
prettyString CmpOp
cmp) [PyExp
x', PyExp
y']
compilePrimExp v -> m PyExp
f (Imp.UnOpExp UnOp
op PrimExp v
exp1) =
  [Char] -> PyExp -> PyExp
UnOp (UnOp -> [Char]
compileUnOp UnOp
op) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) v.
Monad m =>
(v -> m PyExp) -> PrimExp v -> m PyExp
compilePrimExp v -> m PyExp
f PrimExp v
exp1
compilePrimExp v -> m PyExp
f (Imp.FunExp [Char]
h [PrimExp v]
args PrimType
_) =
  [Char] -> [PyExp] -> PyExp
simpleCall (Text -> [Char]
T.unpack (Text -> Text
futharkFun (forall a. Pretty a => a -> Text
prettyText [Char]
h))) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *) v.
Monad m =>
(v -> m PyExp) -> PrimExp v -> m PyExp
compilePrimExp v -> m PyExp
f) [PrimExp v]
args

compileExp :: Imp.Exp -> CompilerM op s PyExp
compileExp :: forall op s. Exp -> CompilerM op s PyExp
compileExp = forall (m :: * -> *) v.
Monad m =>
(v -> m PyExp) -> PrimExp v -> m PyExp
compilePrimExp forall op s. VName -> CompilerM op s PyExp
compileVar

errorMsgString :: Imp.ErrorMsg Imp.Exp -> CompilerM op s (T.Text, [PyExp])
errorMsgString :: forall op s. ErrorMsg Exp -> CompilerM op s (Text, [PyExp])
errorMsgString (Imp.ErrorMsg [ErrorMsgPart Exp]
parts) = do
  let onPart :: ErrorMsgPart Exp -> CompilerM op s (a, PyExp)
onPart (Imp.ErrorString Text
s) = forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
"%s", Text -> PyExp
String Text
s)
      onPart (Imp.ErrorVal IntType {} Exp
x) = (a
"%d",) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
x
      onPart (Imp.ErrorVal FloatType {} Exp
x) = (a
"%f",) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
x
      onPart (Imp.ErrorVal PrimType
Imp.Bool Exp
x) = (a
"%r",) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
x
      onPart (Imp.ErrorVal Unit {} Exp
x) = (a
"%r",) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
x
  ([Text]
formatstrs, [PyExp]
formatargs) <- forall a b. [(a, b)] -> ([a], [b])
unzip forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall {a} {op} {s}.
IsString a =>
ErrorMsgPart Exp -> CompilerM op s (a, PyExp)
onPart [ErrorMsgPart Exp]
parts
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Monoid a => [a] -> a
mconcat [Text]
formatstrs, [PyExp]
formatargs)

compileCode :: Imp.Code op -> CompilerM op s ()
compileCode :: forall op s. Code op -> CompilerM op s ()
compileCode Imp.DebugPrint {} =
  forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
compileCode Imp.TracePrint {} =
  forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
compileCode (Imp.Op op
op) =
  forall (m :: * -> *) a. Monad m => m (m a) -> m a
join forall a b. (a -> b) -> a -> b
$ forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall op s. CompilerEnv op s -> OpCompiler op s
envOpCompiler forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure op
op
compileCode (Imp.If TExp Bool
cond Code op
tb Code op
fb) = do
  PyExp
cond' <- forall op s. Exp -> CompilerM op s PyExp
compileExp forall a b. (a -> b) -> a -> b
$ forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Bool
cond
  [PyStmt]
tb' <- forall op s. CompilerM op s () -> CompilerM op s [PyStmt]
collect forall a b. (a -> b) -> a -> b
$ forall op s. Code op -> CompilerM op s ()
compileCode Code op
tb
  [PyStmt]
fb' <- forall op s. CompilerM op s () -> CompilerM op s [PyStmt]
collect forall a b. (a -> b) -> a -> b
$ forall op s. Code op -> CompilerM op s ()
compileCode Code op
fb
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$ PyExp -> [PyStmt] -> [PyStmt] -> PyStmt
If PyExp
cond' [PyStmt]
tb' [PyStmt]
fb'
compileCode (Code op
c1 Imp.:>>: Code op
c2) = do
  forall op s. Code op -> CompilerM op s ()
compileCode Code op
c1
  forall op s. Code op -> CompilerM op s ()
compileCode Code op
c2
compileCode (Imp.While TExp Bool
cond Code op
body) = do
  PyExp
cond' <- forall op s. Exp -> CompilerM op s PyExp
compileExp forall a b. (a -> b) -> a -> b
$ forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Bool
cond
  [PyStmt]
body' <- forall op s. CompilerM op s () -> CompilerM op s [PyStmt]
collect forall a b. (a -> b) -> a -> b
$ forall op s. Code op -> CompilerM op s ()
compileCode Code op
body
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$ PyExp -> [PyStmt] -> PyStmt
While PyExp
cond' [PyStmt]
body'
compileCode (Imp.For VName
i Exp
bound Code op
body) = do
  PyExp
bound' <- forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
bound
  let i' :: [Char]
i' = VName -> [Char]
compileName VName
i
  [PyStmt]
body' <- forall op s. CompilerM op s () -> CompilerM op s [PyStmt]
collect forall a b. (a -> b) -> a -> b
$ forall op s. Code op -> CompilerM op s ()
compileCode Code op
body
  [Char]
counter <- forall a. Pretty a => a -> [Char]
prettyString forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"counter"
  [Char]
one <- forall a. Pretty a => a -> [Char]
prettyString forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"one"
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$ PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
i') forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall (PrimType -> [Char]
compilePrimToNp (forall v. PrimExp v -> PrimType
Imp.primExpType Exp
bound)) [Integer -> PyExp
Integer Integer
0]
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$ PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
one) forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall (PrimType -> [Char]
compilePrimToNp (forall v. PrimExp v -> PrimType
Imp.primExpType Exp
bound)) [Integer -> PyExp
Integer Integer
1]
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$
    [Char] -> PyExp -> [PyStmt] -> PyStmt
For [Char]
counter ([Char] -> [PyExp] -> PyExp
simpleCall [Char]
"range" [PyExp
bound']) forall a b. (a -> b) -> a -> b
$
      [PyStmt]
body' forall a. [a] -> [a] -> [a]
++ [[Char] -> PyExp -> PyExp -> PyStmt
AssignOp [Char]
"+" ([Char] -> PyExp
Var [Char]
i') ([Char] -> PyExp
Var [Char]
one)]
compileCode (Imp.SetScalar VName
name Exp
exp1) =
  forall op s. PyStmt -> CompilerM op s ()
stm forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< PyExp -> PyExp -> PyStmt
Assign forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall op s. VName -> CompilerM op s PyExp
compileVar VName
name forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
exp1
compileCode Imp.DeclareMem {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
compileCode (Imp.DeclareScalar VName
v Volatility
_ PrimType
Unit) = do
  PyExp
v' <- forall op s. VName -> CompilerM op s PyExp
compileVar VName
v
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$ PyExp -> PyExp -> PyStmt
Assign PyExp
v' forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp
Var [Char]
"True"
compileCode Imp.DeclareScalar {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
compileCode (Imp.DeclareArray VName
name PrimType
t ArrayContents
vs) = do
  let arr_name :: [Char]
arr_name = VName -> [Char]
compileName VName
name forall a. Semigroup a => a -> a -> a
<> [Char]
"_arr"
  -- It is important to store the Numpy array in a temporary variable
  -- to prevent it from going "out-of-scope" before calling
  -- unwrapArray (which internally uses the .ctype method); see
  -- https://docs.scipy.org/doc/numpy/reference/generated/numpy.ndarray.ctypes.html
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$ PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
arr_name) forall a b. (a -> b) -> a -> b
$ case ArrayContents
vs of
    Imp.ArrayValues [PrimValue]
vs' ->
      PyExp -> [PyArg] -> PyExp
Call
        ([Char] -> PyExp
Var [Char]
"np.array")
        [ PyExp -> PyArg
Arg forall a b. (a -> b) -> a -> b
$ [PyExp] -> PyExp
List forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map PrimValue -> PyExp
compilePrimValue [PrimValue]
vs',
          [Char] -> PyExp -> PyArg
ArgKeyword [Char]
"dtype" forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp
Var forall a b. (a -> b) -> a -> b
$ PrimType -> [Char]
compilePrimToNp PrimType
t
        ]
    Imp.ArrayZeros Int
n ->
      PyExp -> [PyArg] -> PyExp
Call
        ([Char] -> PyExp
Var [Char]
"np.zeros")
        [ PyExp -> PyArg
Arg forall a b. (a -> b) -> a -> b
$ Integer -> PyExp
Integer forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n,
          [Char] -> PyExp -> PyArg
ArgKeyword [Char]
"dtype" forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp
Var forall a b. (a -> b) -> a -> b
$ PrimType -> [Char]
compilePrimToNp PrimType
t
        ]
  PyExp
name' <- forall op s. VName -> CompilerM op s PyExp
compileVar VName
name
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$ PyExp -> PyExp -> PyStmt
Assign PyExp
name' forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"unwrapArray" [[Char] -> PyExp
Var [Char]
arr_name]
compileCode (Imp.Comment Text
s Code op
code) = do
  [PyStmt]
code' <- forall op s. CompilerM op s () -> CompilerM op s [PyStmt]
collect forall a b. (a -> b) -> a -> b
$ forall op s. Code op -> CompilerM op s ()
compileCode Code op
code
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$ [Char] -> [PyStmt] -> PyStmt
Comment (Text -> [Char]
T.unpack Text
s) [PyStmt]
code'
compileCode (Imp.Assert Exp
e ErrorMsg Exp
msg (SrcLoc
loc, [SrcLoc]
locs)) = do
  PyExp
e' <- forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
e
  (Text
formatstr, [PyExp]
formatargs) <- forall op s. ErrorMsg Exp -> CompilerM op s (Text, [PyExp])
errorMsgString ErrorMsg Exp
msg
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$
    PyExp -> PyExp -> PyStmt
Assert
      PyExp
e'
      ( [Char] -> PyExp -> PyExp -> PyExp
BinOp
          [Char]
"%"
          (Text -> PyExp
String forall a b. (a -> b) -> a -> b
$ Text
"Error: " forall a. Semigroup a => a -> a -> a
<> Text
formatstr forall a. Semigroup a => a -> a -> a
<> Text
"\n\nBacktrace:\n" forall a. Semigroup a => a -> a -> a
<> Text
stacktrace)
          ([PyExp] -> PyExp
Tuple [PyExp]
formatargs)
      )
  where
    stacktrace :: Text
stacktrace = Int -> [Text] -> Text
prettyStacktrace Int
0 forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a. Located a => a -> Text
locText forall a b. (a -> b) -> a -> b
$ SrcLoc
loc forall a. a -> [a] -> [a]
: [SrcLoc]
locs
compileCode (Imp.Call [VName]
dests Name
fname [Arg]
args) = do
  [PyExp]
args' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall {op} {s}. Arg -> CompilerM op s PyExp
compileArg [Arg]
args
  PyExp
dests' <- [PyExp] -> PyExp
tupleOrSingle forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall op s. VName -> CompilerM op s PyExp
compileVar [VName]
dests
  let fname' :: Text
fname'
        | Name -> Bool
isBuiltInFunction Name
fname = Text -> Text
futharkFun (forall a. Pretty a => a -> Text
prettyText Name
fname)
        | Bool
otherwise = Text
"self." forall a. Semigroup a => a -> a -> a
<> Text -> Text
futharkFun (forall a. Pretty a => a -> Text
prettyText Name
fname)
      call' :: PyExp
call' = [Char] -> [PyExp] -> PyExp
simpleCall (Text -> [Char]
T.unpack Text
fname') [PyExp]
args'
  -- If the function returns nothing (is called only for side
  -- effects), take care not to assign to an empty tuple.
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$
    if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [VName]
dests
      then PyExp -> PyStmt
Exp PyExp
call'
      else PyExp -> PyExp -> PyStmt
Assign PyExp
dests' PyExp
call'
  where
    compileArg :: Arg -> CompilerM op s PyExp
compileArg (Imp.MemArg VName
m) = forall op s. VName -> CompilerM op s PyExp
compileVar VName
m
    compileArg (Imp.ExpArg Exp
e) = forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
e
compileCode (Imp.SetMem VName
dest VName
src Space
_) =
  forall op s. PyStmt -> CompilerM op s ()
stm forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< PyExp -> PyExp -> PyStmt
Assign forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall op s. VName -> CompilerM op s PyExp
compileVar VName
dest forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall op s. VName -> CompilerM op s PyExp
compileVar VName
src
compileCode (Imp.Allocate VName
name (Imp.Count (Imp.TPrimExp Exp
e)) (Imp.Space [Char]
space)) =
  forall (m :: * -> *) a. Monad m => m (m a) -> m a
join forall a b. (a -> b) -> a -> b
$
    forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall op s. CompilerEnv op s -> Allocate op s
envAllocate
      forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall op s. VName -> CompilerM op s PyExp
compileVar VName
name
      forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
e
      forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure [Char]
space
compileCode (Imp.Allocate VName
name (Imp.Count (Imp.TPrimExp Exp
e)) Space
_) = do
  PyExp
e' <- forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
e
  let allocate' :: PyExp
allocate' = [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"allocateMem" [PyExp
e']
  forall op s. PyStmt -> CompilerM op s ()
stm forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< PyExp -> PyExp -> PyStmt
Assign forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall op s. VName -> CompilerM op s PyExp
compileVar VName
name forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure PyExp
allocate'
compileCode (Imp.Free VName
name Space
_) =
  forall op s. PyStmt -> CompilerM op s ()
stm forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< PyExp -> PyExp -> PyStmt
Assign forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall op s. VName -> CompilerM op s PyExp
compileVar VName
name forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure PyExp
None
compileCode (Imp.Copy PrimType
_ VName
dest (Imp.Count TExp Int64
destoffset) Space
DefaultSpace VName
src (Imp.Count TExp Int64
srcoffset) Space
DefaultSpace (Imp.Count TExp Int64
size)) = do
  PyExp
destoffset' <- forall op s. Exp -> CompilerM op s PyExp
compileExp forall a b. (a -> b) -> a -> b
$ forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Int64
destoffset
  PyExp
srcoffset' <- forall op s. Exp -> CompilerM op s PyExp
compileExp forall a b. (a -> b) -> a -> b
$ forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Int64
srcoffset
  PyExp
dest' <- forall op s. VName -> CompilerM op s PyExp
compileVar VName
dest
  PyExp
src' <- forall op s. VName -> CompilerM op s PyExp
compileVar VName
src
  PyExp
size' <- forall op s. Exp -> CompilerM op s PyExp
compileExp forall a b. (a -> b) -> a -> b
$ forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Int64
size
  let offset_call1 :: PyExp
offset_call1 = [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"addressOffset" [PyExp
dest', PyExp
destoffset', [Char] -> PyExp
Var [Char]
"ct.c_byte"]
  let offset_call2 :: PyExp
offset_call2 = [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"addressOffset" [PyExp
src', PyExp
srcoffset', [Char] -> PyExp
Var [Char]
"ct.c_byte"]
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$ PyExp -> PyStmt
Exp forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"ct.memmove" [PyExp
offset_call1, PyExp
offset_call2, PyExp
size']
compileCode (Imp.Copy PrimType
pt VName
dest (Imp.Count TExp Int64
destoffset) Space
destspace VName
src (Imp.Count TExp Int64
srcoffset) Space
srcspace (Imp.Count TExp Int64
size)) = do
  Copy op s
copy <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall op s. CompilerEnv op s -> Copy op s
envCopy
  forall (m :: * -> *) a. Monad m => m (m a) -> m a
join forall a b. (a -> b) -> a -> b
$
    Copy op s
copy
      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall op s. VName -> CompilerM op s PyExp
compileVar VName
dest
      forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall op s. Exp -> CompilerM op s PyExp
compileExp (forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Int64
destoffset)
      forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure Space
destspace
      forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall op s. VName -> CompilerM op s PyExp
compileVar VName
src
      forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall op s. Exp -> CompilerM op s PyExp
compileExp (forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Int64
srcoffset)
      forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure Space
srcspace
      forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall op s. Exp -> CompilerM op s PyExp
compileExp (forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Int64
size)
      forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure PrimType
pt
compileCode (Imp.Write VName
_ Count Elements (TExp Int64)
_ PrimType
Unit Space
_ Volatility
_ Exp
_) = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
compileCode (Imp.Write VName
dest (Imp.Count TExp Int64
idx) PrimType
elemtype (Imp.Space [Char]
space) Volatility
_ Exp
elemexp) =
  forall (m :: * -> *) a. Monad m => m (m a) -> m a
join forall a b. (a -> b) -> a -> b
$
    forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall op s. CompilerEnv op s -> WriteScalar op s
envWriteScalar
      forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall op s. VName -> CompilerM op s PyExp
compileVar VName
dest
      forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall op s. Exp -> CompilerM op s PyExp
compileExp (forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Int64
idx)
      forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure PrimType
elemtype
      forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure [Char]
space
      forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
elemexp
compileCode (Imp.Write VName
dest (Imp.Count TExp Int64
idx) PrimType
elemtype Space
_ Volatility
_ Exp
elemexp) = do
  PyExp
idx' <- forall op s. Exp -> CompilerM op s PyExp
compileExp forall a b. (a -> b) -> a -> b
$ forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Int64
idx
  PyExp
elemexp' <- PrimType -> PyExp -> PyExp
toStorage PrimType
elemtype forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
elemexp
  PyExp
dest' <- forall op s. VName -> CompilerM op s PyExp
compileVar VName
dest
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$ PyExp -> PyStmt
Exp forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"writeScalarArray" [PyExp
dest', PyExp
idx', PyExp
elemexp']
compileCode (Imp.Read VName
x VName
_ Count Elements (TExp Int64)
_ PrimType
Unit Space
_ Volatility
_) =
  forall op s. PyStmt -> CompilerM op s ()
stm forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< PyExp -> PyExp -> PyStmt
Assign forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall op s. VName -> CompilerM op s PyExp
compileVar VName
x forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure (PrimValue -> PyExp
compilePrimValue PrimValue
UnitValue)
compileCode (Imp.Read VName
x VName
src (Imp.Count TExp Int64
iexp) PrimType
restype (Imp.Space [Char]
space) Volatility
_) = do
  PyExp
x' <- forall op s. VName -> CompilerM op s PyExp
compileVar VName
x
  PyExp
e <-
    forall (m :: * -> *) a. Monad m => m (m a) -> m a
join forall a b. (a -> b) -> a -> b
$
      forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall op s. CompilerEnv op s -> ReadScalar op s
envReadScalar
        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall op s. VName -> CompilerM op s PyExp
compileVar VName
src
        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall op s. Exp -> CompilerM op s PyExp
compileExp (forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Int64
iexp)
        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure PrimType
restype
        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure [Char]
space
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$ PyExp -> PyExp -> PyStmt
Assign PyExp
x' PyExp
e
compileCode (Imp.Read VName
x VName
src (Imp.Count TExp Int64
iexp) PrimType
bt Space
_ Volatility
_) = do
  PyExp
x' <- forall op s. VName -> CompilerM op s PyExp
compileVar VName
x
  PyExp
iexp' <- forall op s. Exp -> CompilerM op s PyExp
compileExp forall a b. (a -> b) -> a -> b
$ forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Int64
iexp
  let bt' :: [Char]
bt' = PrimType -> [Char]
compilePrimType PrimType
bt
  PyExp
src' <- forall op s. VName -> CompilerM op s PyExp
compileVar VName
src
  forall op s. PyStmt -> CompilerM op s ()
stm forall a b. (a -> b) -> a -> b
$ PyExp -> PyExp -> PyStmt
Assign PyExp
x' forall a b. (a -> b) -> a -> b
$ PrimType -> PyExp -> PyExp
fromStorage PrimType
bt forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"indexArray" [PyExp
src', PyExp
iexp', [Char] -> PyExp
Var [Char]
bt']
compileCode Code op
Imp.Skip = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()