{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Safe #-}
{-# LANGUAGE Strict #-}
-- |
--
-- This module implements a transformation from source to core
-- Futhark.
--
module Futhark.Internalise (internaliseProg) where

import Control.Monad.State
import Control.Monad.Reader
import Data.Bitraversable
import qualified Data.Map.Strict as M
import qualified Data.Set as S
import Data.List (find, intercalate, intersperse, nub, transpose)
import qualified Data.List.NonEmpty as NE

import Language.Futhark as E hiding (TypeArg)
import Language.Futhark.Semantic (Imports)
import Futhark.IR.SOACS as I hiding (stmPattern)
import Futhark.Transform.Rename as I
import Futhark.MonadFreshNames
import Futhark.Tools
import Futhark.Util (splitAt3)

import Futhark.Internalise.Monad as I
import Futhark.Internalise.AccurateSizes
import Futhark.Internalise.TypesValues
import Futhark.Internalise.Bindings
import Futhark.Internalise.Lambdas
import Futhark.Internalise.Defunctorise as Defunctorise
import Futhark.Internalise.Defunctionalise as Defunctionalise
import Futhark.Internalise.Monomorphise as Monomorphise

-- | Convert a program in source Futhark to a program in the Futhark
-- core language.
internaliseProg :: MonadFreshNames m =>
                   Bool -> Imports -> m (I.Prog SOACS)
internaliseProg :: Bool -> Imports -> m (Prog SOACS)
internaliseProg Bool
always_safe Imports
prog = do
  [Dec]
prog_decs <- Imports -> m [Dec]
forall (m :: * -> *). MonadFreshNames m => Imports -> m [Dec]
Defunctorise.transformProg Imports
prog
  [ValBind]
prog_decs' <- [Dec] -> m [ValBind]
forall (m :: * -> *). MonadFreshNames m => [Dec] -> m [ValBind]
Monomorphise.transformProg [Dec]
prog_decs
  [ValBind]
prog_decs'' <- [ValBind] -> m [ValBind]
forall (m :: * -> *). MonadFreshNames m => [ValBind] -> m [ValBind]
Defunctionalise.transformProg [ValBind]
prog_decs'
  (Stms SOACS
consts, [FunDef SOACS]
funs) <-
    Bool -> InternaliseM () -> m (Stms SOACS, [FunDef SOACS])
forall (m :: * -> *).
MonadFreshNames m =>
Bool -> InternaliseM () -> m (Stms SOACS, [FunDef SOACS])
runInternaliseM Bool
always_safe ([ValBind] -> InternaliseM ()
internaliseValBinds [ValBind]
prog_decs'')
  Prog SOACS -> m (Prog SOACS)
forall lore (m :: * -> *).
(Renameable lore, MonadFreshNames m) =>
Prog lore -> m (Prog lore)
I.renameProg (Prog SOACS -> m (Prog SOACS)) -> Prog SOACS -> m (Prog SOACS)
forall a b. (a -> b) -> a -> b
$ Stms SOACS -> [FunDef SOACS] -> Prog SOACS
forall lore. Stms lore -> [FunDef lore] -> Prog lore
I.Prog Stms SOACS
consts [FunDef SOACS]
funs

internaliseAttr :: E.AttrInfo -> Attr
internaliseAttr :: AttrInfo -> Attr
internaliseAttr (E.AttrAtom Name
v) = Name -> Attr
I.AttrAtom Name
v
internaliseAttr (E.AttrComp Name
f [AttrInfo]
attrs) = Name -> [Attr] -> Attr
I.AttrComp Name
f ([Attr] -> Attr) -> [Attr] -> Attr
forall a b. (a -> b) -> a -> b
$ (AttrInfo -> Attr) -> [AttrInfo] -> [Attr]
forall a b. (a -> b) -> [a] -> [b]
map AttrInfo -> Attr
internaliseAttr [AttrInfo]
attrs

internaliseAttrs :: [E.AttrInfo] -> Attrs
internaliseAttrs :: [AttrInfo] -> Attrs
internaliseAttrs = [Attrs] -> Attrs
forall a. Monoid a => [a] -> a
mconcat ([Attrs] -> Attrs)
-> ([AttrInfo] -> [Attrs]) -> [AttrInfo] -> Attrs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AttrInfo -> Attrs) -> [AttrInfo] -> [Attrs]
forall a b. (a -> b) -> [a] -> [b]
map (Attr -> Attrs
oneAttr (Attr -> Attrs) -> (AttrInfo -> Attr) -> AttrInfo -> Attrs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AttrInfo -> Attr
internaliseAttr)

internaliseValBinds :: [E.ValBind] -> InternaliseM ()
internaliseValBinds :: [ValBind] -> InternaliseM ()
internaliseValBinds = (ValBind -> InternaliseM ()) -> [ValBind] -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ValBind -> InternaliseM ()
internaliseValBind

internaliseFunName :: VName -> [E.Pattern] -> InternaliseM Name
internaliseFunName :: VName -> [Pattern] -> InternaliseM Name
internaliseFunName VName
ofname [] = Name -> InternaliseM Name
forall (m :: * -> *) a. Monad m => a -> m a
return (Name -> InternaliseM Name) -> Name -> InternaliseM Name
forall a b. (a -> b) -> a -> b
$ String -> Name
nameFromString (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ VName -> String
forall a. Pretty a => a -> String
pretty VName
ofname String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"f"
internaliseFunName VName
ofname [Pattern]
_  = do
  Maybe
  (Name, Closure, Closure, [DeclType], [Param DeclType],
   [(SubExp, Type)] -> Maybe [DeclExtType])
info <- VName -> InternaliseM (Maybe FunInfo)
lookupFunction' VName
ofname
  -- In some rare cases involving local functions, the same function
  -- name may be re-used in multiple places.  We check whether the
  -- function name has already been used, and generate a new one if
  -- so.
  case Maybe
  (Name, Closure, Closure, [DeclType], [Param DeclType],
   [(SubExp, Type)] -> Maybe [DeclExtType])
info of
    Just (Name, Closure, Closure, [DeclType], [Param DeclType],
 [(SubExp, Type)] -> Maybe [DeclExtType])
_ -> String -> Name
nameFromString (String -> Name) -> (VName -> String) -> VName -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> String
forall a. Pretty a => a -> String
pretty (VName -> Name) -> InternaliseM VName -> InternaliseM Name
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString (VName -> String
baseString VName
ofname)
    Maybe
  (Name, Closure, Closure, [DeclType], [Param DeclType],
   [(SubExp, Type)] -> Maybe [DeclExtType])
Nothing -> Name -> InternaliseM Name
forall (m :: * -> *) a. Monad m => a -> m a
return (Name -> InternaliseM Name) -> Name -> InternaliseM Name
forall a b. (a -> b) -> a -> b
$ String -> Name
nameFromString (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ VName -> String
forall a. Pretty a => a -> String
pretty VName
ofname

internaliseValBind :: E.ValBind -> InternaliseM ()
internaliseValBind :: ValBind -> InternaliseM ()
internaliseValBind fb :: ValBind
fb@(E.ValBind Maybe (Info EntryPoint)
entry VName
fname Maybe (TypeExp VName)
retdecl (Info (StructType
rettype, Closure
_)) [TypeParamBase VName]
tparams [Pattern]
params ExpBase Info VName
body Maybe DocComment
_ [AttrInfo]
attrs SrcLoc
loc) = do
  InternaliseM () -> InternaliseM ()
forall a. InternaliseM a -> InternaliseM a
localConstsScope (InternaliseM () -> InternaliseM ())
-> InternaliseM () -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ [TypeParamBase VName]
-> [Pattern]
-> ([FParam] -> [[FParam]] -> InternaliseM ())
-> InternaliseM ()
forall a.
[TypeParamBase VName]
-> [Pattern]
-> ([FParam] -> [[FParam]] -> InternaliseM a)
-> InternaliseM a
bindingParams [TypeParamBase VName]
tparams [Pattern]
params (([FParam] -> [[FParam]] -> InternaliseM ()) -> InternaliseM ())
-> ([FParam] -> [[FParam]] -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \[FParam]
shapeparams [[FParam]]
params' -> do
    [DeclExtType]
rettype_bad <- StructType -> InternaliseM [DeclExtType]
internaliseReturnType StructType
rettype
    let rettype' :: [DeclExtType]
rettype' = [DeclExtType] -> [DeclExtType]
forall u. [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts [DeclExtType]
rettype_bad

    let shapenames :: Closure
shapenames = (Param DeclType -> VName) -> [Param DeclType] -> Closure
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
[FParam]
shapeparams
        normal_params :: Closure
normal_params = Closure
shapenames Closure -> Closure -> Closure
forall a. [a] -> [a] -> [a]
++ (Param DeclType -> VName) -> [Param DeclType] -> Closure
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName ([[Param DeclType]] -> [Param DeclType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Param DeclType]]
[[FParam]]
params')
        normal_param_names :: Names
normal_param_names = Closure -> Names
namesFromList Closure
normal_params

    Name
fname' <- VName -> [Pattern] -> InternaliseM Name
internaliseFunName VName
fname [Pattern]
params

    Body
body' <- do
      ErrorMsg SubExp
msg <- case Maybe (TypeExp VName)
retdecl of
               Just TypeExp VName
dt -> [ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg ([ErrorMsgPart SubExp] -> ErrorMsg SubExp)
-> ([ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp]
-> ErrorMsg SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
                          (ErrorMsgPart SubExp
"Function return value does not match shape of type "ErrorMsgPart SubExp
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. a -> [a] -> [a]
:) ([ErrorMsgPart SubExp] -> ErrorMsg SubExp)
-> InternaliseM [ErrorMsgPart SubExp]
-> InternaliseM (ErrorMsg SubExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
                          TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
dt
               Maybe (TypeExp VName)
Nothing -> ErrorMsg SubExp -> InternaliseM (ErrorMsg SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (ErrorMsg SubExp -> InternaliseM (ErrorMsg SubExp))
-> ErrorMsg SubExp -> InternaliseM (ErrorMsg SubExp)
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [ErrorMsgPart SubExp
"Function return value does not match shape of declared return type."]
      ExpBase Info VName -> InternaliseM Body
internaliseBody ExpBase Info VName
body InternaliseM Body
-> (Body -> InternaliseM Body) -> InternaliseM Body
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=
        ErrorMsg SubExp -> SrcLoc -> [ExtType] -> Body -> InternaliseM Body
ensureResultExtShape ErrorMsg SubExp
msg SrcLoc
loc ((DeclExtType -> ExtType) -> [DeclExtType] -> [ExtType]
forall a b. (a -> b) -> [a] -> [b]
map DeclExtType -> ExtType
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl [DeclExtType]
rettype')

    Names
constants <- InternaliseM Names
allConsts
    let free_in_fun :: Names
free_in_fun = Body -> Names
forall a. FreeIn a => a -> Names
freeIn Body
body'
                      Names -> Names -> Names
`namesSubtract` Names
normal_param_names
                      Names -> Names -> Names
`namesSubtract` Names
constants

    [Param DeclType]
used_free_params <- Closure
-> (VName -> InternaliseM (Param DeclType))
-> InternaliseM [Param DeclType]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (Names -> Closure
namesToList Names
free_in_fun) ((VName -> InternaliseM (Param DeclType))
 -> InternaliseM [Param DeclType])
-> (VName -> InternaliseM (Param DeclType))
-> InternaliseM [Param DeclType]
forall a b. (a -> b) -> a -> b
$ \VName
v -> do
      Type
v_t <- VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
v
      Param DeclType -> InternaliseM (Param DeclType)
forall (m :: * -> *) a. Monad m => a -> m a
return (Param DeclType -> InternaliseM (Param DeclType))
-> Param DeclType -> InternaliseM (Param DeclType)
forall a b. (a -> b) -> a -> b
$ VName -> DeclType -> Param DeclType
forall dec. VName -> dec -> Param dec
Param VName
v (DeclType -> Param DeclType) -> DeclType -> Param DeclType
forall a b. (a -> b) -> a -> b
$ Type -> Uniqueness -> DeclType
forall shape.
TypeBase shape NoUniqueness
-> Uniqueness -> TypeBase shape Uniqueness
toDecl Type
v_t Uniqueness
Nonunique

    let free_shape_params :: [Param DeclType]
free_shape_params = (VName -> Param DeclType) -> Closure -> [Param DeclType]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> DeclType -> Param DeclType
forall dec. VName -> dec -> Param dec
`Param` PrimType -> DeclType
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int32) (Closure -> [Param DeclType]) -> Closure -> [Param DeclType]
forall a b. (a -> b) -> a -> b
$
                            (Param DeclType -> Closure) -> [Param DeclType] -> Closure
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Shape -> Closure
I.shapeVars (Shape -> Closure)
-> (Param DeclType -> Shape) -> Param DeclType -> Closure
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape (Type -> Shape)
-> (Param DeclType -> Type) -> Param DeclType -> Shape
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param DeclType -> Type
forall dec. Typed dec => Param dec -> Type
I.paramType) [Param DeclType]
used_free_params
        free_params :: [Param DeclType]
free_params = [Param DeclType] -> [Param DeclType]
forall a. Eq a => [a] -> [a]
nub ([Param DeclType] -> [Param DeclType])
-> [Param DeclType] -> [Param DeclType]
forall a b. (a -> b) -> a -> b
$ [Param DeclType]
free_shape_params [Param DeclType] -> [Param DeclType] -> [Param DeclType]
forall a. [a] -> [a] -> [a]
++ [Param DeclType]
used_free_params
        all_params :: [Param DeclType]
all_params = [Param DeclType]
free_params [Param DeclType] -> [Param DeclType] -> [Param DeclType]
forall a. [a] -> [a] -> [a]
++ [Param DeclType]
[FParam]
shapeparams [Param DeclType] -> [Param DeclType] -> [Param DeclType]
forall a. [a] -> [a] -> [a]
++ [[Param DeclType]] -> [Param DeclType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Param DeclType]]
[[FParam]]
params'

    let fd :: FunDef SOACS
fd = Maybe EntryPoint
-> Attrs
-> Name
-> [RetType SOACS]
-> [FParam]
-> Body
-> FunDef SOACS
forall lore.
Maybe EntryPoint
-> Attrs
-> Name
-> [RetType lore]
-> [FParam lore]
-> BodyT lore
-> FunDef lore
I.FunDef Maybe EntryPoint
forall a. Maybe a
Nothing ([AttrInfo] -> Attrs
internaliseAttrs [AttrInfo]
attrs) Name
fname'
             [DeclExtType]
[RetType SOACS]
rettype' [Param DeclType]
[FParam]
all_params Body
body'

    if [[Param DeclType]] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[Param DeclType]]
[[FParam]]
params'
      then VName -> FunDef SOACS -> InternaliseM ()
bindConstant VName
fname FunDef SOACS
fd
      else VName -> FunDef SOACS -> FunInfo -> InternaliseM ()
bindFunction VName
fname FunDef SOACS
fd
           (Name
fname',
            (Param DeclType -> VName) -> [Param DeclType] -> Closure
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
free_params,
            Closure
shapenames,
            (Param DeclType -> DeclType) -> [Param DeclType] -> [DeclType]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> DeclType
forall t. DeclTyped t => t -> DeclType
declTypeOf ([Param DeclType] -> [DeclType]) -> [Param DeclType] -> [DeclType]
forall a b. (a -> b) -> a -> b
$ [[Param DeclType]] -> [Param DeclType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Param DeclType]]
[[FParam]]
params',
            [Param DeclType]
[FParam]
all_params,
            [DeclExtType]
-> [Param DeclType] -> [(SubExp, Type)] -> Maybe [DeclExtType]
forall rt dec.
(IsRetType rt, Typed dec) =>
[rt] -> [Param dec] -> [(SubExp, Type)] -> Maybe [rt]
applyRetType [DeclExtType]
rettype' [Param DeclType]
all_params)

  case Maybe (Info EntryPoint)
entry of Just (Info EntryPoint
entry') -> EntryPoint -> ValBind -> InternaliseM ()
generateEntryPoint EntryPoint
entry' ValBind
fb
                Maybe (Info EntryPoint)
Nothing -> () -> InternaliseM ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

  where
    -- | Recompute existential sizes to start from zero.
    -- Necessary because some convoluted constructions will start
    -- them from somewhere else.
    zeroExts :: [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts [TypeBase ExtShape u]
ts = [TypeBase ExtShape u]
-> [TypeBase ExtShape u] -> [TypeBase ExtShape u]
forall u.
[TypeBase ExtShape u]
-> [TypeBase ExtShape u] -> [TypeBase ExtShape u]
generaliseExtTypes [TypeBase ExtShape u]
ts [TypeBase ExtShape u]
ts

allDimsFreshInType :: MonadFreshNames m => E.PatternType -> m E.PatternType
allDimsFreshInType :: PatternType -> m PatternType
allDimsFreshInType = (DimDecl VName -> m (DimDecl VName))
-> (Aliasing -> m Aliasing) -> PatternType -> m PatternType
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse DimDecl VName -> m (DimDecl VName)
forall (f :: * -> *).
MonadFreshNames f =>
DimDecl VName -> f (DimDecl VName)
onDim Aliasing -> m Aliasing
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  where onDim :: DimDecl VName -> f (DimDecl VName)
onDim (E.NamedDim QualName VName
v) =
          QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
E.NamedDim (QualName VName -> DimDecl VName)
-> (VName -> QualName VName) -> VName -> DimDecl VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> QualName VName
forall v. v -> QualName v
E.qualName (VName -> DimDecl VName) -> f VName -> f (DimDecl VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> f VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName (VName -> String
baseString (VName -> String) -> VName -> String
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
E.qualLeaf QualName VName
v)
        onDim DimDecl VName
_ = DimDecl VName -> f (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure DimDecl VName
forall vn. DimDecl vn
AnyDim

-- | Replace all named dimensions with a fresh name, and remove all
-- constant dimensions.  The point is to remove the constraints, but
-- keep the names around.  We use this for constructing the entry
-- point parameters.
allDimsFreshInPat :: MonadFreshNames m => E.Pattern -> m E.Pattern
allDimsFreshInPat :: Pattern -> m Pattern
allDimsFreshInPat (PatternAscription Pattern
p TypeDeclBase Info VName
_ SrcLoc
_) =
  Pattern -> m Pattern
forall (m :: * -> *). MonadFreshNames m => Pattern -> m Pattern
allDimsFreshInPat Pattern
p
allDimsFreshInPat (PatternParens Pattern
p SrcLoc
_) =
  Pattern -> m Pattern
forall (m :: * -> *). MonadFreshNames m => Pattern -> m Pattern
allDimsFreshInPat Pattern
p
allDimsFreshInPat (Id VName
v (Info PatternType
t) SrcLoc
loc) =
  VName -> Info PatternType -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase f vn
Id VName
v (Info PatternType -> SrcLoc -> Pattern)
-> m (Info PatternType) -> m (SrcLoc -> Pattern)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> m PatternType -> m (Info PatternType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PatternType -> m PatternType
forall (m :: * -> *).
MonadFreshNames m =>
PatternType -> m PatternType
allDimsFreshInType PatternType
t) m (SrcLoc -> Pattern) -> m SrcLoc -> m Pattern
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> m SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
allDimsFreshInPat (TuplePattern [Pattern]
ps SrcLoc
loc) =
  [Pattern] -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
[PatternBase f vn] -> SrcLoc -> PatternBase f vn
TuplePattern ([Pattern] -> SrcLoc -> Pattern)
-> m [Pattern] -> m (SrcLoc -> Pattern)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Pattern -> m Pattern) -> [Pattern] -> m [Pattern]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Pattern -> m Pattern
forall (m :: * -> *). MonadFreshNames m => Pattern -> m Pattern
allDimsFreshInPat [Pattern]
ps m (SrcLoc -> Pattern) -> m SrcLoc -> m Pattern
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> m SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
allDimsFreshInPat (RecordPattern [(Name, Pattern)]
ps SrcLoc
loc) =
  [(Name, Pattern)] -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
[(Name, PatternBase f vn)] -> SrcLoc -> PatternBase f vn
RecordPattern ([(Name, Pattern)] -> SrcLoc -> Pattern)
-> m [(Name, Pattern)] -> m (SrcLoc -> Pattern)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Name, Pattern) -> m (Name, Pattern))
-> [(Name, Pattern)] -> m [(Name, Pattern)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Pattern -> m Pattern) -> (Name, Pattern) -> m (Name, Pattern)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Pattern -> m Pattern
forall (m :: * -> *). MonadFreshNames m => Pattern -> m Pattern
allDimsFreshInPat) [(Name, Pattern)]
ps m (SrcLoc -> Pattern) -> m SrcLoc -> m Pattern
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> m SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
allDimsFreshInPat (Wildcard (Info PatternType
t) SrcLoc
loc) =
  Info PatternType -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
f PatternType -> SrcLoc -> PatternBase f vn
Wildcard (Info PatternType -> SrcLoc -> Pattern)
-> m (Info PatternType) -> m (SrcLoc -> Pattern)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> m PatternType -> m (Info PatternType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PatternType -> m PatternType
forall (m :: * -> *).
MonadFreshNames m =>
PatternType -> m PatternType
allDimsFreshInType PatternType
t) m (SrcLoc -> Pattern) -> m SrcLoc -> m Pattern
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> m SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
allDimsFreshInPat (PatternLit ExpBase Info VName
e (Info PatternType
t) SrcLoc
loc) =
  ExpBase Info VName -> Info PatternType -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
ExpBase f vn -> f PatternType -> SrcLoc -> PatternBase f vn
PatternLit ExpBase Info VName
e (Info PatternType -> SrcLoc -> Pattern)
-> m (Info PatternType) -> m (SrcLoc -> Pattern)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> m PatternType -> m (Info PatternType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PatternType -> m PatternType
forall (m :: * -> *).
MonadFreshNames m =>
PatternType -> m PatternType
allDimsFreshInType PatternType
t) m (SrcLoc -> Pattern) -> m SrcLoc -> m Pattern
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> m SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
allDimsFreshInPat (PatternConstr Name
c (Info PatternType
t) [Pattern]
pats SrcLoc
loc) =
  Name -> Info PatternType -> [Pattern] -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
Name
-> f PatternType
-> [PatternBase f vn]
-> SrcLoc
-> PatternBase f vn
PatternConstr Name
c (Info PatternType -> [Pattern] -> SrcLoc -> Pattern)
-> m (Info PatternType) -> m ([Pattern] -> SrcLoc -> Pattern)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> m PatternType -> m (Info PatternType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PatternType -> m PatternType
forall (m :: * -> *).
MonadFreshNames m =>
PatternType -> m PatternType
allDimsFreshInType PatternType
t) m ([Pattern] -> SrcLoc -> Pattern)
-> m [Pattern] -> m (SrcLoc -> Pattern)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
  (Pattern -> m Pattern) -> [Pattern] -> m [Pattern]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Pattern -> m Pattern
forall (m :: * -> *). MonadFreshNames m => Pattern -> m Pattern
allDimsFreshInPat [Pattern]
pats m (SrcLoc -> Pattern) -> m SrcLoc -> m Pattern
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> m SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc

generateEntryPoint :: E.EntryPoint -> E.ValBind -> InternaliseM ()
generateEntryPoint :: EntryPoint -> ValBind -> InternaliseM ()
generateEntryPoint (E.EntryPoint [EntryType]
e_paramts EntryType
e_rettype) ValBind
vb = InternaliseM () -> InternaliseM ()
forall a. InternaliseM a -> InternaliseM a
localConstsScope (InternaliseM () -> InternaliseM ())
-> InternaliseM () -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ do
  let (E.ValBind Maybe (Info EntryPoint)
_ VName
ofname Maybe (TypeExp VName)
_ (Info (StructType
rettype, Closure
_)) [TypeParamBase VName]
_ [Pattern]
params ExpBase Info VName
_ Maybe DocComment
_ [AttrInfo]
attrs SrcLoc
loc) = ValBind
vb
  -- We replace all shape annotations, so there should be no constant
  -- parameters here.
  [Pattern]
params_fresh <- (Pattern -> InternaliseM Pattern)
-> [Pattern] -> InternaliseM [Pattern]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Pattern -> InternaliseM Pattern
forall (m :: * -> *). MonadFreshNames m => Pattern -> m Pattern
allDimsFreshInPat [Pattern]
params
  let tparams :: [TypeParamBase VName]
tparams = (VName -> TypeParamBase VName) -> Closure -> [TypeParamBase VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SrcLoc -> TypeParamBase VName
forall vn. vn -> SrcLoc -> TypeParamBase vn
`E.TypeParamDim` SrcLoc
forall a. Monoid a => a
mempty) (Closure -> [TypeParamBase VName])
-> Closure -> [TypeParamBase VName]
forall a b. (a -> b) -> a -> b
$ Set VName -> Closure
forall a. Set a -> [a]
S.toList (Set VName -> Closure) -> Set VName -> Closure
forall a b. (a -> b) -> a -> b
$
                [Set VName] -> Set VName
forall a. Monoid a => [a] -> a
mconcat ([Set VName] -> Set VName) -> [Set VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ (Pattern -> Set VName) -> [Pattern] -> [Set VName]
forall a b. (a -> b) -> [a] -> [b]
map Pattern -> Set VName
E.patternDimNames [Pattern]
params_fresh
  [TypeParamBase VName]
-> [Pattern]
-> ([FParam] -> [[FParam]] -> InternaliseM ())
-> InternaliseM ()
forall a.
[TypeParamBase VName]
-> [Pattern]
-> ([FParam] -> [[FParam]] -> InternaliseM a)
-> InternaliseM a
bindingParams [TypeParamBase VName]
tparams [Pattern]
params_fresh (([FParam] -> [[FParam]] -> InternaliseM ()) -> InternaliseM ())
-> ([FParam] -> [[FParam]] -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \[FParam]
shapeparams [[FParam]]
params' -> do
    [[DeclExtType]]
entry_rettype <- StructType -> InternaliseM [[DeclExtType]]
internaliseEntryReturnType (StructType -> InternaliseM [[DeclExtType]])
-> StructType -> InternaliseM [[DeclExtType]]
forall a b. (a -> b) -> a -> b
$ StructType -> StructType
forall vn as. TypeBase (DimDecl vn) as -> TypeBase (DimDecl vn) as
anySizes StructType
rettype
    let entry' :: EntryPoint
entry' = [(EntryType, [FParam])]
-> (EntryType, [[DeclExtType]]) -> EntryPoint
entryPoint ([EntryType]
-> [[Param DeclType]] -> [(EntryType, [Param DeclType])]
forall a b. [a] -> [b] -> [(a, b)]
zip [EntryType]
e_paramts [[Param DeclType]]
[[FParam]]
params') (EntryType
e_rettype, [[DeclExtType]]
entry_rettype)
        args :: [SubExp]
args = (Param DeclType -> SubExp) -> [Param DeclType] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
I.Var (VName -> SubExp)
-> (Param DeclType -> VName) -> Param DeclType -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName) ([Param DeclType] -> [SubExp]) -> [Param DeclType] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ [[Param DeclType]] -> [Param DeclType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Param DeclType]]
[[FParam]]
params'

    Body
entry_body <- InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
m (Body (Lore m)) -> m (Body (Lore m))
insertStmsM (InternaliseM (Body (Lore InternaliseM))
 -> InternaliseM (Body (Lore InternaliseM)))
-> InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ do
      -- Special case the (rare) situation where the entry point is
      -- not a function.
      Maybe [SubExp]
maybe_const <- VName -> InternaliseM (Maybe [SubExp])
lookupConst VName
ofname
      [SubExp]
vals <- case Maybe [SubExp]
maybe_const of
                Just [SubExp]
ses ->
                  [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [SubExp]
ses
                Maybe [SubExp]
Nothing ->
                  ([SubExp], [ExtType]) -> [SubExp]
forall a b. (a, b) -> a
fst (([SubExp], [ExtType]) -> [SubExp])
-> InternaliseM ([SubExp], [ExtType]) -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String
-> QualName VName
-> [SubExp]
-> SrcLoc
-> InternaliseM ([SubExp], [ExtType])
funcall String
"entry_result" (VName -> QualName VName
forall v. v -> QualName v
E.qualName VName
ofname) [SubExp]
args SrcLoc
loc
      [SubExp]
ctx <- [DeclExtType] -> [[SubExp]] -> [SubExp]
forall u a. [TypeBase ExtShape u] -> [[a]] -> [a]
extractShapeContext ([[DeclExtType]] -> [DeclExtType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[DeclExtType]]
entry_rettype) ([[SubExp]] -> [SubExp])
-> InternaliseM [[SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
             (SubExp -> InternaliseM [SubExp])
-> [SubExp] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Type -> [SubExp]) -> InternaliseM Type -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Type -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
I.arrayDims (InternaliseM Type -> InternaliseM [SubExp])
-> (SubExp -> InternaliseM Type) -> SubExp -> InternaliseM [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType) [SubExp]
vals
      [SubExp] -> InternaliseM (Body (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
[SubExp] -> m (Body (Lore m))
resultBodyM ([SubExp]
ctx [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
vals)

    FunDef SOACS -> InternaliseM ()
addFunDef (FunDef SOACS -> InternaliseM ())
-> FunDef SOACS -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
      Maybe EntryPoint
-> Attrs
-> Name
-> [RetType SOACS]
-> [FParam]
-> Body
-> FunDef SOACS
forall lore.
Maybe EntryPoint
-> Attrs
-> Name
-> [RetType lore]
-> [FParam lore]
-> BodyT lore
-> FunDef lore
I.FunDef (EntryPoint -> Maybe EntryPoint
forall a. a -> Maybe a
Just EntryPoint
entry') ([AttrInfo] -> Attrs
internaliseAttrs [AttrInfo]
attrs)
      (VName -> Name
baseName VName
ofname)
      ([[DeclExtType]] -> [DeclExtType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[DeclExtType]]
entry_rettype)
      ([Param DeclType]
[FParam]
shapeparams [Param DeclType] -> [Param DeclType] -> [Param DeclType]
forall a. [a] -> [a] -> [a]
++ [[Param DeclType]] -> [Param DeclType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Param DeclType]]
[[FParam]]
params') Body
entry_body

entryPoint :: [(E.EntryType, [I.FParam])]
           -> (E.EntryType,
               [[I.TypeBase ExtShape Uniqueness]])
           -> I.EntryPoint
entryPoint :: [(EntryType, [FParam])]
-> (EntryType, [[DeclExtType]]) -> EntryPoint
entryPoint [(EntryType, [FParam])]
params (EntryType
eret, [[DeclExtType]]
crets) =
  (((EntryType, [Param DeclType]) -> [EntryPointType])
-> [(EntryType, [Param DeclType])] -> [EntryPointType]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((EntryType, [DeclExtType]) -> [EntryPointType]
forall (t :: * -> *) a.
Foldable t =>
(EntryType, t a) -> [EntryPointType]
entryPointType ((EntryType, [DeclExtType]) -> [EntryPointType])
-> ((EntryType, [Param DeclType]) -> (EntryType, [DeclExtType]))
-> (EntryType, [Param DeclType])
-> [EntryPointType]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (EntryType, [Param DeclType]) -> (EntryType, [DeclExtType])
forall dec a.
DeclTyped dec =>
(a, [Param dec]) -> (a, [DeclExtType])
preParam) [(EntryType, [Param DeclType])]
[(EntryType, [FParam])]
params,
   case (StructType -> Maybe [StructType]
forall dim as. TypeBase dim as -> Maybe [TypeBase dim as]
isTupleRecord (StructType -> Maybe [StructType])
-> StructType -> Maybe [StructType]
forall a b. (a -> b) -> a -> b
$ EntryType -> StructType
entryType EntryType
eret,
         EntryType -> Maybe (TypeExp VName)
entryAscribed EntryType
eret) of
     (Just [StructType]
ts, Just (E.TETuple [TypeExp VName]
e_ts SrcLoc
_)) ->
       ((EntryType, [DeclExtType]) -> [EntryPointType])
-> [(EntryType, [DeclExtType])] -> [EntryPointType]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (EntryType, [DeclExtType]) -> [EntryPointType]
forall (t :: * -> *) a.
Foldable t =>
(EntryType, t a) -> [EntryPointType]
entryPointType ([(EntryType, [DeclExtType])] -> [EntryPointType])
-> [(EntryType, [DeclExtType])] -> [EntryPointType]
forall a b. (a -> b) -> a -> b
$
       [EntryType] -> [[DeclExtType]] -> [(EntryType, [DeclExtType])]
forall a b. [a] -> [b] -> [(a, b)]
zip ((StructType -> Maybe (TypeExp VName) -> EntryType)
-> [StructType] -> [Maybe (TypeExp VName)] -> [EntryType]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith StructType -> Maybe (TypeExp VName) -> EntryType
E.EntryType [StructType]
ts ((TypeExp VName -> Maybe (TypeExp VName))
-> [TypeExp VName] -> [Maybe (TypeExp VName)]
forall a b. (a -> b) -> [a] -> [b]
map TypeExp VName -> Maybe (TypeExp VName)
forall a. a -> Maybe a
Just [TypeExp VName]
e_ts)) [[DeclExtType]]
crets
     (Just [StructType]
ts, Maybe (TypeExp VName)
Nothing) ->
       ((EntryType, [DeclExtType]) -> [EntryPointType])
-> [(EntryType, [DeclExtType])] -> [EntryPointType]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (EntryType, [DeclExtType]) -> [EntryPointType]
forall (t :: * -> *) a.
Foldable t =>
(EntryType, t a) -> [EntryPointType]
entryPointType ([(EntryType, [DeclExtType])] -> [EntryPointType])
-> [(EntryType, [DeclExtType])] -> [EntryPointType]
forall a b. (a -> b) -> a -> b
$
       [EntryType] -> [[DeclExtType]] -> [(EntryType, [DeclExtType])]
forall a b. [a] -> [b] -> [(a, b)]
zip ((StructType -> EntryType) -> [StructType] -> [EntryType]
forall a b. (a -> b) -> [a] -> [b]
map (StructType -> Maybe (TypeExp VName) -> EntryType
`E.EntryType` Maybe (TypeExp VName)
forall a. Maybe a
Nothing) [StructType]
ts) [[DeclExtType]]
crets
     (Maybe [StructType], Maybe (TypeExp VName))
_ ->
       (EntryType, [DeclExtType]) -> [EntryPointType]
forall (t :: * -> *) a.
Foldable t =>
(EntryType, t a) -> [EntryPointType]
entryPointType (EntryType
eret, [[DeclExtType]] -> [DeclExtType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[DeclExtType]]
crets))
  where preParam :: (a, [Param dec]) -> (a, [DeclExtType])
preParam (a
e_t, [Param dec]
ps) = (a
e_t, [DeclType] -> [DeclExtType]
forall u. [TypeBase Shape u] -> [TypeBase ExtShape u]
staticShapes ([DeclType] -> [DeclExtType]) -> [DeclType] -> [DeclExtType]
forall a b. (a -> b) -> a -> b
$ (Param dec -> DeclType) -> [Param dec] -> [DeclType]
forall a b. (a -> b) -> [a] -> [b]
map Param dec -> DeclType
forall dec. DeclTyped dec => Param dec -> DeclType
I.paramDeclType [Param dec]
ps)

        entryPointType :: (EntryType, t a) -> [EntryPointType]
entryPointType (EntryType
t, t a
ts)
          | E.Scalar (E.Prim E.Unsigned{}) <- EntryType -> StructType
E.entryType EntryType
t =
              [EntryPointType
I.TypeUnsigned]
          | E.Array ()
_ Uniqueness
_ (E.Prim E.Unsigned{}) ShapeDecl (DimDecl VName)
_ <- EntryType -> StructType
E.entryType EntryType
t =
              [EntryPointType
I.TypeUnsigned]
          | E.Scalar E.Prim{} <- EntryType -> StructType
E.entryType EntryType
t =
              [EntryPointType
I.TypeDirect]
          | E.Array ()
_ Uniqueness
_ E.Prim{} ShapeDecl (DimDecl VName)
_ <- EntryType -> StructType
E.entryType EntryType
t =
              [EntryPointType
I.TypeDirect]
          | Bool
otherwise =
              [String -> Int -> EntryPointType
I.TypeOpaque String
desc (Int -> EntryPointType) -> Int -> EntryPointType
forall a b. (a -> b) -> a -> b
$ t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length t a
ts]
          where desc :: String
desc = String
-> (TypeExp VName -> String) -> Maybe (TypeExp VName) -> String
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (TypeBase () () -> String
forall a. Pretty a => a -> String
pretty TypeBase () ()
t') TypeExp VName -> String
forall vn. (Eq vn, IsName vn) => TypeExp vn -> String
typeExpOpaqueName (Maybe (TypeExp VName) -> String)
-> Maybe (TypeExp VName) -> String
forall a b. (a -> b) -> a -> b
$ EntryType -> Maybe (TypeExp VName)
E.entryAscribed EntryType
t
                t' :: TypeBase () ()
t' = StructType -> TypeBase () ()
forall vn as. TypeBase (DimDecl vn) as -> TypeBase () as
noSizes (EntryType -> StructType
E.entryType EntryType
t) TypeBase () () -> Uniqueness -> TypeBase () ()
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`E.setUniqueness` Uniqueness
Nonunique

        -- | We remove dimension arguments such that we hopefully end
        -- up with a simpler type name for the entry point.  The
        -- intent is that if an entry point uses a type 'nasty [w] [h]',
        -- then we should turn that into an opaque type just called
        -- 'nasty'.  Also, we try to give arrays of opaques a nicer name.
        typeExpOpaqueName :: TypeExp vn -> String
typeExpOpaqueName (TEApply TypeExp vn
te TypeArgExpDim{} SrcLoc
_) =
          TypeExp vn -> String
typeExpOpaqueName TypeExp vn
te
        typeExpOpaqueName (TEArray TypeExp vn
te DimExp vn
_ SrcLoc
_) =
          let (Int
d, TypeExp vn
te') = TypeExp vn -> (Int, TypeExp vn)
forall vn. TypeExp vn -> (Int, TypeExp vn)
withoutDims TypeExp vn
te
          in String
"arr_" String -> String -> String
forall a. [a] -> [a] -> [a]
++ TypeExp vn -> String
typeExpOpaqueName TypeExp vn
te' String -> String -> String
forall a. [a] -> [a] -> [a]
++
             String
"_" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
d) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"d"
        typeExpOpaqueName TypeExp vn
te = TypeExp vn -> String
forall a. Pretty a => a -> String
pretty TypeExp vn
te

        withoutDims :: TypeExp vn -> (Int, TypeExp vn)
withoutDims (TEArray TypeExp vn
te DimExp vn
_ SrcLoc
_) =
          let (Int
d, TypeExp vn
te') = TypeExp vn -> (Int, TypeExp vn)
withoutDims TypeExp vn
te
          in (Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1, TypeExp vn
te')
        withoutDims TypeExp vn
te = (Int
0::Int, TypeExp vn
te)

internaliseIdent :: E.Ident -> InternaliseM I.VName
internaliseIdent :: Ident -> InternaliseM VName
internaliseIdent (E.Ident VName
name (Info PatternType
tp) SrcLoc
loc) =
  case PatternType
tp of
    E.Scalar E.Prim{} -> VName -> InternaliseM VName
forall (m :: * -> *) a. Monad m => a -> m a
return VName
name
    PatternType
_ -> String -> InternaliseM VName
forall a. HasCallStack => String -> a
error (String -> InternaliseM VName) -> String -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$ String
"Futhark.Internalise.internaliseIdent: asked to internalise non-prim-typed ident '"
         String -> String -> String
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Pretty a => a -> String
pretty VName
name String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" of type " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PatternType -> String
forall a. Pretty a => a -> String
pretty PatternType
tp String -> String -> String
forall a. [a] -> [a] -> [a]
++
         String
" at " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SrcLoc -> String
forall a. Located a => a -> String
locStr SrcLoc
loc String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"."

internaliseBody :: E.Exp -> InternaliseM Body
internaliseBody :: ExpBase Info VName -> InternaliseM Body
internaliseBody ExpBase Info VName
e = InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
m (Body (Lore m)) -> m (Body (Lore m))
insertStmsM (InternaliseM (Body (Lore InternaliseM))
 -> InternaliseM (Body (Lore InternaliseM)))
-> InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody ([SubExp] -> Body) -> InternaliseM [SubExp] -> InternaliseM Body
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"res" ExpBase Info VName
e

bodyFromStms :: InternaliseM (Result, a)
             -> InternaliseM (Body, a)
bodyFromStms :: InternaliseM ([SubExp], a) -> InternaliseM (Body, a)
bodyFromStms InternaliseM ([SubExp], a)
m = do
  (([SubExp]
res, a
a), Stms SOACS
stms) <- InternaliseM ([SubExp], a)
-> InternaliseM (([SubExp], a), Stms (Lore InternaliseM))
forall (m :: * -> *) a.
MonadBinder m =>
m a -> m (a, Stms (Lore m))
collectStms InternaliseM ([SubExp], a)
m
  (,a
a) (Body -> (Body, a)) -> InternaliseM Body -> InternaliseM (Body, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Stms (Lore InternaliseM)
-> [SubExp] -> InternaliseM (Body (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
Stms (Lore m) -> [SubExp] -> m (Body (Lore m))
mkBodyM Stms (Lore InternaliseM)
Stms SOACS
stms [SubExp]
res

internaliseExp :: String -> E.Exp -> InternaliseM [I.SubExp]

internaliseExp :: String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc (E.Parens ExpBase Info VName
e SrcLoc
_) =
  String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
e

internaliseExp String
desc (E.QualParens (QualName VName, SrcLoc)
_ ExpBase Info VName
e SrcLoc
_) =
  String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
e

internaliseExp String
desc (E.StringLit [Word8]
vs SrcLoc
_) =
  (SubExp -> [SubExp])
-> InternaliseM SubExp -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> [SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InternaliseM SubExp -> InternaliseM [SubExp])
-> InternaliseM SubExp -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
desc (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
  BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Type -> BasicOp
I.ArrayLit ((Word8 -> SubExp) -> [Word8] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map Word8 -> SubExp
forall v. IsValue v => v -> SubExp
constant [Word8]
vs) (Type -> BasicOp) -> Type -> BasicOp
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int8

internaliseExp String
_ (E.Var (E.QualName Closure
_ VName
name) (Info PatternType
t) SrcLoc
loc) = do
  Maybe [SubExp]
subst <- VName -> InternaliseM (Maybe [SubExp])
lookupSubst VName
name
  case Maybe [SubExp]
subst of
    Just [SubExp]
substs -> [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [SubExp]
substs
    Maybe [SubExp]
Nothing     -> do
      -- If this identifier is the name of a constant, we have to turn it
      -- into a call to the corresponding function.
      Maybe [SubExp]
is_const <- VName -> InternaliseM (Maybe [SubExp])
lookupConst VName
name
      case Maybe [SubExp]
is_const of
        Just [SubExp]
ses -> [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [SubExp]
ses
        Maybe [SubExp]
Nothing -> (SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
:[]) (SubExp -> [SubExp]) -> (VName -> SubExp) -> VName -> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
I.Var (VName -> [SubExp]) -> InternaliseM VName -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ident -> InternaliseM VName
internaliseIdent (VName -> Info PatternType -> SrcLoc -> Ident
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> IdentBase f vn
E.Ident VName
name (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t) SrcLoc
loc)

internaliseExp String
desc (E.Index ExpBase Info VName
e [DimIndexBase Info VName]
idxs (Info PatternType
ret, Info Closure
retext) SrcLoc
loc) = do
  Closure
vs <- String -> ExpBase Info VName -> InternaliseM Closure
internaliseExpToVars String
"indexed" ExpBase Info VName
e
  [SubExp]
dims <- case Closure
vs of []  -> [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [] -- Will this happen?
                     VName
v:Closure
_ -> Type -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
I.arrayDims (Type -> [SubExp]) -> InternaliseM Type -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
v
  ([DimIndex SubExp]
idxs', Certificates
cs) <- SrcLoc
-> [SubExp]
-> [DimIndexBase Info VName]
-> InternaliseM ([DimIndex SubExp], Certificates)
internaliseSlice SrcLoc
loc [SubExp]
dims [DimIndexBase Info VName]
idxs
  let index :: VName -> InternaliseM (ExpT SOACS)
index VName
v = do Type
v_t <- VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
v
                   ExpT SOACS -> InternaliseM (ExpT SOACS)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpT SOACS -> InternaliseM (ExpT SOACS))
-> ExpT SOACS -> InternaliseM (ExpT SOACS)
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ VName -> [DimIndex SubExp] -> BasicOp
I.Index VName
v ([DimIndex SubExp] -> BasicOp) -> [DimIndex SubExp] -> BasicOp
forall a b. (a -> b) -> a -> b
$ Type -> [DimIndex SubExp] -> [DimIndex SubExp]
fullSlice Type
v_t [DimIndex SubExp]
idxs'
  [SubExp]
ses <- Certificates -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
cs (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String -> [Exp (Lore InternaliseM)] -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> [Exp (Lore m)] -> m [SubExp]
letSubExps String
desc ([ExpT SOACS] -> InternaliseM [SubExp])
-> InternaliseM [ExpT SOACS] -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (VName -> InternaliseM (ExpT SOACS))
-> Closure -> InternaliseM [ExpT SOACS]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM (ExpT SOACS)
index Closure
vs
  StructType -> Closure -> [SubExp] -> InternaliseM ()
bindExtSizes (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret) Closure
retext [SubExp]
ses
  [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [SubExp]
ses

-- XXX: we map empty records and tuples to bools, because otherwise
-- arrays of unit will lose their sizes.
internaliseExp String
_ (E.TupLit [] SrcLoc
_) =
  [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
True]
internaliseExp String
_ (E.RecordLit [] SrcLoc
_) =
  [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
True]

internaliseExp String
desc (E.TupLit [ExpBase Info VName]
es SrcLoc
_) = [[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[SubExp]] -> [SubExp])
-> InternaliseM [[SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ExpBase Info VName -> InternaliseM [SubExp])
-> [ExpBase Info VName] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc) [ExpBase Info VName]
es

internaliseExp String
desc (E.RecordLit [FieldBase Info VName]
orig_fields SrcLoc
_) =
  ((Name, [SubExp]) -> [SubExp]) -> [(Name, [SubExp])] -> [SubExp]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Name, [SubExp]) -> [SubExp]
forall a b. (a, b) -> b
snd ([(Name, [SubExp])] -> [SubExp])
-> ([Map Name [SubExp]] -> [(Name, [SubExp])])
-> [Map Name [SubExp]]
-> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name [SubExp] -> [(Name, [SubExp])]
forall a. Map Name a -> [(Name, a)]
sortFields (Map Name [SubExp] -> [(Name, [SubExp])])
-> ([Map Name [SubExp]] -> Map Name [SubExp])
-> [Map Name [SubExp]]
-> [(Name, [SubExp])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Map Name [SubExp]] -> Map Name [SubExp]
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
M.unions ([Map Name [SubExp]] -> [SubExp])
-> InternaliseM [Map Name [SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (FieldBase Info VName -> InternaliseM (Map Name [SubExp]))
-> [FieldBase Info VName] -> InternaliseM [Map Name [SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM FieldBase Info VName -> InternaliseM (Map Name [SubExp])
internaliseField [FieldBase Info VName]
orig_fields
  where internaliseField :: FieldBase Info VName -> InternaliseM (Map Name [SubExp])
internaliseField (E.RecordFieldExplicit Name
name ExpBase Info VName
e SrcLoc
_) =
          Name -> [SubExp] -> Map Name [SubExp]
forall k a. k -> a -> Map k a
M.singleton Name
name ([SubExp] -> Map Name [SubExp])
-> InternaliseM [SubExp] -> InternaliseM (Map Name [SubExp])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
e
        internaliseField (E.RecordFieldImplicit VName
name Info PatternType
t SrcLoc
loc) =
          FieldBase Info VName -> InternaliseM (Map Name [SubExp])
internaliseField (FieldBase Info VName -> InternaliseM (Map Name [SubExp]))
-> FieldBase Info VName -> InternaliseM (Map Name [SubExp])
forall a b. (a -> b) -> a -> b
$ Name -> ExpBase Info VName -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
E.RecordFieldExplicit (VName -> Name
baseName VName
name)
          (QualName VName -> Info PatternType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
E.Var (VName -> QualName VName
forall v. v -> QualName v
E.qualName VName
name) Info PatternType
t SrcLoc
loc) SrcLoc
loc

internaliseExp String
desc (E.ArrayLit [ExpBase Info VName]
es (Info PatternType
arr_t) SrcLoc
loc)
  -- If this is a multidimensional array literal of primitives, we
  -- treat it specially by flattening it out followed by a reshape.
  -- This cuts down on the amount of statements that are produced, and
  -- thus allows us to efficiently handle huge array literals - a
  -- corner case, but an important one.
  | Just (([Int]
eshape,[ExpBase Info VName]
e'):[([Int], [ExpBase Info VName])]
es') <- (ExpBase Info VName -> Maybe ([Int], [ExpBase Info VName]))
-> [ExpBase Info VName] -> Maybe [([Int], [ExpBase Info VName])]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ExpBase Info VName -> Maybe ([Int], [ExpBase Info VName])
isArrayLiteral [ExpBase Info VName]
es,
    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [Int] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
eshape,
    (([Int], [ExpBase Info VName]) -> Bool)
-> [([Int], [ExpBase Info VName])] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (([Int]
eshape[Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
==) ([Int] -> Bool)
-> (([Int], [ExpBase Info VName]) -> [Int])
-> ([Int], [ExpBase Info VName])
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Int], [ExpBase Info VName]) -> [Int]
forall a b. (a, b) -> a
fst) [([Int], [ExpBase Info VName])]
es',
    Just PatternType
basetype <- Int -> PatternType -> Maybe PatternType
forall dim as. Int -> TypeBase dim as -> Maybe (TypeBase dim as)
E.peelArray ([Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
eshape) PatternType
arr_t = do
      let flat_lit :: ExpBase Info VName
flat_lit = [ExpBase Info VName]
-> Info PatternType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
[ExpBase f vn] -> f PatternType -> SrcLoc -> ExpBase f vn
E.ArrayLit ([ExpBase Info VName]
e' [ExpBase Info VName]
-> [ExpBase Info VName] -> [ExpBase Info VName]
forall a. [a] -> [a] -> [a]
++ (([Int], [ExpBase Info VName]) -> [ExpBase Info VName])
-> [([Int], [ExpBase Info VName])] -> [ExpBase Info VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Int], [ExpBase Info VName]) -> [ExpBase Info VName]
forall a b. (a, b) -> b
snd [([Int], [ExpBase Info VName])]
es') (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
basetype) SrcLoc
loc
          new_shape :: [Int]
new_shape = [ExpBase Info VName] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ExpBase Info VName]
esInt -> [Int] -> [Int]
forall a. a -> [a] -> [a]
:[Int]
eshape
      Closure
flat_arrs <- String -> ExpBase Info VName -> InternaliseM Closure
internaliseExpToVars String
"flat_literal" ExpBase Info VName
flat_lit
      Closure -> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM Closure
flat_arrs ((VName -> InternaliseM SubExp) -> InternaliseM [SubExp])
-> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
flat_arr -> do
        Type
flat_arr_t <- VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
flat_arr
        let new_shape' :: ShapeChange SubExp
new_shape' = ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp
reshapeOuter ((Int -> DimChange SubExp) -> [Int] -> ShapeChange SubExp
forall a b. (a -> b) -> [a] -> [b]
map (SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimNew (SubExp -> DimChange SubExp)
-> (Int -> SubExp) -> Int -> DimChange SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> SubExp
forall v. IsValue v => v -> SubExp
constant) [Int]
new_shape)
                         Int
1 (Shape -> ShapeChange SubExp) -> Shape -> ShapeChange SubExp
forall a b. (a -> b) -> a -> b
$ Type -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape Type
flat_arr_t
        String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
desc (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
I.Reshape ShapeChange SubExp
new_shape' VName
flat_arr

  | Bool
otherwise = do
      [DeclExtType]
arr_t_ext <- StructType -> InternaliseM [DeclExtType]
internaliseReturnType (StructType -> InternaliseM [DeclExtType])
-> StructType -> InternaliseM [DeclExtType]
forall a b. (a -> b) -> a -> b
$ PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
arr_t
      [[SubExp]]
es' <- (ExpBase Info VName -> InternaliseM [SubExp])
-> [ExpBase Info VName] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"arr_elem") [ExpBase Info VName]
es

      [Type]
rowtypes <-
        case (DeclExtType -> Maybe Type) -> [DeclExtType] -> Maybe [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Type -> Type) -> Maybe Type -> Maybe Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
rowType (Maybe Type -> Maybe Type)
-> (DeclExtType -> Maybe Type) -> DeclExtType -> Maybe Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExtType -> Maybe Type
hasStaticShape (ExtType -> Maybe Type)
-> (DeclExtType -> ExtType) -> DeclExtType -> Maybe Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DeclExtType -> ExtType
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl) [DeclExtType]
arr_t_ext of
          Just [Type]
ts -> [Type] -> InternaliseM [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Type]
ts
          Maybe [Type]
Nothing ->
            -- XXX: the monomorphiser may create single-element array
            -- literals with an unknown row type.  In those cases we
            -- need to look at the types of the actual elements.
            -- Fixing this in the monomorphiser is a lot more tricky
            -- than just working around it here.
            case [[SubExp]]
es' of
              [] -> String -> InternaliseM [Type]
forall a. HasCallStack => String -> a
error (String -> InternaliseM [Type]) -> String -> InternaliseM [Type]
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp ArrayLit: existential type: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PatternType -> String
forall a. Pretty a => a -> String
pretty PatternType
arr_t
              [SubExp]
e':[[SubExp]]
_ -> (SubExp -> InternaliseM Type) -> [SubExp] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType [SubExp]
e'

      let arraylit :: [SubExp] -> Type -> InternaliseM (ExpT SOACS)
arraylit [SubExp]
ks Type
rt = do
            [SubExp]
ks' <- (SubExp -> InternaliseM SubExp)
-> [SubExp] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (ErrorMsg SubExp
-> SrcLoc -> Type -> String -> SubExp -> InternaliseM SubExp
ensureShape
                         ErrorMsg SubExp
"shape of element differs from shape of first element"
                         SrcLoc
loc Type
rt String
"elem_reshaped") [SubExp]
ks
            ExpT SOACS -> InternaliseM (ExpT SOACS)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpT SOACS -> InternaliseM (ExpT SOACS))
-> ExpT SOACS -> InternaliseM (ExpT SOACS)
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Type -> BasicOp
I.ArrayLit [SubExp]
ks' Type
rt

      String -> [Exp (Lore InternaliseM)] -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> [Exp (Lore m)] -> m [SubExp]
letSubExps String
desc ([ExpT SOACS] -> InternaliseM [SubExp])
-> InternaliseM [ExpT SOACS] -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<
        if [[SubExp]] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[SubExp]]
es'
        then (Type -> InternaliseM (ExpT SOACS))
-> [Type] -> InternaliseM [ExpT SOACS]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([SubExp] -> Type -> InternaliseM (ExpT SOACS)
arraylit []) [Type]
rowtypes
        else ([SubExp] -> Type -> InternaliseM (ExpT SOACS))
-> [[SubExp]] -> [Type] -> InternaliseM [ExpT SOACS]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM [SubExp] -> Type -> InternaliseM (ExpT SOACS)
arraylit ([[SubExp]] -> [[SubExp]]
forall a. [[a]] -> [[a]]
transpose [[SubExp]]
es') [Type]
rowtypes

  where isArrayLiteral :: E.Exp -> Maybe ([Int],[E.Exp])
        isArrayLiteral :: ExpBase Info VName -> Maybe ([Int], [ExpBase Info VName])
isArrayLiteral (E.ArrayLit [ExpBase Info VName]
inner_es Info PatternType
_ SrcLoc
_) = do
          ([Int]
eshape,[ExpBase Info VName]
e):[([Int], [ExpBase Info VName])]
inner_es' <- (ExpBase Info VName -> Maybe ([Int], [ExpBase Info VName]))
-> [ExpBase Info VName] -> Maybe [([Int], [ExpBase Info VName])]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ExpBase Info VName -> Maybe ([Int], [ExpBase Info VName])
isArrayLiteral [ExpBase Info VName]
inner_es
          Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ (([Int], [ExpBase Info VName]) -> Bool)
-> [([Int], [ExpBase Info VName])] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (([Int]
eshape[Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
==) ([Int] -> Bool)
-> (([Int], [ExpBase Info VName]) -> [Int])
-> ([Int], [ExpBase Info VName])
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Int], [ExpBase Info VName]) -> [Int]
forall a b. (a, b) -> a
fst) [([Int], [ExpBase Info VName])]
inner_es'
          ([Int], [ExpBase Info VName])
-> Maybe ([Int], [ExpBase Info VName])
forall (m :: * -> *) a. Monad m => a -> m a
return ([ExpBase Info VName] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ExpBase Info VName]
inner_esInt -> [Int] -> [Int]
forall a. a -> [a] -> [a]
:[Int]
eshape, [ExpBase Info VName]
e [ExpBase Info VName]
-> [ExpBase Info VName] -> [ExpBase Info VName]
forall a. [a] -> [a] -> [a]
++ (([Int], [ExpBase Info VName]) -> [ExpBase Info VName])
-> [([Int], [ExpBase Info VName])] -> [ExpBase Info VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Int], [ExpBase Info VName]) -> [ExpBase Info VName]
forall a b. (a, b) -> b
snd [([Int], [ExpBase Info VName])]
inner_es')
        isArrayLiteral ExpBase Info VName
e =
          ([Int], [ExpBase Info VName])
-> Maybe ([Int], [ExpBase Info VName])
forall a. a -> Maybe a
Just ([], [ExpBase Info VName
e])

internaliseExp String
desc (E.Range ExpBase Info VName
start Maybe (ExpBase Info VName)
maybe_second Inclusiveness (ExpBase Info VName)
end (Info PatternType
ret, Info Closure
retext) SrcLoc
loc) = do
  SubExp
start' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"range_start" ExpBase Info VName
start
  SubExp
end' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"range_end" (ExpBase Info VName -> InternaliseM SubExp)
-> ExpBase Info VName -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ case Inclusiveness (ExpBase Info VName)
end of
    DownToExclusive ExpBase Info VName
e -> ExpBase Info VName
e
    ToInclusive ExpBase Info VName
e -> ExpBase Info VName
e
    UpToExclusive ExpBase Info VName
e -> ExpBase Info VName
e
  Maybe SubExp
maybe_second' <-
    (ExpBase Info VName -> InternaliseM SubExp)
-> Maybe (ExpBase Info VName) -> InternaliseM (Maybe SubExp)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"range_second") Maybe (ExpBase Info VName)
maybe_second

  -- Construct an error message in case the range is invalid.
  let conv :: SubExp -> InternaliseM SubExp
conv = case ExpBase Info VName -> PatternType
E.typeOf ExpBase Info VName
start of
               E.Scalar (E.Prim (E.Unsigned IntType
_)) -> IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
IntType -> SubExp -> m SubExp
asIntZ IntType
Int32
               PatternType
_ -> IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int32
  SubExp
start'_i32 <- SubExp -> InternaliseM SubExp
conv SubExp
start'
  SubExp
end'_i32 <- SubExp -> InternaliseM SubExp
conv SubExp
end'
  Maybe SubExp
maybe_second'_i32 <- (SubExp -> InternaliseM SubExp)
-> Maybe SubExp -> InternaliseM (Maybe SubExp)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse SubExp -> InternaliseM SubExp
conv Maybe SubExp
maybe_second'
  let errmsg :: ErrorMsg SubExp
errmsg =
        [ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg ([ErrorMsgPart SubExp] -> ErrorMsg SubExp)
-> [ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a b. (a -> b) -> a -> b
$
        [ErrorMsgPart SubExp
"Range "] [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++
        [SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt32 SubExp
start'_i32] [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++
        (case Maybe SubExp
maybe_second'_i32 of
           Maybe SubExp
Nothing -> []
           Just SubExp
second_i32 -> [ErrorMsgPart SubExp
"..", SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt32 SubExp
second_i32]) [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++
        (case Inclusiveness (ExpBase Info VName)
end of
           DownToExclusive{} -> [ErrorMsgPart SubExp
"..>"]
           ToInclusive{} -> [ErrorMsgPart SubExp
"..."]
           UpToExclusive{} -> [ErrorMsgPart SubExp
"..<"]) [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++
        [SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt32 SubExp
end'_i32, ErrorMsgPart SubExp
" is invalid."]

  (IntType
it, CmpOp
le_op, CmpOp
lt_op) <-
    case ExpBase Info VName -> PatternType
E.typeOf ExpBase Info VName
start of
      E.Scalar (E.Prim (E.Signed IntType
it)) -> (IntType, CmpOp, CmpOp) -> InternaliseM (IntType, CmpOp, CmpOp)
forall (m :: * -> *) a. Monad m => a -> m a
return (IntType
it, IntType -> CmpOp
CmpSle IntType
it, IntType -> CmpOp
CmpSlt IntType
it)
      E.Scalar (E.Prim (E.Unsigned IntType
it)) -> (IntType, CmpOp, CmpOp) -> InternaliseM (IntType, CmpOp, CmpOp)
forall (m :: * -> *) a. Monad m => a -> m a
return (IntType
it, IntType -> CmpOp
CmpUle IntType
it, IntType -> CmpOp
CmpUlt IntType
it)
      PatternType
start_t -> String -> InternaliseM (IntType, CmpOp, CmpOp)
forall a. HasCallStack => String -> a
error (String -> InternaliseM (IntType, CmpOp, CmpOp))
-> String -> InternaliseM (IntType, CmpOp, CmpOp)
forall a b. (a -> b) -> a -> b
$ String
"Start value in range has type " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PatternType -> String
forall a. Pretty a => a -> String
pretty PatternType
start_t

  let one :: SubExp
one = IntType -> Integer -> SubExp
intConst IntType
it Integer
1
      negone :: SubExp
negone = IntType -> Integer -> SubExp
intConst IntType
it (-Integer
1)
      default_step :: SubExp
default_step = case Inclusiveness (ExpBase Info VName)
end of DownToExclusive{} -> SubExp
negone
                                 ToInclusive{} -> SubExp
one
                                 UpToExclusive{} -> SubExp
one

  (SubExp
step, SubExp
step_zero) <- case Maybe SubExp
maybe_second' of
    Just SubExp
second' -> do
      SubExp
subtracted_step <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"subtracted_step" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
        BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
it Overflow
I.OverflowWrap) SubExp
second' SubExp
start'
      SubExp
step_zero <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"step_zero" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
start' SubExp
second'
      (SubExp, SubExp) -> InternaliseM (SubExp, SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
subtracted_step, SubExp
step_zero)
    Maybe SubExp
Nothing ->
      (SubExp, SubExp) -> InternaliseM (SubExp, SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
default_step, Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
False)

  SubExp
step_sign <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"s_sign" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (IntType -> UnOp
I.SSignum IntType
it) SubExp
step
  SubExp
step_sign_i32 <- IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int32 SubExp
step_sign

  SubExp
bounds_invalid_downwards <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"bounds_invalid_downwards" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
                              BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp CmpOp
le_op SubExp
start' SubExp
end'
  SubExp
bounds_invalid_upwards <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"bounds_invalid_upwards" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
                            BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp CmpOp
lt_op SubExp
end' SubExp
start'

  (SubExp
distance, SubExp
step_wrong_dir, SubExp
bounds_invalid) <- case Inclusiveness (ExpBase Info VName)
end of
    DownToExclusive{} -> do
      SubExp
step_wrong_dir <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"step_wrong_dir" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
                        BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
one
      SubExp
distance <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"distance" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
                  BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
it Overflow
I.OverflowWrap) SubExp
start' SubExp
end'
      SubExp
distance_i32 <- IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
IntType -> SubExp -> m SubExp
asIntZ IntType
Int32 SubExp
distance
      (SubExp, SubExp, SubExp) -> InternaliseM (SubExp, SubExp, SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
distance_i32, SubExp
step_wrong_dir, SubExp
bounds_invalid_downwards)
    UpToExclusive{} -> do
      SubExp
step_wrong_dir <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"step_wrong_dir" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
                        BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
negone
      SubExp
distance <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"distance" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
it Overflow
I.OverflowWrap) SubExp
end' SubExp
start'
      SubExp
distance_i32 <- IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
IntType -> SubExp -> m SubExp
asIntZ IntType
Int32 SubExp
distance
      (SubExp, SubExp, SubExp) -> InternaliseM (SubExp, SubExp, SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
distance_i32, SubExp
step_wrong_dir, SubExp
bounds_invalid_upwards)
    ToInclusive{} -> do
      SubExp
downwards <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"downwards" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
                   BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
negone
      SubExp
distance_downwards_exclusive <-
        String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"distance_downwards_exclusive" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
        BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
it Overflow
I.OverflowWrap) SubExp
start' SubExp
end'
      SubExp
distance_upwards_exclusive <-
        String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"distance_upwards_exclusive" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
        BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
it Overflow
I.OverflowWrap) SubExp
end' SubExp
start'

      SubExp
bounds_invalid <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"bounds_invalid" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
                        SubExp -> Body -> Body -> IfDec (BranchType SOACS) -> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If SubExp
downwards
                        ([SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [SubExp
bounds_invalid_downwards])
                        ([SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [SubExp
bounds_invalid_upwards]) (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
                        [Type] -> IfDec ExtType
ifCommon [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.Bool]
      SubExp
distance_exclusive <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"distance_exclusive" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
                            SubExp -> Body -> Body -> IfDec (BranchType SOACS) -> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If SubExp
downwards
                            ([SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [SubExp
distance_downwards_exclusive])
                            ([SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [SubExp
distance_upwards_exclusive]) (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
                            [Type] -> IfDec ExtType
ifCommon [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim (PrimType -> Type) -> PrimType -> Type
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it]
      SubExp
distance_exclusive_i32 <- IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
IntType -> SubExp -> m SubExp
asIntZ IntType
Int32 SubExp
distance_exclusive
      SubExp
distance <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"distance" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
                  BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Add IntType
Int32 Overflow
I.OverflowWrap)
                  SubExp
distance_exclusive_i32 (IntType -> Integer -> SubExp
intConst IntType
Int32 Integer
1)
      (SubExp, SubExp, SubExp) -> InternaliseM (SubExp, SubExp, SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
distance, Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
False, SubExp
bounds_invalid)

  SubExp
step_invalid <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"step_invalid" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
                  BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogOr SubExp
step_wrong_dir SubExp
step_zero

  SubExp
invalid <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"range_invalid" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
             BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogOr SubExp
step_invalid SubExp
bounds_invalid
  SubExp
valid <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"valid" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
I.Not SubExp
invalid
  Certificates
cs <- String
-> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certificates
assert String
"range_valid_c" SubExp
valid ErrorMsg SubExp
errmsg SrcLoc
loc

  SubExp
step_i32 <- IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int32 SubExp
step
  SubExp
pos_step <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"pos_step" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
              BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Mul IntType
Int32 Overflow
I.OverflowWrap) SubExp
step_i32 SubExp
step_sign_i32

  SubExp
num_elems <- Certificates -> InternaliseM SubExp -> InternaliseM SubExp
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
cs (InternaliseM SubExp -> InternaliseM SubExp)
-> InternaliseM SubExp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
               String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"num_elems" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
               BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Safety -> BinOp
SDivUp IntType
Int32 Safety
I.Unsafe) SubExp
distance SubExp
pos_step

  SubExp
se <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
desc (BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> SubExp -> SubExp -> IntType -> BasicOp
I.Iota SubExp
num_elems SubExp
start' SubExp
step IntType
it)
  StructType -> Closure -> [SubExp] -> InternaliseM ()
bindExtSizes (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret) Closure
retext [SubExp
se]
  [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [SubExp
se]

internaliseExp String
desc (E.Ascript ExpBase Info VName
e TypeDeclBase Info VName
_ SrcLoc
_) =
  String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
e

internaliseExp String
desc (E.Coerce ExpBase Info VName
e (TypeDecl TypeExp VName
dt (Info StructType
et)) (Info PatternType
ret, Info Closure
retext) SrcLoc
loc) = do
  [SubExp]
ses <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
e
  [DeclExtType]
ts <- StructType -> InternaliseM [DeclExtType]
internaliseReturnType StructType
et
  [ErrorMsgPart SubExp]
dt' <- TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
dt
  StructType -> Closure -> [SubExp] -> InternaliseM ()
bindExtSizes (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret) Closure
retext [SubExp]
ses
  [(SubExp, DeclExtType)]
-> ((SubExp, DeclExtType) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([SubExp] -> [DeclExtType] -> [(SubExp, DeclExtType)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
ses [DeclExtType]
ts) (((SubExp, DeclExtType) -> InternaliseM SubExp)
 -> InternaliseM [SubExp])
-> ((SubExp, DeclExtType) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \(SubExp
e',DeclExtType
t') -> do
    [SubExp]
dims <- Type -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
arrayDims (Type -> [SubExp]) -> InternaliseM Type -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
e'
    let parts :: [ErrorMsgPart SubExp]
parts = [ErrorMsgPart SubExp
"Value of (core language) shape ("] [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++
                ErrorMsgPart SubExp
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. a -> [a] -> [a]
intersperse ErrorMsgPart SubExp
", " ((SubExp -> ErrorMsgPart SubExp)
-> [SubExp] -> [ErrorMsgPart SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt32 [SubExp]
dims) [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++
                [ErrorMsgPart SubExp
") cannot match shape of type `"] [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
dt' [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"`."]
    ErrorMsg SubExp
-> SrcLoc -> ExtType -> String -> SubExp -> InternaliseM SubExp
ensureExtShape ([ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [ErrorMsgPart SubExp]
parts) SrcLoc
loc (DeclExtType -> ExtType
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl DeclExtType
t') String
desc SubExp
e'

internaliseExp String
desc (E.Negate ExpBase Info VName
e SrcLoc
_) = do
  SubExp
e' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"negate_arg" ExpBase Info VName
e
  Type
et <- SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
e'
  case Type
et of I.Prim (I.IntType IntType
t) ->
               String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
t Overflow
I.OverflowWrap) (IntType -> Integer -> SubExp
I.intConst IntType
t Integer
0) SubExp
e'
             I.Prim (I.FloatType FloatType
t) ->
               String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (FloatType -> BinOp
I.FSub FloatType
t) (FloatType -> Double -> SubExp
I.floatConst FloatType
t Double
0) SubExp
e'
             Type
_ -> String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error String
"Futhark.Internalise.internaliseExp: non-numeric type in Negate"

internaliseExp String
desc e :: ExpBase Info VName
e@E.Apply{} = do
  (QualName VName
qfname, [(ExpBase Info VName, Maybe VName)]
args, StructType
ret, Closure
retext) <- ExpBase Info VName
-> InternaliseM
     (QualName VName, [(ExpBase Info VName, Maybe VName)], StructType,
      Closure)
findFuncall ExpBase Info VName
e
  -- Argument evaluation is outermost-in so that any existential sizes
  -- created by function applications can be brought into scope.
  let fname :: Name
fname = String -> Name
nameFromString (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ Name -> String
forall a. Pretty a => a -> String
pretty (Name -> String) -> Name -> String
forall a b. (a -> b) -> a -> b
$ VName -> Name
baseName (VName -> Name) -> VName -> Name
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname
      loc :: SrcLoc
loc = ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
e
      arg_desc :: String
arg_desc = Name -> String
nameToString Name
fname String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_arg"

  -- Some functions are magical (overloaded) and we handle that here.
  [SubExp]
ses <-
    case () of
      -- Overloaded functions never take array arguments (except
      -- equality, but those cannot be existential), so we can safely
      -- ignore the existential dimensions.
      () | Just String -> InternaliseM [SubExp]
internalise <- QualName VName
-> [ExpBase Info VName]
-> SrcLoc
-> Maybe (String -> InternaliseM [SubExp])
isOverloadedFunction QualName VName
qfname (((ExpBase Info VName, Maybe VName) -> ExpBase Info VName)
-> [(ExpBase Info VName, Maybe VName)] -> [ExpBase Info VName]
forall a b. (a -> b) -> [a] -> [b]
map (ExpBase Info VName, Maybe VName) -> ExpBase Info VName
forall a b. (a, b) -> a
fst [(ExpBase Info VName, Maybe VName)]
args) SrcLoc
loc ->
             String -> InternaliseM [SubExp]
internalise String
desc

         | Just (PrimType
rettype, [PrimType]
_) <- Name
-> Map Name (PrimType, [PrimType]) -> Maybe (PrimType, [PrimType])
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
fname Map Name (PrimType, [PrimType])
I.builtInFunctions -> do
             let tag :: [a] -> [(a, Diet)]
tag [a]
ses = [ (a
se, Diet
I.Observe) | a
se <- [a]
ses ]
             [[SubExp]]
args' <- [[SubExp]] -> [[SubExp]]
forall a. [a] -> [a]
reverse ([[SubExp]] -> [[SubExp]])
-> InternaliseM [[SubExp]] -> InternaliseM [[SubExp]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((ExpBase Info VName, Maybe VName) -> InternaliseM [SubExp])
-> [(ExpBase Info VName, Maybe VName)] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String
-> (ExpBase Info VName, Maybe VName) -> InternaliseM [SubExp]
internaliseArg String
arg_desc) ([(ExpBase Info VName, Maybe VName)]
-> [(ExpBase Info VName, Maybe VName)]
forall a. [a] -> [a]
reverse [(ExpBase Info VName, Maybe VName)]
args)
             let args'' :: [(SubExp, Diet)]
args'' = ([SubExp] -> [(SubExp, Diet)]) -> [[SubExp]] -> [(SubExp, Diet)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap [SubExp] -> [(SubExp, Diet)]
forall a. [a] -> [(a, Diet)]
tag [[SubExp]]
args'
             String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Name
-> [(SubExp, Diet)]
-> [RetType SOACS]
-> (Safety, SrcLoc, [SrcLoc])
-> ExpT SOACS
forall lore.
Name
-> [(SubExp, Diet)]
-> [RetType lore]
-> (Safety, SrcLoc, [SrcLoc])
-> ExpT lore
I.Apply Name
fname [(SubExp, Diet)]
args'' [PrimType -> DeclExtType
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
rettype]
               (Safety
Safe, SrcLoc
loc, [])

         | Bool
otherwise -> do
             [SubExp]
args' <- [[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[SubExp]] -> [SubExp])
-> ([[SubExp]] -> [[SubExp]]) -> [[SubExp]] -> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[SubExp]] -> [[SubExp]]
forall a. [a] -> [a]
reverse ([[SubExp]] -> [SubExp])
-> InternaliseM [[SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((ExpBase Info VName, Maybe VName) -> InternaliseM [SubExp])
-> [(ExpBase Info VName, Maybe VName)] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String
-> (ExpBase Info VName, Maybe VName) -> InternaliseM [SubExp]
internaliseArg String
arg_desc) ([(ExpBase Info VName, Maybe VName)]
-> [(ExpBase Info VName, Maybe VName)]
forall a. [a] -> [a]
reverse [(ExpBase Info VName, Maybe VName)]
args)
             ([SubExp], [ExtType]) -> [SubExp]
forall a b. (a, b) -> a
fst (([SubExp], [ExtType]) -> [SubExp])
-> InternaliseM ([SubExp], [ExtType]) -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String
-> QualName VName
-> [SubExp]
-> SrcLoc
-> InternaliseM ([SubExp], [ExtType])
funcall String
desc QualName VName
qfname [SubExp]
args' SrcLoc
loc

  StructType -> Closure -> [SubExp] -> InternaliseM ()
bindExtSizes StructType
ret Closure
retext [SubExp]
ses
  [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [SubExp]
ses

internaliseExp String
desc (E.LetPat Pattern
pat ExpBase Info VName
e ExpBase Info VName
body (Info PatternType
ret, Info Closure
retext) SrcLoc
loc) = do
  [SubExp]
ses <- String
-> Pattern
-> ExpBase Info VName
-> ExpBase Info VName
-> SrcLoc
-> (ExpBase Info VName -> InternaliseM [SubExp])
-> InternaliseM [SubExp]
forall a.
String
-> Pattern
-> ExpBase Info VName
-> ExpBase Info VName
-> SrcLoc
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat String
desc Pattern
pat ExpBase Info VName
e ExpBase Info VName
body SrcLoc
loc (String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc)
  StructType -> Closure -> [SubExp] -> InternaliseM ()
bindExtSizes (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret) Closure
retext [SubExp]
ses
  [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [SubExp]
ses

internaliseExp String
desc (E.LetFun VName
ofname ([TypeParamBase VName]
tparams, [Pattern]
params, Maybe (TypeExp VName)
retdecl, Info StructType
rettype, ExpBase Info VName
body) ExpBase Info VName
letbody Info PatternType
_ SrcLoc
loc) = do
  ValBind -> InternaliseM ()
internaliseValBind (ValBind -> InternaliseM ()) -> ValBind -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
    Maybe (Info EntryPoint)
-> VName
-> Maybe (TypeExp VName)
-> Info (StructType, Closure)
-> [TypeParamBase VName]
-> [Pattern]
-> ExpBase Info VName
-> Maybe DocComment
-> [AttrInfo]
-> SrcLoc
-> ValBind
forall (f :: * -> *) vn.
Maybe (f EntryPoint)
-> vn
-> Maybe (TypeExp vn)
-> f (StructType, Closure)
-> [TypeParamBase vn]
-> [PatternBase f vn]
-> ExpBase f vn
-> Maybe DocComment
-> [AttrInfo]
-> SrcLoc
-> ValBindBase f vn
E.ValBind Maybe (Info EntryPoint)
forall a. Maybe a
Nothing VName
ofname Maybe (TypeExp VName)
retdecl ((StructType, Closure) -> Info (StructType, Closure)
forall a. a -> Info a
Info (StructType
rettype, [])) [TypeParamBase VName]
tparams [Pattern]
params ExpBase Info VName
body Maybe DocComment
forall a. Maybe a
Nothing [AttrInfo]
forall a. Monoid a => a
mempty SrcLoc
loc
  String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
letbody

internaliseExp String
desc (E.DoLoop Closure
sparams Pattern
mergepat ExpBase Info VName
mergeexp LoopFormBase Info VName
form ExpBase Info VName
loopbody (Info (PatternType
ret, Closure
retext)) SrcLoc
loc) = do
  [SubExp]
ses <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"loop_init" ExpBase Info VName
mergeexp
  ((Body
loopbody', (LoopForm SOACS
form', [Param DeclType]
shapepat, [Param DeclType]
mergepat', [SubExp]
mergeinit')), Stms SOACS
initstms) <-
    InternaliseM
  (Body,
   (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     ((Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])),
      Stms (Lore InternaliseM))
forall (m :: * -> *) a.
MonadBinder m =>
m a -> m (a, Stms (Lore m))
collectStms (InternaliseM
   (Body,
    (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
 -> InternaliseM
      ((Body,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])),
       Stms (Lore InternaliseM)))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     ((Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])),
      Stms (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ [SubExp]
-> LoopFormBase Info VName
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
handleForm [SubExp]
ses LoopFormBase Info VName
form

  Stms (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *). MonadBinder m => Stms (Lore m) -> m ()
addStms Stms (Lore InternaliseM)
Stms SOACS
initstms
  [Type]
mergeinit_ts' <- (SubExp -> InternaliseM Type) -> [SubExp] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType [SubExp]
mergeinit'

  let ctxinit :: [SubExp]
ctxinit = Closure -> [Type] -> [Type] -> [SubExp]
forall u0 u1.
Closure -> [TypeBase Shape u0] -> [TypeBase Shape u1] -> [SubExp]
argShapes
                ((Param DeclType -> VName) -> [Param DeclType] -> Closure
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
shapepat)
                ((Param DeclType -> Type) -> [Param DeclType] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> Type
forall dec. Typed dec => Param dec -> Type
I.paramType [Param DeclType]
mergepat')
                [Type]
mergeinit_ts'
      ctxmerge :: [(Param DeclType, SubExp)]
ctxmerge = [Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
shapepat [SubExp]
ctxinit
      valmerge :: [(Param DeclType, SubExp)]
valmerge = [Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
mergepat' [SubExp]
mergeinit'
      dropCond :: Closure -> Closure
dropCond = case LoopFormBase Info VName
form of E.While{} -> Int -> Closure -> Closure
forall a. Int -> [a] -> [a]
drop Int
1
                              LoopFormBase Info VName
_         -> Closure -> Closure
forall a. a -> a
id

  -- Ensure that the result of the loop matches the shapes of the
  -- merge parameters.  XXX: Ideally they should already match (by
  -- the source language type rules), but some of our
  -- transformations (esp. defunctionalisation) strips out some size
  -- information.  For a type-correct source program, these reshapes
  -- should simplify away.
  let merge :: [(Param DeclType, SubExp)]
merge = [(Param DeclType, SubExp)]
ctxmerge [(Param DeclType, SubExp)]
-> [(Param DeclType, SubExp)] -> [(Param DeclType, SubExp)]
forall a. [a] -> [a] -> [a]
++ [(Param DeclType, SubExp)]
valmerge
      merge_ts :: [Type]
merge_ts = ((Param DeclType, SubExp) -> Type)
-> [(Param DeclType, SubExp)] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Param DeclType -> Type
forall dec. Typed dec => Param dec -> Type
I.paramType (Param DeclType -> Type)
-> ((Param DeclType, SubExp) -> Param DeclType)
-> (Param DeclType, SubExp)
-> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param DeclType, SubExp) -> Param DeclType
forall a b. (a, b) -> a
fst) [(Param DeclType, SubExp)]
merge
  Body
loopbody'' <- Scope SOACS -> InternaliseM Body -> InternaliseM Body
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param DeclType] -> Scope SOACS
forall lore dec.
(FParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
scopeOfFParams ([Param DeclType] -> Scope SOACS)
-> [Param DeclType] -> Scope SOACS
forall a b. (a -> b) -> a -> b
$ ((Param DeclType, SubExp) -> Param DeclType)
-> [(Param DeclType, SubExp)] -> [Param DeclType]
forall a b. (a -> b) -> [a] -> [b]
map (Param DeclType, SubExp) -> Param DeclType
forall a b. (a, b) -> a
fst [(Param DeclType, SubExp)]
merge) (InternaliseM Body -> InternaliseM Body)
-> InternaliseM Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$
                LoopForm SOACS -> InternaliseM Body -> InternaliseM Body
forall lore a (m :: * -> *) b.
(Scoped lore a, LocalScope lore m) =>
a -> m b -> m b
inScopeOf LoopForm SOACS
form' (InternaliseM Body -> InternaliseM Body)
-> InternaliseM Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$ InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
m (Body (Lore m)) -> m (Body (Lore m))
insertStmsM (InternaliseM (Body (Lore InternaliseM))
 -> InternaliseM (Body (Lore InternaliseM)))
-> InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$
    [SubExp] -> InternaliseM Body
forall (m :: * -> *).
MonadBinder m =>
[SubExp] -> m (Body (Lore m))
resultBodyM
    ([SubExp] -> InternaliseM Body)
-> InternaliseM [SubExp] -> InternaliseM Body
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ErrorMsg SubExp
-> SrcLoc -> Closure -> [Type] -> [SubExp] -> InternaliseM [SubExp]
forall u.
Typed (TypeBase Shape u) =>
ErrorMsg SubExp
-> SrcLoc
-> Closure
-> [TypeBase Shape u]
-> [SubExp]
-> InternaliseM [SubExp]
ensureArgShapes
        ErrorMsg SubExp
"shape of loop result does not match shapes in loop parameter"
        SrcLoc
loc (((Param DeclType, SubExp) -> VName)
-> [(Param DeclType, SubExp)] -> Closure
forall a b. (a -> b) -> [a] -> [b]
map (Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName (Param DeclType -> VName)
-> ((Param DeclType, SubExp) -> Param DeclType)
-> (Param DeclType, SubExp)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param DeclType, SubExp) -> Param DeclType
forall a b. (a, b) -> a
fst) [(Param DeclType, SubExp)]
ctxmerge) [Type]
merge_ts
    ([SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Body (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *). MonadBinder m => Body (Lore m) -> m [SubExp]
bodyBind Body (Lore InternaliseM)
Body
loopbody'

  [SubExp]
loop_res <- (VName -> SubExp) -> Closure -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var (Closure -> [SubExp])
-> (Closure -> Closure) -> Closure -> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Closure -> Closure
dropCond (Closure -> [SubExp])
-> InternaliseM Closure -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
              String -> Exp (Lore InternaliseM) -> InternaliseM Closure
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Closure
letTupExp String
desc ([(FParam, SubExp)]
-> [(FParam, SubExp)] -> LoopForm SOACS -> Body -> ExpT SOACS
forall lore.
[(FParam lore, SubExp)]
-> [(FParam lore, SubExp)]
-> LoopForm lore
-> BodyT lore
-> ExpT lore
I.DoLoop [(Param DeclType, SubExp)]
[(FParam, SubExp)]
ctxmerge [(Param DeclType, SubExp)]
[(FParam, SubExp)]
valmerge LoopForm SOACS
form' Body
loopbody'')
  StructType -> Closure -> [SubExp] -> InternaliseM ()
bindExtSizes (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret) Closure
retext [SubExp]
loop_res
  [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [SubExp]
loop_res

  where
    sparams' :: [TypeParamBase VName]
sparams' = (VName -> TypeParamBase VName) -> Closure -> [TypeParamBase VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SrcLoc -> TypeParamBase VName
forall vn. vn -> SrcLoc -> TypeParamBase vn
`TypeParamDim` SrcLoc
forall a. Monoid a => a
mempty) Closure
sparams

    forLoop :: [[Param DeclType]]
-> [Param DeclType]
-> [SubExp]
-> LoopForm SOACS
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forLoop [[Param DeclType]]
nested_mergepat [Param DeclType]
shapepat [SubExp]
mergeinit LoopForm SOACS
form' = do
      let mergepat' :: [Param DeclType]
mergepat' = [[Param DeclType]] -> [Param DeclType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Param DeclType]]
nested_mergepat
      InternaliseM
  ([SubExp],
   (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a. InternaliseM ([SubExp], a) -> InternaliseM (Body, a)
bodyFromStms (InternaliseM
   ([SubExp],
    (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     ([SubExp],
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ LoopForm SOACS
-> InternaliseM
     ([SubExp],
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     ([SubExp],
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall lore a (m :: * -> *) b.
(Scoped lore a, LocalScope lore m) =>
a -> m b -> m b
inScopeOf LoopForm SOACS
form' (InternaliseM
   ([SubExp],
    (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
 -> InternaliseM
      ([SubExp],
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     ([SubExp],
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     ([SubExp],
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ do
        [SubExp]
ses <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"loopres" ExpBase Info VName
loopbody
        [Type]
sets <- (SubExp -> InternaliseM Type) -> [SubExp] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType [SubExp]
ses
        let shapeargs :: [SubExp]
shapeargs = Closure -> [Type] -> [Type] -> [SubExp]
forall u0 u1.
Closure -> [TypeBase Shape u0] -> [TypeBase Shape u1] -> [SubExp]
argShapes
                        ((Param DeclType -> VName) -> [Param DeclType] -> Closure
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
shapepat)
                        ((Param DeclType -> Type) -> [Param DeclType] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> Type
forall dec. Typed dec => Param dec -> Type
I.paramType [Param DeclType]
mergepat')
                        [Type]
sets
        ([SubExp],
 (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     ([SubExp],
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall (m :: * -> *) a. Monad m => a -> m a
return ([SubExp]
shapeargs [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
ses,
                (LoopForm SOACS
form',
                 [Param DeclType]
shapepat,
                 [Param DeclType]
mergepat',
                 [SubExp]
mergeinit))

    handleForm :: [SubExp]
-> LoopFormBase Info VName
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
handleForm [SubExp]
mergeinit (E.ForIn Pattern
x ExpBase Info VName
arr) = do
      Closure
arr' <- String -> ExpBase Info VName -> InternaliseM Closure
internaliseExpToVars String
"for_in_arr" ExpBase Info VName
arr
      [Type]
arr_ts <- (VName -> InternaliseM Type) -> Closure -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType Closure
arr'
      let w :: SubExp
w = Int -> [Type] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [Type]
arr_ts

      VName
i <- String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"i"

      [TypeParamBase VName]
-> [Pattern]
-> ([FParam]
    -> [[FParam]]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a.
[TypeParamBase VName]
-> [Pattern]
-> ([FParam] -> [[FParam]] -> InternaliseM a)
-> InternaliseM a
bindingParams [TypeParamBase VName]
sparams' [Pattern
mergepat] (([FParam]
  -> [[FParam]]
  -> InternaliseM
       (Body,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> ([FParam]
    -> [[FParam]]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$
        \[FParam]
shapepat [[FParam]]
nested_mergepat ->
        [Pattern]
-> [Type]
-> ([LParam]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a.
[Pattern]
-> [Type] -> ([LParam] -> InternaliseM a) -> InternaliseM a
bindingLambdaParams [Pattern
x] ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
rowType [Type]
arr_ts) (([LParam]
  -> InternaliseM
       (Body,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> ([LParam]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ \[LParam]
x_params -> do
          let loopvars :: [(Param Type, VName)]
loopvars = [Param Type] -> Closure -> [(Param Type, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param Type]
[LParam]
x_params Closure
arr'
          [[Param DeclType]]
-> [Param DeclType]
-> [SubExp]
-> LoopForm SOACS
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forLoop [[Param DeclType]]
[[FParam]]
nested_mergepat [Param DeclType]
[FParam]
shapepat [SubExp]
mergeinit (LoopForm SOACS
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> LoopForm SOACS
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$
            VName -> IntType -> SubExp -> [(LParam, VName)] -> LoopForm SOACS
forall lore.
VName
-> IntType -> SubExp -> [(LParam lore, VName)] -> LoopForm lore
I.ForLoop VName
i IntType
Int32 SubExp
w [(Param Type, VName)]
[(LParam, VName)]
loopvars

    handleForm [SubExp]
mergeinit (E.For Ident
i ExpBase Info VName
num_iterations) = do
      SubExp
num_iterations' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"upper_bound" ExpBase Info VName
num_iterations
      VName
i' <- Ident -> InternaliseM VName
internaliseIdent Ident
i
      Type
num_iterations_t <- SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
I.subExpType SubExp
num_iterations'
      IntType
it <- case Type
num_iterations_t of
              I.Prim (IntType IntType
it) -> IntType -> InternaliseM IntType
forall (m :: * -> *) a. Monad m => a -> m a
return IntType
it
              Type
_                   -> String -> InternaliseM IntType
forall a. HasCallStack => String -> a
error String
"internaliseExp DoLoop: invalid type"

      [TypeParamBase VName]
-> [Pattern]
-> ([FParam]
    -> [[FParam]]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a.
[TypeParamBase VName]
-> [Pattern]
-> ([FParam] -> [[FParam]] -> InternaliseM a)
-> InternaliseM a
bindingParams [TypeParamBase VName]
sparams' [Pattern
mergepat] (([FParam]
  -> [[FParam]]
  -> InternaliseM
       (Body,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> ([FParam]
    -> [[FParam]]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$
        \[FParam]
shapepat [[FParam]]
nested_mergepat ->
          [[Param DeclType]]
-> [Param DeclType]
-> [SubExp]
-> LoopForm SOACS
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forLoop [[Param DeclType]]
[[FParam]]
nested_mergepat [Param DeclType]
[FParam]
shapepat [SubExp]
mergeinit (LoopForm SOACS
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> LoopForm SOACS
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$
          VName -> IntType -> SubExp -> [(LParam, VName)] -> LoopForm SOACS
forall lore.
VName
-> IntType -> SubExp -> [(LParam lore, VName)] -> LoopForm lore
I.ForLoop VName
i' IntType
it SubExp
num_iterations' []

    handleForm [SubExp]
mergeinit (E.While ExpBase Info VName
cond) =
      [TypeParamBase VName]
-> [Pattern]
-> ([FParam]
    -> [[FParam]]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a.
[TypeParamBase VName]
-> [Pattern]
-> ([FParam] -> [[FParam]] -> InternaliseM a)
-> InternaliseM a
bindingParams [TypeParamBase VName]
sparams' [Pattern
mergepat] (([FParam]
  -> [[FParam]]
  -> InternaliseM
       (Body,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> ([FParam]
    -> [[FParam]]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ \[FParam]
shapepat [[FParam]]
nested_mergepat -> do
        [Type]
mergeinit_ts <- (SubExp -> InternaliseM Type) -> [SubExp] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType [SubExp]
mergeinit
        let mergepat' :: [Param DeclType]
mergepat' = [[Param DeclType]] -> [Param DeclType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Param DeclType]]
[[FParam]]
nested_mergepat
        -- We need to insert 'cond' twice - once for the initial
        -- condition (do we enter the loop at all?), and once with the
        -- result values of the loop (do we continue into the next
        -- iteration?).  This is safe, as the type rules for the
        -- external language guarantees that 'cond' does not consume
        -- anything.
        let shapeinit :: [SubExp]
shapeinit = Closure -> [Type] -> [Type] -> [SubExp]
forall u0 u1.
Closure -> [TypeBase Shape u0] -> [TypeBase Shape u1] -> [SubExp]
argShapes
                        ((Param DeclType -> VName) -> [Param DeclType] -> Closure
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
[FParam]
shapepat)
                        ((Param DeclType -> Type) -> [Param DeclType] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> Type
forall dec. Typed dec => Param dec -> Type
I.paramType [Param DeclType]
mergepat')
                        [Type]
mergeinit_ts

        (SubExp
loop_initial_cond, Stms SOACS
init_loop_cond_bnds) <- InternaliseM SubExp
-> InternaliseM (SubExp, Stms (Lore InternaliseM))
forall (m :: * -> *) a.
MonadBinder m =>
m a -> m (a, Stms (Lore m))
collectStms (InternaliseM SubExp
 -> InternaliseM (SubExp, Stms (Lore InternaliseM)))
-> InternaliseM SubExp
-> InternaliseM (SubExp, Stms (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ do
          [(Param DeclType, SubExp)]
-> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
[FParam]
shapepat [SubExp]
shapeinit) (((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ())
-> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \(Param DeclType
p, SubExp
se) ->
            Closure -> Exp (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
Closure -> Exp (Lore m) -> m ()
letBindNames [Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
          [(Param DeclType, SubExp)]
-> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([[Param DeclType]] -> [Param DeclType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Param DeclType]]
[[FParam]]
nested_mergepat) [SubExp]
mergeinit) (((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ())
-> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \(Param DeclType
p, SubExp
se) ->
            Bool -> InternaliseM () -> InternaliseM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (SubExp
se SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== VName -> SubExp
I.Var (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p)) (InternaliseM () -> InternaliseM ())
-> InternaliseM () -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
            Closure -> Exp (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
Closure -> Exp (Lore m) -> m ()
letBindNames [Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
            case SubExp
se of I.Var VName
v | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Type -> Bool
forall shape u. TypeBase shape u -> Bool
primType (Type -> Bool) -> Type -> Bool
forall a b. (a -> b) -> a -> b
$ Param DeclType -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param DeclType
p ->
                                   ShapeChange SubExp -> VName -> BasicOp
Reshape ((SubExp -> DimChange SubExp) -> [SubExp] -> ShapeChange SubExp
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimCoercion ([SubExp] -> ShapeChange SubExp) -> [SubExp] -> ShapeChange SubExp
forall a b. (a -> b) -> a -> b
$ Type -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
arrayDims (Type -> [SubExp]) -> Type -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Param DeclType -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param DeclType
p) VName
v
                       SubExp
_ -> SubExp -> BasicOp
SubExp SubExp
se
          String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"loop_cond" ExpBase Info VName
cond

        Stms (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *). MonadBinder m => Stms (Lore m) -> m ()
addStms Stms (Lore InternaliseM)
Stms SOACS
init_loop_cond_bnds

        InternaliseM
  ([SubExp],
   (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a. InternaliseM ([SubExp], a) -> InternaliseM (Body, a)
bodyFromStms (InternaliseM
   ([SubExp],
    (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     ([SubExp],
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ do
          [SubExp]
ses <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"loopres" ExpBase Info VName
loopbody
          [Type]
sets <- (SubExp -> InternaliseM Type) -> [SubExp] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType [SubExp]
ses
          Param DeclType
loop_while <- String -> DeclType -> InternaliseM (Param DeclType)
forall (m :: * -> *) dec.
MonadFreshNames m =>
String -> dec -> m (Param dec)
newParam String
"loop_while" (DeclType -> InternaliseM (Param DeclType))
-> DeclType -> InternaliseM (Param DeclType)
forall a b. (a -> b) -> a -> b
$ PrimType -> DeclType
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.Bool
          let shapeargs :: [SubExp]
shapeargs = Closure -> [Type] -> [Type] -> [SubExp]
forall u0 u1.
Closure -> [TypeBase Shape u0] -> [TypeBase Shape u1] -> [SubExp]
argShapes
                          ((Param DeclType -> VName) -> [Param DeclType] -> Closure
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
[FParam]
shapepat)
                          ((Param DeclType -> Type) -> [Param DeclType] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> Type
forall dec. Typed dec => Param dec -> Type
I.paramType [Param DeclType]
mergepat')
                          [Type]
sets

          -- Careful not to clobber anything.
          Body
loop_end_cond_body <- Body -> InternaliseM Body
forall lore (m :: * -> *).
(Renameable lore, MonadFreshNames m) =>
Body lore -> m (Body lore)
renameBody (Body -> InternaliseM Body)
-> (InternaliseM Body -> InternaliseM Body)
-> InternaliseM Body
-> InternaliseM Body
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< InternaliseM Body -> InternaliseM Body
forall (m :: * -> *).
MonadBinder m =>
m (Body (Lore m)) -> m (Body (Lore m))
insertStmsM (InternaliseM Body -> InternaliseM Body)
-> InternaliseM Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$ do
            [(Param DeclType, SubExp)]
-> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
[FParam]
shapepat [SubExp]
shapeargs) (((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ())
-> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \(Param DeclType
p, SubExp
se) ->
              Bool -> InternaliseM () -> InternaliseM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (SubExp
se SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== VName -> SubExp
I.Var (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p)) (InternaliseM () -> InternaliseM ())
-> InternaliseM () -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
              Closure -> Exp (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
Closure -> Exp (Lore m) -> m ()
letBindNames [Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
            [(Param DeclType, SubExp)]
-> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([[Param DeclType]] -> [Param DeclType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Param DeclType]]
[[FParam]]
nested_mergepat) [SubExp]
ses) (((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ())
-> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \(Param DeclType
p, SubExp
se) ->
              Bool -> InternaliseM () -> InternaliseM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (SubExp
se SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== VName -> SubExp
I.Var (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p)) (InternaliseM () -> InternaliseM ())
-> InternaliseM () -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
              Closure -> Exp (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
Closure -> Exp (Lore m) -> m ()
letBindNames [Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
              case SubExp
se of I.Var VName
v | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Type -> Bool
forall shape u. TypeBase shape u -> Bool
primType (Type -> Bool) -> Type -> Bool
forall a b. (a -> b) -> a -> b
$ Param DeclType -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param DeclType
p ->
                                     ShapeChange SubExp -> VName -> BasicOp
Reshape ((SubExp -> DimChange SubExp) -> [SubExp] -> ShapeChange SubExp
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimCoercion ([SubExp] -> ShapeChange SubExp) -> [SubExp] -> ShapeChange SubExp
forall a b. (a -> b) -> a -> b
$ Type -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
arrayDims (Type -> [SubExp]) -> Type -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Param DeclType -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param DeclType
p) VName
v
                         SubExp
_ -> SubExp -> BasicOp
SubExp SubExp
se
            [SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody ([SubExp] -> Body) -> InternaliseM [SubExp] -> InternaliseM Body
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"loop_cond" ExpBase Info VName
cond
          [SubExp]
loop_end_cond <- Body (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *). MonadBinder m => Body (Lore m) -> m [SubExp]
bodyBind Body (Lore InternaliseM)
Body
loop_end_cond_body

          ([SubExp],
 (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     ([SubExp],
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall (m :: * -> *) a. Monad m => a -> m a
return ([SubExp]
shapeargs[SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++[SubExp]
loop_end_cond[SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++[SubExp]
ses,
                  (VName -> LoopForm SOACS
forall lore. VName -> LoopForm lore
I.WhileLoop (VName -> LoopForm SOACS) -> VName -> LoopForm SOACS
forall a b. (a -> b) -> a -> b
$ Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName Param DeclType
loop_while,
                   [Param DeclType]
[FParam]
shapepat,
                   Param DeclType
loop_while Param DeclType -> [Param DeclType] -> [Param DeclType]
forall a. a -> [a] -> [a]
: [Param DeclType]
mergepat',
                   SubExp
loop_initial_cond SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: [SubExp]
mergeinit))

internaliseExp String
desc (E.LetWith Ident
name Ident
src [DimIndexBase Info VName]
idxs ExpBase Info VName
ve ExpBase Info VName
body Info PatternType
t SrcLoc
loc) = do
  let pat :: Pattern
pat = VName -> Info PatternType -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase f vn
E.Id (Ident -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
E.identName Ident
name) (Ident -> Info PatternType
forall (f :: * -> *) vn. IdentBase f vn -> f PatternType
E.identType Ident
name) SrcLoc
loc
      src_t :: Info PatternType
src_t = PatternType -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
E.fromStruct (PatternType -> PatternType)
-> Info PatternType -> Info PatternType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ident -> Info PatternType
forall (f :: * -> *) vn. IdentBase f vn -> f PatternType
E.identType Ident
src
      e :: ExpBase Info VName
e = ExpBase Info VName
-> [DimIndexBase Info VName]
-> ExpBase Info VName
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> [DimIndexBase f vn] -> ExpBase f vn -> SrcLoc -> ExpBase f vn
E.Update (QualName VName -> Info PatternType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
E.Var (VName -> QualName VName
forall v. v -> QualName v
E.qualName (VName -> QualName VName) -> VName -> QualName VName
forall a b. (a -> b) -> a -> b
$ Ident -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
E.identName Ident
src) Info PatternType
src_t SrcLoc
loc) [DimIndexBase Info VName]
idxs ExpBase Info VName
ve SrcLoc
loc
  String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc (ExpBase Info VName -> InternaliseM [SubExp])
-> ExpBase Info VName -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Pattern
-> ExpBase Info VName
-> ExpBase Info VName
-> (Info PatternType, Info Closure)
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
PatternBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> (f PatternType, f Closure)
-> SrcLoc
-> ExpBase f vn
E.LetPat Pattern
pat ExpBase Info VName
e ExpBase Info VName
body (Info PatternType
t, Closure -> Info Closure
forall a. a -> Info a
Info []) SrcLoc
loc

internaliseExp String
desc (E.Update ExpBase Info VName
src [DimIndexBase Info VName]
slice ExpBase Info VName
ve SrcLoc
loc) = do
  [SubExp]
ves <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"lw_val" ExpBase Info VName
ve
  Closure
srcs <- String -> ExpBase Info VName -> InternaliseM Closure
internaliseExpToVars String
"src" ExpBase Info VName
src
  [SubExp]
dims <- case Closure
srcs of
            [] -> [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [] -- Will this happen?
            VName
v:Closure
_ -> Type -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
I.arrayDims (Type -> [SubExp]) -> InternaliseM Type -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
v
  ([DimIndex SubExp]
idxs', Certificates
cs) <- SrcLoc
-> [SubExp]
-> [DimIndexBase Info VName]
-> InternaliseM ([DimIndex SubExp], Certificates)
internaliseSlice SrcLoc
loc [SubExp]
dims [DimIndexBase Info VName]
slice

  let comb :: VName -> SubExp -> InternaliseM VName
comb VName
sname SubExp
ve' = do
        Type
sname_t <- VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
sname
        let full_slice :: [DimIndex SubExp]
full_slice = Type -> [DimIndex SubExp] -> [DimIndex SubExp]
fullSlice Type
sname_t [DimIndex SubExp]
idxs'
            rowtype :: Type
rowtype = Type
sname_t Type -> [SubExp] -> Type
forall oldshape u.
TypeBase oldshape u -> [SubExp] -> TypeBase Shape u
`setArrayDims` [DimIndex SubExp] -> [SubExp]
forall d. Slice d -> [d]
sliceDims [DimIndex SubExp]
full_slice
        SubExp
ve'' <- ErrorMsg SubExp
-> SrcLoc -> Type -> String -> SubExp -> InternaliseM SubExp
ensureShape ErrorMsg SubExp
"shape of value does not match shape of source array"
                SrcLoc
loc Type
rowtype String
"lw_val_correct_shape" SubExp
ve'
        String
-> VName
-> [DimIndex SubExp]
-> Exp (Lore InternaliseM)
-> InternaliseM VName
forall (m :: * -> *).
MonadBinder m =>
String -> VName -> [DimIndex SubExp] -> Exp (Lore m) -> m VName
letInPlace String
desc VName
sname [DimIndex SubExp]
full_slice (Exp (Lore InternaliseM) -> InternaliseM VName)
-> Exp (Lore InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
ve''
  Certificates -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
cs (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> Closure -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var (Closure -> [SubExp])
-> InternaliseM Closure -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> SubExp -> InternaliseM VName)
-> Closure -> [SubExp] -> InternaliseM Closure
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM VName -> SubExp -> InternaliseM VName
comb Closure
srcs [SubExp]
ves

internaliseExp String
desc (E.RecordUpdate ExpBase Info VName
src [Name]
fields ExpBase Info VName
ve Info PatternType
_ SrcLoc
_) = do
  [SubExp]
src' <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
src
  [SubExp]
ve' <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
ve
  StructType
-> [Name] -> [SubExp] -> [SubExp] -> InternaliseM [SubExp]
forall a. StructType -> [Name] -> [a] -> [a] -> InternaliseM [a]
replace (ExpBase Info VName -> PatternType
E.typeOf ExpBase Info VName
src PatternType -> () -> StructType
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` ()) [Name]
fields [SubExp]
ve' [SubExp]
src'
  where replace :: StructType -> [Name] -> [a] -> [a] -> InternaliseM [a]
replace (E.Scalar (E.Record Map Name StructType
m)) (Name
f:[Name]
fs) [a]
ve' [a]
src'
          | Just StructType
t <- Name -> Map Name StructType -> Maybe StructType
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f Map Name StructType
m = do
          Int
i <- ([Int] -> Int) -> InternaliseM [Int] -> InternaliseM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (InternaliseM [Int] -> InternaliseM Int)
-> InternaliseM [Int] -> InternaliseM Int
forall a b. (a -> b) -> a -> b
$ ((Name, StructType) -> InternaliseM Int)
-> [(Name, StructType)] -> InternaliseM [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (StructType -> InternaliseM Int
internalisedTypeSize (StructType -> InternaliseM Int)
-> ((Name, StructType) -> StructType)
-> (Name, StructType)
-> InternaliseM Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, StructType) -> StructType
forall a b. (a, b) -> b
snd) ([(Name, StructType)] -> InternaliseM [Int])
-> [(Name, StructType)] -> InternaliseM [Int]
forall a b. (a -> b) -> a -> b
$
               ((Name, StructType) -> Bool)
-> [(Name, StructType)] -> [(Name, StructType)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
/=Name
f) (Name -> Bool)
-> ((Name, StructType) -> Name) -> (Name, StructType) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, StructType) -> Name
forall a b. (a, b) -> a
fst) ([(Name, StructType)] -> [(Name, StructType)])
-> [(Name, StructType)] -> [(Name, StructType)]
forall a b. (a -> b) -> a -> b
$ Map Name StructType -> [(Name, StructType)]
forall a. Map Name a -> [(Name, a)]
sortFields Map Name StructType
m
          Int
k <- StructType -> InternaliseM Int
internalisedTypeSize StructType
t
          let ([a]
bef, [a]
to_update, [a]
aft) = Int -> Int -> [a] -> ([a], [a], [a])
forall a. Int -> Int -> [a] -> ([a], [a], [a])
splitAt3 Int
i Int
k [a]
src'
          [a]
src'' <- StructType -> [Name] -> [a] -> [a] -> InternaliseM [a]
replace StructType
t [Name]
fs [a]
ve' [a]
to_update
          [a] -> InternaliseM [a]
forall (m :: * -> *) a. Monad m => a -> m a
return ([a] -> InternaliseM [a]) -> [a] -> InternaliseM [a]
forall a b. (a -> b) -> a -> b
$ [a]
bef [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
src'' [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
aft
        replace StructType
_ [Name]
_ [a]
ve' [a]
_ = [a] -> InternaliseM [a]
forall (m :: * -> *) a. Monad m => a -> m a
return [a]
ve'

internaliseExp String
desc (E.Attr AttrInfo
attr ExpBase Info VName
e SrcLoc
_) =
  (InternaliseEnv -> InternaliseEnv)
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local InternaliseEnv -> InternaliseEnv
f (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
e
  where attrs :: Attrs
attrs = Attr -> Attrs
oneAttr (Attr -> Attrs) -> Attr -> Attrs
forall a b. (a -> b) -> a -> b
$ AttrInfo -> Attr
internaliseAttr AttrInfo
attr
        f :: InternaliseEnv -> InternaliseEnv
f InternaliseEnv
env | Attr
"unsafe" Attr -> Attrs -> Bool
`inAttrs` Attrs
attrs,
                Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ InternaliseEnv -> Bool
envSafe InternaliseEnv
env =
                  InternaliseEnv
env { envDoBoundsChecks :: Bool
envDoBoundsChecks = Bool
False }
              | Bool
otherwise =
                  InternaliseEnv
env { envAttrs :: Attrs
envAttrs = InternaliseEnv -> Attrs
envAttrs InternaliseEnv
env Attrs -> Attrs -> Attrs
forall a. Semigroup a => a -> a -> a
<> Attrs
attrs }

internaliseExp String
desc (E.Assert ExpBase Info VName
e1 ExpBase Info VName
e2 (Info String
check) SrcLoc
loc) = do
  SubExp
e1' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"assert_cond" ExpBase Info VName
e1
  Certificates
c <- String
-> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certificates
assert String
"assert_c" SubExp
e1' ([ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [String -> ErrorMsgPart SubExp
forall a. String -> ErrorMsgPart a
ErrorString (String -> ErrorMsgPart SubExp) -> String -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ String
"Assertion is false: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
check]) SrcLoc
loc
  -- Make sure there are some bindings to certify.
  Certificates -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
c (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ (SubExp -> InternaliseM SubExp)
-> [SubExp] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM SubExp
forall (m :: * -> *). MonadBinder m => SubExp -> m SubExp
rebind ([SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
e2
  where rebind :: SubExp -> m SubExp
rebind SubExp
v = do
          VName
v' <- String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"assert_res"
          Closure -> Exp (Lore m) -> m ()
forall (m :: * -> *).
MonadBinder m =>
Closure -> Exp (Lore m) -> m ()
letBindNames [VName
v'] (Exp (Lore m) -> m ()) -> Exp (Lore m) -> m ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
v
          SubExp -> m SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> m SubExp) -> SubExp -> m SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
v'

internaliseExp String
_ (E.Constr Name
c [ExpBase Info VName]
es (Info (E.Scalar (E.Sum Map Name [PatternType]
fs))) SrcLoc
_) = do
  ([DeclExtType]
ts, Map Name (Int, [Int])
constr_map) <- Map Name [StructType]
-> InternaliseM ([DeclExtType], Map Name (Int, [Int]))
internaliseSumType (Map Name [StructType]
 -> InternaliseM ([DeclExtType], Map Name (Int, [Int])))
-> Map Name [StructType]
-> InternaliseM ([DeclExtType], Map Name (Int, [Int]))
forall a b. (a -> b) -> a -> b
$ ([PatternType] -> [StructType])
-> Map Name [PatternType] -> Map Name [StructType]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((PatternType -> StructType) -> [PatternType] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct) Map Name [PatternType]
fs
  [SubExp]
es' <- [[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[SubExp]] -> [SubExp])
-> InternaliseM [[SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ExpBase Info VName -> InternaliseM [SubExp])
-> [ExpBase Info VName] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"payload") [ExpBase Info VName]
es

  let noExt :: p -> m SubExp
noExt p
_ = SubExp -> m SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> m SubExp) -> SubExp -> m SubExp
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int32 Integer
0
  [Type]
ts' <- (Int -> InternaliseM SubExp) -> [ExtType] -> InternaliseM [Type]
forall (m :: * -> *) u.
Monad m =>
(Int -> m SubExp) -> [TypeBase ExtShape u] -> m [TypeBase Shape u]
instantiateShapes Int -> InternaliseM SubExp
forall (m :: * -> *) p. Monad m => p -> m SubExp
noExt ([ExtType] -> InternaliseM [Type])
-> [ExtType] -> InternaliseM [Type]
forall a b. (a -> b) -> a -> b
$ (DeclExtType -> ExtType) -> [DeclExtType] -> [ExtType]
forall a b. (a -> b) -> [a] -> [b]
map DeclExtType -> ExtType
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
fromDecl [DeclExtType]
ts

  case Name -> Map Name (Int, [Int]) -> Maybe (Int, [Int])
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
c Map Name (Int, [Int])
constr_map of
    Just (Int
i, [Int]
js) ->
      (IntType -> Integer -> SubExp
intConst IntType
Int8 (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
i)SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
:) ([SubExp] -> [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [Type] -> [(Int, SubExp)] -> InternaliseM [SubExp]
forall (f :: * -> *) a.
(Num a, MonadBinder f, Eq a) =>
a -> [Type] -> [(a, SubExp)] -> f [SubExp]
clauses Int
0 [Type]
ts' ([Int] -> [SubExp] -> [(Int, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
js [SubExp]
es')
    Maybe (Int, [Int])
Nothing ->
      String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error String
"internaliseExp Constr: missing constructor"

  where clauses :: a -> [Type] -> [(a, SubExp)] -> f [SubExp]
clauses a
j (Type
t:[Type]
ts) [(a, SubExp)]
js_to_es
          | Just SubExp
e <- a
j a -> [(a, SubExp)] -> Maybe SubExp
forall a b. Eq a => a -> [(a, b)] -> Maybe b
`lookup` [(a, SubExp)]
js_to_es =
              (SubExp
eSubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
:) ([SubExp] -> [SubExp]) -> f [SubExp] -> f [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> [Type] -> [(a, SubExp)] -> f [SubExp]
clauses (a
ja -> a -> a
forall a. Num a => a -> a -> a
+a
1) [Type]
ts [(a, SubExp)]
js_to_es
          | Bool
otherwise = do
              SubExp
blank <- String -> Exp (Lore f) -> f SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"zero" (Exp (Lore f) -> f SubExp) -> f (Exp (Lore f)) -> f SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Type -> f (Exp (Lore f))
forall (m :: * -> *). MonadBinder m => Type -> m (Exp (Lore m))
eBlank Type
t
              (SubExp
blankSubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
:) ([SubExp] -> [SubExp]) -> f [SubExp] -> f [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> [Type] -> [(a, SubExp)] -> f [SubExp]
clauses (a
ja -> a -> a
forall a. Num a => a -> a -> a
+a
1) [Type]
ts [(a, SubExp)]
js_to_es
        clauses a
_ [] [(a, SubExp)]
_ =
          [SubExp] -> f [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return []

internaliseExp String
_ (E.Constr Name
_ [ExpBase Info VName]
_ (Info PatternType
t) SrcLoc
loc) =
  String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error (String -> InternaliseM [SubExp])
-> String -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp: constructor with type " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PatternType -> String
forall a. Pretty a => a -> String
pretty PatternType
t String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" at " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SrcLoc -> String
forall a. Located a => a -> String
locStr SrcLoc
loc

internaliseExp String
desc (E.Match ExpBase Info VName
e NonEmpty (CaseBase Info VName)
cs (Info PatternType
ret, Info Closure
retext) SrcLoc
_) = do
  [SubExp]
ses <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp (String
desc String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_scrutinee") ExpBase Info VName
e
  [SubExp]
res <-
    case NonEmpty (CaseBase Info VName)
-> (CaseBase Info VName, Maybe (NonEmpty (CaseBase Info VName)))
forall a. NonEmpty a -> (a, Maybe (NonEmpty a))
NE.uncons NonEmpty (CaseBase Info VName)
cs of
    (CasePat Pattern
pCase ExpBase Info VName
eCase SrcLoc
locCase, Maybe (NonEmpty (CaseBase Info VName))
Nothing) -> do
      (SubExp
_, [SubExp]
pertinent) <- Pattern -> [SubExp] -> InternaliseM (SubExp, [SubExp])
generateCond Pattern
pCase [SubExp]
ses
      Pattern
-> [SubExp]
-> ExpBase Info VName
-> SrcLoc
-> (ExpBase Info VName -> InternaliseM [SubExp])
-> InternaliseM [SubExp]
forall a.
Pattern
-> [SubExp]
-> ExpBase Info VName
-> SrcLoc
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat' Pattern
pCase [SubExp]
pertinent ExpBase Info VName
eCase SrcLoc
locCase (String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc)
    (CaseBase Info VName
c, Just NonEmpty (CaseBase Info VName)
cs') -> do
      let CasePat Pattern
pLast ExpBase Info VName
eLast SrcLoc
locLast = NonEmpty (CaseBase Info VName) -> CaseBase Info VName
forall a. NonEmpty a -> a
NE.last NonEmpty (CaseBase Info VName)
cs'
      Body
bFalse <- do
        (SubExp
_, [SubExp]
pertinent) <- Pattern -> [SubExp] -> InternaliseM (SubExp, [SubExp])
generateCond Pattern
pLast [SubExp]
ses
        Body
eLast' <- Pattern
-> [SubExp]
-> ExpBase Info VName
-> SrcLoc
-> (ExpBase Info VName -> InternaliseM Body)
-> InternaliseM Body
forall a.
Pattern
-> [SubExp]
-> ExpBase Info VName
-> SrcLoc
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat' Pattern
pLast [SubExp]
pertinent ExpBase Info VName
eLast SrcLoc
locLast ExpBase Info VName -> InternaliseM Body
internaliseBody
        (Body -> CaseBase Info VName -> InternaliseM Body)
-> Body -> [CaseBase Info VName] -> InternaliseM Body
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\Body
bf CaseBase Info VName
c' -> [InternaliseM (Exp (Lore InternaliseM))]
-> InternaliseM (Body (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
[m (Exp (Lore m))] -> m (Body (Lore m))
eBody ([InternaliseM (Exp (Lore InternaliseM))]
 -> InternaliseM (Body (Lore InternaliseM)))
-> [InternaliseM (Exp (Lore InternaliseM))]
-> InternaliseM (Body (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ InternaliseM (ExpT SOACS) -> [InternaliseM (ExpT SOACS)]
forall (m :: * -> *) a. Monad m => a -> m a
return (InternaliseM (ExpT SOACS) -> [InternaliseM (ExpT SOACS)])
-> InternaliseM (ExpT SOACS) -> [InternaliseM (ExpT SOACS)]
forall a b. (a -> b) -> a -> b
$ [SubExp]
-> CaseBase Info VName -> Body -> InternaliseM (ExpT SOACS)
generateCaseIf [SubExp]
ses CaseBase Info VName
c' Body
bf) Body
eLast' ([CaseBase Info VName] -> InternaliseM Body)
-> [CaseBase Info VName] -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$
          [CaseBase Info VName] -> [CaseBase Info VName]
forall a. [a] -> [a]
reverse ([CaseBase Info VName] -> [CaseBase Info VName])
-> [CaseBase Info VName] -> [CaseBase Info VName]
forall a b. (a -> b) -> a -> b
$ NonEmpty (CaseBase Info VName) -> [CaseBase Info VName]
forall a. NonEmpty a -> [a]
NE.init NonEmpty (CaseBase Info VName)
cs'
      String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (ExpT SOACS -> InternaliseM [SubExp])
-> InternaliseM (ExpT SOACS) -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [SubExp]
-> CaseBase Info VName -> Body -> InternaliseM (ExpT SOACS)
generateCaseIf [SubExp]
ses CaseBase Info VName
c Body
bFalse
  StructType -> Closure -> [SubExp] -> InternaliseM ()
bindExtSizes (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret) Closure
retext [SubExp]
res
  [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [SubExp]
res

-- The "interesting" cases are over, now it's mostly boilerplate.

internaliseExp String
_ (E.Literal PrimValue
v SrcLoc
_) =
  [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimValue
internalisePrimValue PrimValue
v]

internaliseExp String
_ (E.IntLit Integer
v (Info PatternType
t) SrcLoc
_) =
  case PatternType
t of
    E.Scalar (E.Prim (E.Signed IntType
it)) ->
      [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
I.IntValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> IntValue
forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
v]
    E.Scalar (E.Prim (E.Unsigned IntType
it)) ->
      [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
I.IntValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> IntValue
forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
v]
    E.Scalar (E.Prim (E.FloatType FloatType
ft)) ->
      [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
I.FloatValue (FloatValue -> PrimValue) -> FloatValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ FloatType -> Integer -> FloatValue
forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
ft Integer
v]
    PatternType
_ -> String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error (String -> InternaliseM [SubExp])
-> String -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp: nonsensical type for integer literal: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PatternType -> String
forall a. Pretty a => a -> String
pretty PatternType
t

internaliseExp String
_ (E.FloatLit Double
v (Info PatternType
t) SrcLoc
_) =
  case PatternType
t of
    E.Scalar (E.Prim (E.FloatType FloatType
ft)) ->
      [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
I.FloatValue (FloatValue -> PrimValue) -> FloatValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ FloatType -> Double -> FloatValue
forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
ft Double
v]
    PatternType
_ -> String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error (String -> InternaliseM [SubExp])
-> String -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp: nonsensical type for float literal: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PatternType -> String
forall a. Pretty a => a -> String
pretty PatternType
t

internaliseExp String
desc (E.If ExpBase Info VName
ce ExpBase Info VName
te ExpBase Info VName
fe (Info PatternType
ret, Info Closure
retext) SrcLoc
_) = do
  [SubExp]
ses <- String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (ExpT SOACS -> InternaliseM [SubExp])
-> InternaliseM (ExpT SOACS) -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<
         InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
(MonadBinder m, BranchType (Lore m) ~ ExtType) =>
m (Exp (Lore m))
-> m (Body (Lore m)) -> m (Body (Lore m)) -> m (Exp (Lore m))
eIf (BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS)
-> (SubExp -> BasicOp) -> SubExp -> ExpT SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
SubExp (SubExp -> ExpT SOACS)
-> InternaliseM SubExp -> InternaliseM (ExpT SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"cond" ExpBase Info VName
ce)
         (ExpBase Info VName -> InternaliseM Body
internaliseBody ExpBase Info VName
te) (ExpBase Info VName -> InternaliseM Body
internaliseBody ExpBase Info VName
fe)
  StructType -> Closure -> [SubExp] -> InternaliseM ()
bindExtSizes (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret) Closure
retext [SubExp]
ses
  [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [SubExp]
ses

-- Builtin operators are handled specially because they are
-- overloaded.
internaliseExp String
desc (E.BinOp (QualName VName
op, SrcLoc
_) Info PatternType
_ (ExpBase Info VName
xe,Info (StructType, Maybe VName)
_) (ExpBase Info VName
ye,Info (StructType, Maybe VName)
_) Info PatternType
_ Info Closure
_ SrcLoc
loc)
  | Just String -> InternaliseM [SubExp]
internalise <- QualName VName
-> [ExpBase Info VName]
-> SrcLoc
-> Maybe (String -> InternaliseM [SubExp])
isOverloadedFunction QualName VName
op [ExpBase Info VName
xe, ExpBase Info VName
ye] SrcLoc
loc =
      String -> InternaliseM [SubExp]
internalise String
desc

-- User-defined operators are just the same as a function call.
internaliseExp String
desc (E.BinOp (QualName VName
op, SrcLoc
oploc) (Info PatternType
t)
                     (ExpBase Info VName
xarg, Info (StructType
xt,Maybe VName
xext)) (ExpBase Info VName
yarg, Info (StructType
yt,Maybe VName
yext))
                     Info PatternType
_ (Info Closure
retext) SrcLoc
loc) =
  String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc (ExpBase Info VName -> InternaliseM [SubExp])
-> ExpBase Info VName -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
  ExpBase Info VName
-> ExpBase Info VName
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info Closure)
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f Closure)
-> SrcLoc
-> ExpBase f vn
E.Apply (ExpBase Info VName
-> ExpBase Info VName
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info Closure)
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f Closure)
-> SrcLoc
-> ExpBase f vn
E.Apply (QualName VName -> Info PatternType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
E.Var QualName VName
op (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t) SrcLoc
oploc) ExpBase Info VName
xarg ((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (StructType -> Diet
forall shape as. TypeBase shape as -> Diet
E.diet StructType
xt, Maybe VName
xext))
           (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ [PatternType] -> PatternType -> PatternType
forall as dim.
Monoid as =>
[TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType [StructType -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
E.fromStruct StructType
yt] PatternType
t, Closure -> Info Closure
forall a. a -> Info a
Info []) SrcLoc
loc)
          ExpBase Info VName
yarg ((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (StructType -> Diet
forall shape as. TypeBase shape as -> Diet
E.diet StructType
yt, Maybe VName
yext)) (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t, Closure -> Info Closure
forall a. a -> Info a
Info Closure
retext) SrcLoc
loc

internaliseExp String
desc (E.Project Name
k ExpBase Info VName
e (Info PatternType
rt) SrcLoc
_) = do
  Int
n <- StructType -> InternaliseM Int
internalisedTypeSize (StructType -> InternaliseM Int) -> StructType -> InternaliseM Int
forall a b. (a -> b) -> a -> b
$ PatternType
rt PatternType -> () -> StructType
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` ()
  Int
i' <- ([Int] -> Int) -> InternaliseM [Int] -> InternaliseM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (InternaliseM [Int] -> InternaliseM Int)
-> InternaliseM [Int] -> InternaliseM Int
forall a b. (a -> b) -> a -> b
$ (StructType -> InternaliseM Int)
-> [StructType] -> InternaliseM [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM StructType -> InternaliseM Int
internalisedTypeSize ([StructType] -> InternaliseM [Int])
-> [StructType] -> InternaliseM [Int]
forall a b. (a -> b) -> a -> b
$
        case ExpBase Info VName -> PatternType
E.typeOf ExpBase Info VName
e PatternType -> () -> StructType
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` () of
               E.Scalar (Record Map Name StructType
fs) ->
                 ((Name, StructType) -> StructType)
-> [(Name, StructType)] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map (Name, StructType) -> StructType
forall a b. (a, b) -> b
snd ([(Name, StructType)] -> [StructType])
-> [(Name, StructType)] -> [StructType]
forall a b. (a -> b) -> a -> b
$ ((Name, StructType) -> Bool)
-> [(Name, StructType)] -> [(Name, StructType)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
/=Name
k) (Name -> Bool)
-> ((Name, StructType) -> Name) -> (Name, StructType) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, StructType) -> Name
forall a b. (a, b) -> a
fst) ([(Name, StructType)] -> [(Name, StructType)])
-> [(Name, StructType)] -> [(Name, StructType)]
forall a b. (a -> b) -> a -> b
$ Map Name StructType -> [(Name, StructType)]
forall a. Map Name a -> [(Name, a)]
sortFields Map Name StructType
fs
               StructType
t -> [StructType
t]
  Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
take Int
n ([SubExp] -> [SubExp])
-> ([SubExp] -> [SubExp]) -> [SubExp] -> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
drop Int
i' ([SubExp] -> [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
e

internaliseExp String
_ e :: ExpBase Info VName
e@E.Lambda{} =
  String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error (String -> InternaliseM [SubExp])
-> String -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp: Unexpected lambda at " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SrcLoc -> String
forall a. Located a => a -> String
locStr (ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
e)

internaliseExp String
_ e :: ExpBase Info VName
e@E.OpSection{} =
  String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error (String -> InternaliseM [SubExp])
-> String -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp: Unexpected operator section at " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SrcLoc -> String
forall a. Located a => a -> String
locStr (ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
e)

internaliseExp String
_ e :: ExpBase Info VName
e@E.OpSectionLeft{} =
  String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error (String -> InternaliseM [SubExp])
-> String -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp: Unexpected left operator section at " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SrcLoc -> String
forall a. Located a => a -> String
locStr (ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
e)

internaliseExp String
_ e :: ExpBase Info VName
e@E.OpSectionRight{} =
  String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error (String -> InternaliseM [SubExp])
-> String -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp: Unexpected right operator section at " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SrcLoc -> String
forall a. Located a => a -> String
locStr (ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
e)

internaliseExp String
_ e :: ExpBase Info VName
e@E.ProjectSection{} =
  String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error (String -> InternaliseM [SubExp])
-> String -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp: Unexpected projection section at " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SrcLoc -> String
forall a. Located a => a -> String
locStr (ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
e)

internaliseExp String
_ e :: ExpBase Info VName
e@E.IndexSection{} =
  String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error (String -> InternaliseM [SubExp])
-> String -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp: Unexpected index section at " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SrcLoc -> String
forall a. Located a => a -> String
locStr (ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
e)

internaliseArg :: String -> (E.Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg :: String
-> (ExpBase Info VName, Maybe VName) -> InternaliseM [SubExp]
internaliseArg String
desc (ExpBase Info VName
arg, Maybe VName
argdim) = do
  [SubExp]
arg' <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
arg
  case ([SubExp]
arg', Maybe VName
argdim) of
    ([SubExp
se], Just VName
d) -> Closure -> Exp (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
Closure -> Exp (Lore m) -> m ()
letBindNames [VName
d] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
    ([SubExp], Maybe VName)
_ -> () -> InternaliseM ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
  [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [SubExp]
arg'

generateCond :: E.Pattern -> [I.SubExp] -> InternaliseM (I.SubExp, [I.SubExp])
generateCond :: Pattern -> [SubExp] -> InternaliseM (SubExp, [SubExp])
generateCond Pattern
orig_p [SubExp]
orig_ses = do
  ([SubExp]
cmps, [SubExp]
pertinent, [SubExp]
_) <- Pattern -> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
compares Pattern
orig_p [SubExp]
orig_ses
  SubExp
cmp <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"matches" (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [SubExp] -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => [SubExp] -> m (Exp (Lore m))
eAll [SubExp]
cmps
  (SubExp, [SubExp]) -> InternaliseM (SubExp, [SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
cmp, [SubExp]
pertinent)
  where
    -- Literals are always primitive values.
    compares :: Pattern -> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
compares (E.PatternLit ExpBase Info VName
e Info PatternType
_ SrcLoc
_) (SubExp
se:[SubExp]
ses) = do
      SubExp
e' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"constant" ExpBase Info VName
e
      PrimType
t' <- Type -> PrimType
forall shape u. TypeBase shape u -> PrimType
I.elemType (Type -> PrimType) -> InternaliseM Type -> InternaliseM PrimType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
se
      SubExp
cmp <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"match_lit" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
t') SubExp
e' SubExp
se
      ([SubExp], [SubExp], [SubExp])
-> InternaliseM ([SubExp], [SubExp], [SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return ([SubExp
cmp], [SubExp
se], [SubExp]
ses)

    compares (E.PatternConstr Name
c (Info (E.Scalar (E.Sum Map Name [PatternType]
fs))) [Pattern]
pats SrcLoc
_) (SubExp
se:[SubExp]
ses) = do
      ([DeclExtType]
payload_ts, Map Name (Int, [Int])
m) <- Map Name [StructType]
-> InternaliseM ([DeclExtType], Map Name (Int, [Int]))
internaliseSumType (Map Name [StructType]
 -> InternaliseM ([DeclExtType], Map Name (Int, [Int])))
-> Map Name [StructType]
-> InternaliseM ([DeclExtType], Map Name (Int, [Int]))
forall a b. (a -> b) -> a -> b
$ ([PatternType] -> [StructType])
-> Map Name [PatternType] -> Map Name [StructType]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((PatternType -> StructType) -> [PatternType] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct) Map Name [PatternType]
fs
      case Name -> Map Name (Int, [Int]) -> Maybe (Int, [Int])
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
c Map Name (Int, [Int])
m of
        Just (Int
i, [Int]
payload_is) -> do
          let i' :: SubExp
i' = IntType -> Integer -> SubExp
intConst IntType
Int8 (Integer -> SubExp) -> Integer -> SubExp
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
i
          let ([SubExp]
payload_ses, [SubExp]
ses') = Int -> [SubExp] -> ([SubExp], [SubExp])
forall a. Int -> [a] -> ([a], [a])
splitAt ([DeclExtType] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DeclExtType]
payload_ts) [SubExp]
ses
          SubExp
cmp <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"match_constr" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
int8) SubExp
i' SubExp
se
          ([SubExp]
cmps, [SubExp]
pertinent, [SubExp]
_) <- [Pattern]
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
comparesMany [Pattern]
pats ([SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp]))
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
forall a b. (a -> b) -> a -> b
$ (Int -> SubExp) -> [Int] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map ([SubExp]
payload_ses[SubExp] -> Int -> SubExp
forall a. [a] -> Int -> a
!!) [Int]
payload_is
          ([SubExp], [SubExp], [SubExp])
-> InternaliseM ([SubExp], [SubExp], [SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
cmp SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: [SubExp]
cmps, [SubExp]
pertinent, [SubExp]
ses')
        Maybe (Int, [Int])
Nothing ->
          String -> InternaliseM ([SubExp], [SubExp], [SubExp])
forall a. HasCallStack => String -> a
error String
"generateCond: missing constructor"

    compares (E.PatternConstr Name
_ (Info PatternType
t) [Pattern]
_ SrcLoc
_) [SubExp]
_ =
      String -> InternaliseM ([SubExp], [SubExp], [SubExp])
forall a. HasCallStack => String -> a
error (String -> InternaliseM ([SubExp], [SubExp], [SubExp]))
-> String -> InternaliseM ([SubExp], [SubExp], [SubExp])
forall a b. (a -> b) -> a -> b
$ String
"generateCond: PatternConstr has nonsensical type: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PatternType -> String
forall a. Pretty a => a -> String
pretty PatternType
t

    compares (E.Id VName
_ Info PatternType
t SrcLoc
loc) [SubExp]
ses =
      Pattern -> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
compares (Info PatternType -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
f PatternType -> SrcLoc -> PatternBase f vn
E.Wildcard Info PatternType
t SrcLoc
loc) [SubExp]
ses

    compares (E.Wildcard (Info PatternType
t) SrcLoc
_) [SubExp]
ses = do
      Int
n <- StructType -> InternaliseM Int
internalisedTypeSize (StructType -> InternaliseM Int) -> StructType -> InternaliseM Int
forall a b. (a -> b) -> a -> b
$ PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
t
      let ([SubExp]
id_ses, [SubExp]
rest_ses) = Int -> [SubExp] -> ([SubExp], [SubExp])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n [SubExp]
ses
      ([SubExp], [SubExp], [SubExp])
-> InternaliseM ([SubExp], [SubExp], [SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return ([], [SubExp]
id_ses, [SubExp]
rest_ses)

    compares (E.PatternParens Pattern
pat SrcLoc
_) [SubExp]
ses =
      Pattern -> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
compares Pattern
pat [SubExp]
ses

    compares (E.TuplePattern [Pattern]
pats SrcLoc
_) [SubExp]
ses =
      [Pattern]
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
comparesMany [Pattern]
pats [SubExp]
ses

    compares (E.RecordPattern [(Name, Pattern)]
fs SrcLoc
_) [SubExp]
ses =
      [Pattern]
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
comparesMany (((Name, Pattern) -> Pattern) -> [(Name, Pattern)] -> [Pattern]
forall a b. (a -> b) -> [a] -> [b]
map (Name, Pattern) -> Pattern
forall a b. (a, b) -> b
snd ([(Name, Pattern)] -> [Pattern]) -> [(Name, Pattern)] -> [Pattern]
forall a b. (a -> b) -> a -> b
$ Map Name Pattern -> [(Name, Pattern)]
forall a. Map Name a -> [(Name, a)]
E.sortFields (Map Name Pattern -> [(Name, Pattern)])
-> Map Name Pattern -> [(Name, Pattern)]
forall a b. (a -> b) -> a -> b
$ [(Name, Pattern)] -> Map Name Pattern
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, Pattern)]
fs) [SubExp]
ses

    compares (E.PatternAscription Pattern
pat TypeDeclBase Info VName
_ SrcLoc
_) [SubExp]
ses =
      Pattern -> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
compares Pattern
pat [SubExp]
ses

    compares Pattern
pat [] =
      String -> InternaliseM ([SubExp], [SubExp], [SubExp])
forall a. HasCallStack => String -> a
error (String -> InternaliseM ([SubExp], [SubExp], [SubExp]))
-> String -> InternaliseM ([SubExp], [SubExp], [SubExp])
forall a b. (a -> b) -> a -> b
$ String
"generateCond: No values left for pattern " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Pattern -> String
forall a. Pretty a => a -> String
pretty Pattern
pat

    comparesMany :: [Pattern]
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
comparesMany [] [SubExp]
ses = ([SubExp], [SubExp], [SubExp])
-> InternaliseM ([SubExp], [SubExp], [SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return ([], [], [SubExp]
ses)
    comparesMany (Pattern
pat:[Pattern]
pats) [SubExp]
ses = do
      ([SubExp]
cmps1, [SubExp]
pertinent1, [SubExp]
ses') <- Pattern -> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
compares Pattern
pat [SubExp]
ses
      ([SubExp]
cmps2, [SubExp]
pertinent2, [SubExp]
ses'') <- [Pattern]
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
comparesMany [Pattern]
pats [SubExp]
ses'
      ([SubExp], [SubExp], [SubExp])
-> InternaliseM ([SubExp], [SubExp], [SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return ([SubExp]
cmps1 [SubExp] -> [SubExp] -> [SubExp]
forall a. Semigroup a => a -> a -> a
<> [SubExp]
cmps2,
              [SubExp]
pertinent1 [SubExp] -> [SubExp] -> [SubExp]
forall a. Semigroup a => a -> a -> a
<> [SubExp]
pertinent2,
              [SubExp]
ses'')

generateCaseIf :: [I.SubExp] -> Case -> I.Body -> InternaliseM I.Exp
generateCaseIf :: [SubExp]
-> CaseBase Info VName -> Body -> InternaliseM (ExpT SOACS)
generateCaseIf [SubExp]
ses (CasePat Pattern
p ExpBase Info VName
eCase SrcLoc
loc) Body
bFail = do
  (SubExp
cond, [SubExp]
pertinent) <- Pattern -> [SubExp] -> InternaliseM (SubExp, [SubExp])
generateCond Pattern
p [SubExp]
ses
  Body
eCase' <- Pattern
-> [SubExp]
-> ExpBase Info VName
-> SrcLoc
-> (ExpBase Info VName -> InternaliseM Body)
-> InternaliseM Body
forall a.
Pattern
-> [SubExp]
-> ExpBase Info VName
-> SrcLoc
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat' Pattern
p [SubExp]
pertinent ExpBase Info VName
eCase SrcLoc
loc ExpBase Info VName -> InternaliseM Body
internaliseBody
  InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
(MonadBinder m, BranchType (Lore m) ~ ExtType) =>
m (Exp (Lore m))
-> m (Body (Lore m)) -> m (Body (Lore m)) -> m (Exp (Lore m))
eIf (SubExp -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => SubExp -> m (Exp (Lore m))
eSubExp SubExp
cond) (Body -> InternaliseM Body
forall (m :: * -> *) a. Monad m => a -> m a
return Body
eCase') (Body -> InternaliseM Body
forall (m :: * -> *) a. Monad m => a -> m a
return Body
bFail)

internalisePat :: String -> E.Pattern -> E.Exp
               -> E.Exp -> SrcLoc -> (E.Exp -> InternaliseM a)
               -> InternaliseM a
internalisePat :: String
-> Pattern
-> ExpBase Info VName
-> ExpBase Info VName
-> SrcLoc
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat String
desc Pattern
p ExpBase Info VName
e ExpBase Info VName
body SrcLoc
loc ExpBase Info VName -> InternaliseM a
m = do
  [SubExp]
ses <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc' ExpBase Info VName
e
  Pattern
-> [SubExp]
-> ExpBase Info VName
-> SrcLoc
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
forall a.
Pattern
-> [SubExp]
-> ExpBase Info VName
-> SrcLoc
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat' Pattern
p [SubExp]
ses ExpBase Info VName
body SrcLoc
loc ExpBase Info VName -> InternaliseM a
m
  where desc' :: String
desc' = case Set Ident -> [Ident]
forall a. Set a -> [a]
S.toList (Set Ident -> [Ident]) -> Set Ident -> [Ident]
forall a b. (a -> b) -> a -> b
$ Pattern -> Set Ident
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set (IdentBase f vn)
E.patternIdents Pattern
p of
                  [Ident
v] -> VName -> String
baseString (VName -> String) -> VName -> String
forall a b. (a -> b) -> a -> b
$ Ident -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
E.identName Ident
v
                  [Ident]
_ -> String
desc

internalisePat' :: E.Pattern -> [I.SubExp]
                -> E.Exp -> SrcLoc -> (E.Exp -> InternaliseM a)
                -> InternaliseM a
internalisePat' :: Pattern
-> [SubExp]
-> ExpBase Info VName
-> SrcLoc
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat' Pattern
p [SubExp]
ses ExpBase Info VName
body SrcLoc
loc ExpBase Info VName -> InternaliseM a
m = do
  [ExtType]
t <- [Type] -> [ExtType]
forall u. [TypeBase Shape u] -> [TypeBase ExtShape u]
I.staticShapes ([Type] -> [ExtType])
-> InternaliseM [Type] -> InternaliseM [ExtType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> InternaliseM Type) -> [SubExp] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
I.subExpType [SubExp]
ses
  Pattern
-> [ExtType]
-> (Closure -> MatchPattern -> InternaliseM a)
-> InternaliseM a
forall a.
Pattern
-> [ExtType]
-> (Closure -> MatchPattern -> InternaliseM a)
-> InternaliseM a
stmPattern Pattern
p [ExtType]
t ((Closure -> MatchPattern -> InternaliseM a) -> InternaliseM a)
-> (Closure -> MatchPattern -> InternaliseM a) -> InternaliseM a
forall a b. (a -> b) -> a -> b
$ \Closure
pat_names MatchPattern
match -> do
    [SubExp]
ses' <- MatchPattern
match SrcLoc
loc [SubExp]
ses
    [(VName, SubExp)]
-> ((VName, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Closure -> [SubExp] -> [(VName, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip Closure
pat_names [SubExp]
ses') (((VName, SubExp) -> InternaliseM ()) -> InternaliseM ())
-> ((VName, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \(VName
v,SubExp
se) ->
      Closure -> Exp (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
Closure -> Exp (Lore m) -> m ()
letBindNames [VName
v] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
se
    ExpBase Info VName -> InternaliseM a
m ExpBase Info VName
body

internaliseSlice :: SrcLoc
                 -> [SubExp]
                 -> [E.DimIndex]
                 -> InternaliseM ([I.DimIndex SubExp], Certificates)
internaliseSlice :: SrcLoc
-> [SubExp]
-> [DimIndexBase Info VName]
-> InternaliseM ([DimIndex SubExp], Certificates)
internaliseSlice SrcLoc
loc [SubExp]
dims [DimIndexBase Info VName]
idxs = do
 ([DimIndex SubExp]
idxs', [SubExp]
oks, [[ErrorMsgPart SubExp]]
parts) <- [(DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])]
-> ([DimIndex SubExp], [SubExp], [[ErrorMsgPart SubExp]])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 ([(DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])]
 -> ([DimIndex SubExp], [SubExp], [[ErrorMsgPart SubExp]]))
-> InternaliseM [(DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])]
-> InternaliseM
     ([DimIndex SubExp], [SubExp], [[ErrorMsgPart SubExp]])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp
 -> DimIndexBase Info VName
 -> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp]))
-> [SubExp]
-> [DimIndexBase Info VName]
-> InternaliseM [(DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM SubExp
-> DimIndexBase Info VName
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
internaliseDimIndex [SubExp]
dims [DimIndexBase Info VName]
idxs
 SubExp
ok <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"index_ok" (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [SubExp] -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => [SubExp] -> m (Exp (Lore m))
eAll [SubExp]
oks
 let msg :: ErrorMsg SubExp
msg = [ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg ([ErrorMsgPart SubExp] -> ErrorMsg SubExp)
-> [ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp
"Index ["] [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
-> [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
", "] [[ErrorMsgPart SubExp]]
parts [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++
           [ErrorMsgPart SubExp
"] out of bounds for array of shape ["] [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++
           ErrorMsgPart SubExp
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. a -> [a] -> [a]
intersperse ErrorMsgPart SubExp
"][" ((SubExp -> ErrorMsgPart SubExp)
-> [SubExp] -> [ErrorMsgPart SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt32 ([SubExp] -> [ErrorMsgPart SubExp])
-> [SubExp] -> [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
take ([DimIndexBase Info VName] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimIndexBase Info VName]
idxs) [SubExp]
dims) [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"]."]
 Certificates
c <- String
-> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certificates
assert String
"index_certs" SubExp
ok ErrorMsg SubExp
msg SrcLoc
loc
 ([DimIndex SubExp], Certificates)
-> InternaliseM ([DimIndex SubExp], Certificates)
forall (m :: * -> *) a. Monad m => a -> m a
return ([DimIndex SubExp]
idxs', Certificates
c)

internaliseDimIndex :: SubExp -> E.DimIndex
                    -> InternaliseM (I.DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
internaliseDimIndex :: SubExp
-> DimIndexBase Info VName
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
internaliseDimIndex SubExp
w (E.DimFix ExpBase Info VName
i) = do
  (SubExp
i', IntType
_) <- String -> ExpBase Info VName -> InternaliseM (SubExp, IntType)
internaliseDimExp String
"i" ExpBase Info VName
i
  let lowerBound :: ExpT SOACS
lowerBound = BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
                   CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
I.Int32) (Int32 -> SubExp
forall v. IsValue v => v -> SubExp
I.constant (Int32
0 :: I.Int32)) SubExp
i'
      upperBound :: ExpT SOACS
upperBound = BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
                   CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSlt IntType
I.Int32) SubExp
i' SubExp
w
  SubExp
ok <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"bounds_check" (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinOp
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> m (Exp (Lore m)) -> m (Exp (Lore m)) -> m (Exp (Lore m))
eBinOp BinOp
I.LogAnd (ExpT SOACS -> InternaliseM (ExpT SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ExpT SOACS
lowerBound) (ExpT SOACS -> InternaliseM (ExpT SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ExpT SOACS
upperBound)
  (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
I.DimFix SubExp
i', SubExp
ok, [SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt32 SubExp
i'])

-- Special-case an important common case that otherwise leads to horrible code.
internaliseDimIndex SubExp
w (E.DimSlice Maybe (ExpBase Info VName)
Nothing Maybe (ExpBase Info VName)
Nothing
                       (Just (E.Negate (E.IntLit Integer
1 Info PatternType
_ SrcLoc
_) SrcLoc
_))) = do
  SubExp
w_minus_1 <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"w_minus_1" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
               BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int32 Overflow
I.OverflowWrap) SubExp
w SubExp
one
  (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
I.DimSlice SubExp
w_minus_1 SubExp
w (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int32 (-Integer
1),
          Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
True,
          [ErrorMsgPart SubExp]
forall a. Monoid a => a
mempty)
  where one :: SubExp
one = Int32 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int32
1::Int32)

internaliseDimIndex SubExp
w (E.DimSlice Maybe (ExpBase Info VName)
i Maybe (ExpBase Info VName)
j Maybe (ExpBase Info VName)
s) = do
  SubExp
s' <- InternaliseM SubExp
-> (ExpBase Info VName -> InternaliseM SubExp)
-> Maybe (ExpBase Info VName)
-> InternaliseM SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (SubExp -> InternaliseM SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return SubExp
one) (((SubExp, IntType) -> SubExp)
-> InternaliseM (SubExp, IntType) -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SubExp, IntType) -> SubExp
forall a b. (a, b) -> a
fst (InternaliseM (SubExp, IntType) -> InternaliseM SubExp)
-> (ExpBase Info VName -> InternaliseM (SubExp, IntType))
-> ExpBase Info VName
-> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ExpBase Info VName -> InternaliseM (SubExp, IntType)
internaliseDimExp String
"s") Maybe (ExpBase Info VName)
s
  SubExp
s_sign <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"s_sign" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (IntType -> UnOp
I.SSignum IntType
Int32) SubExp
s'
  SubExp
backwards <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"backwards" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
int32) SubExp
s_sign SubExp
negone
  SubExp
w_minus_1 <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"w_minus_1" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int32 Overflow
I.OverflowWrap) SubExp
w SubExp
one
  let i_def :: InternaliseM SubExp
i_def = String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"i_def" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ SubExp -> Body -> Body -> IfDec (BranchType SOACS) -> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If SubExp
backwards
              ([SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [SubExp
w_minus_1])
              ([SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [SubExp
zero]) (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ [Type] -> IfDec ExtType
ifCommon [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int32]
      j_def :: InternaliseM SubExp
j_def = String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"j_def" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ SubExp -> Body -> Body -> IfDec (BranchType SOACS) -> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If SubExp
backwards
              ([SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [SubExp
negone])
              ([SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [SubExp
w]) (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ [Type] -> IfDec ExtType
ifCommon [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int32]
  SubExp
i' <- InternaliseM SubExp
-> (ExpBase Info VName -> InternaliseM SubExp)
-> Maybe (ExpBase Info VName)
-> InternaliseM SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe InternaliseM SubExp
i_def (((SubExp, IntType) -> SubExp)
-> InternaliseM (SubExp, IntType) -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SubExp, IntType) -> SubExp
forall a b. (a, b) -> a
fst (InternaliseM (SubExp, IntType) -> InternaliseM SubExp)
-> (ExpBase Info VName -> InternaliseM (SubExp, IntType))
-> ExpBase Info VName
-> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ExpBase Info VName -> InternaliseM (SubExp, IntType)
internaliseDimExp String
"i") Maybe (ExpBase Info VName)
i
  SubExp
j' <- InternaliseM SubExp
-> (ExpBase Info VName -> InternaliseM SubExp)
-> Maybe (ExpBase Info VName)
-> InternaliseM SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe InternaliseM SubExp
j_def (((SubExp, IntType) -> SubExp)
-> InternaliseM (SubExp, IntType) -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SubExp, IntType) -> SubExp
forall a b. (a, b) -> a
fst (InternaliseM (SubExp, IntType) -> InternaliseM SubExp)
-> (ExpBase Info VName -> InternaliseM (SubExp, IntType))
-> ExpBase Info VName
-> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ExpBase Info VName -> InternaliseM (SubExp, IntType)
internaliseDimExp String
"j") Maybe (ExpBase Info VName)
j
  SubExp
j_m_i <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"j_m_i" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int32 Overflow
I.OverflowWrap) SubExp
j' SubExp
i'
  -- Something like a division-rounding-up, but accomodating negative
  -- operands.
  let divRounding :: InternaliseM (ExpT SOACS)
-> InternaliseM (ExpT SOACS)
-> InternaliseM (Exp (Lore InternaliseM))
divRounding InternaliseM (ExpT SOACS)
x InternaliseM (ExpT SOACS)
y =
        BinOp
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> m (Exp (Lore m)) -> m (Exp (Lore m)) -> m (Exp (Lore m))
eBinOp (IntType -> Safety -> BinOp
SQuot IntType
Int32 Safety
Unsafe)
        (BinOp
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> m (Exp (Lore m)) -> m (Exp (Lore m)) -> m (Exp (Lore m))
eBinOp (IntType -> Overflow -> BinOp
Add IntType
Int32 Overflow
I.OverflowWrap) InternaliseM (Exp (Lore InternaliseM))
InternaliseM (ExpT SOACS)
x
         (BinOp
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> m (Exp (Lore m)) -> m (Exp (Lore m)) -> m (Exp (Lore m))
eBinOp (IntType -> Overflow -> BinOp
Sub IntType
Int32 Overflow
I.OverflowWrap) InternaliseM (Exp (Lore InternaliseM))
InternaliseM (ExpT SOACS)
y (InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
m (Exp (Lore m)) -> m (Exp (Lore m))
eSignum (InternaliseM (Exp (Lore InternaliseM))
 -> InternaliseM (Exp (Lore InternaliseM)))
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ SubExp -> InternaliseM (Exp (Lore InternaliseM))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Lore m))
toExp SubExp
s'))) InternaliseM (Exp (Lore InternaliseM))
InternaliseM (ExpT SOACS)
y
  SubExp
n <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"n" (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (ExpT SOACS)
-> InternaliseM (ExpT SOACS)
-> InternaliseM (Exp (Lore InternaliseM))
divRounding (SubExp -> InternaliseM (Exp (Lore InternaliseM))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Lore m))
toExp SubExp
j_m_i) (SubExp -> InternaliseM (Exp (Lore InternaliseM))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Lore m))
toExp SubExp
s')

  -- Bounds checks depend on whether we are slicing forwards or
  -- backwards.  If forwards, we must check '0 <= i && i <= j'.  If
  -- backwards, '-1 <= j && j <= i'.  In both cases, we check '0 <=
  -- i+n*s && i+(n-1)*s < w'.  We only check if the slice is nonempty.
  SubExp
empty_slice <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"empty_slice" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
CmpEq PrimType
int32) SubExp
n SubExp
zero

  SubExp
m <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"m" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int32 Overflow
I.OverflowWrap) SubExp
n SubExp
one
  SubExp
m_t_s <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"m_t_s" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Mul IntType
Int32 Overflow
I.OverflowWrap) SubExp
m SubExp
s'
  SubExp
i_p_m_t_s <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"i_p_m_t_s" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Add IntType
Int32 Overflow
I.OverflowWrap) SubExp
i' SubExp
m_t_s
  SubExp
zero_leq_i_p_m_t_s <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"zero_leq_i_p_m_t_s" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
                        BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int32) SubExp
zero SubExp
i_p_m_t_s
  SubExp
i_p_m_t_s_leq_w <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"i_p_m_t_s_leq_w" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
                     BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int32) SubExp
i_p_m_t_s SubExp
w
  SubExp
i_p_m_t_s_lth_w <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"i_p_m_t_s_leq_w" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
                     BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSlt IntType
Int32) SubExp
i_p_m_t_s SubExp
w

  SubExp
zero_lte_i <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"zero_lte_i" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int32) SubExp
zero SubExp
i'
  SubExp
i_lte_j <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"i_lte_j" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int32) SubExp
i' SubExp
j'
  SubExp
forwards_ok <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"forwards_ok" (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<
                 [SubExp] -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => [SubExp] -> m (Exp (Lore m))
eAll [SubExp
zero_lte_i, SubExp
zero_lte_i, SubExp
i_lte_j, SubExp
zero_leq_i_p_m_t_s, SubExp
i_p_m_t_s_lth_w]

  SubExp
negone_lte_j <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"negone_lte_j" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int32) SubExp
negone SubExp
j'
  SubExp
j_lte_i <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"j_lte_i" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int32) SubExp
j' SubExp
i'
  SubExp
backwards_ok <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"backwards_ok" (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<
                  [SubExp] -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => [SubExp] -> m (Exp (Lore m))
eAll
                  [SubExp
negone_lte_j, SubExp
negone_lte_j, SubExp
j_lte_i, SubExp
zero_leq_i_p_m_t_s, SubExp
i_p_m_t_s_leq_w]

  SubExp
slice_ok <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"slice_ok" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ SubExp -> Body -> Body -> IfDec (BranchType SOACS) -> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If SubExp
backwards
              ([SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [SubExp
backwards_ok])
              ([SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [SubExp
forwards_ok]) (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
              [Type] -> IfDec ExtType
ifCommon [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.Bool]
  SubExp
ok_or_empty <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"ok_or_empty" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
                 BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogOr SubExp
empty_slice SubExp
slice_ok

  let parts :: [ErrorMsgPart SubExp]
parts = case (Maybe (ExpBase Info VName)
i, Maybe (ExpBase Info VName)
j, Maybe (ExpBase Info VName)
s) of
                (Maybe (ExpBase Info VName)
_, Maybe (ExpBase Info VName)
_, Just{}) ->
                  [ErrorMsgPart SubExp
-> (ExpBase Info VName -> ErrorMsgPart SubExp)
-> Maybe (ExpBase Info VName)
-> ErrorMsgPart SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp
forall a b. a -> b -> a
const (ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp)
-> ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt32 SubExp
i') Maybe (ExpBase Info VName)
i, ErrorMsgPart SubExp
":",
                   ErrorMsgPart SubExp
-> (ExpBase Info VName -> ErrorMsgPart SubExp)
-> Maybe (ExpBase Info VName)
-> ErrorMsgPart SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp
forall a b. a -> b -> a
const (ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp)
-> ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt32 SubExp
j') Maybe (ExpBase Info VName)
j, ErrorMsgPart SubExp
":",
                   SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt32 SubExp
s']
                (Maybe (ExpBase Info VName)
_, Just{}, Maybe (ExpBase Info VName)
_) ->
                  [ErrorMsgPart SubExp
-> (ExpBase Info VName -> ErrorMsgPart SubExp)
-> Maybe (ExpBase Info VName)
-> ErrorMsgPart SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp
forall a b. a -> b -> a
const (ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp)
-> ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt32 SubExp
i') Maybe (ExpBase Info VName)
i, ErrorMsgPart SubExp
":",
                   SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt32 SubExp
j'] [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++
                   [ErrorMsgPart SubExp]
-> (ExpBase Info VName -> [ErrorMsgPart SubExp])
-> Maybe (ExpBase Info VName)
-> [ErrorMsgPart SubExp]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [ErrorMsgPart SubExp]
forall a. Monoid a => a
mempty ([ErrorMsgPart SubExp]
-> ExpBase Info VName -> [ErrorMsgPart SubExp]
forall a b. a -> b -> a
const [ErrorMsgPart SubExp
":", SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt32 SubExp
s']) Maybe (ExpBase Info VName)
s
                (Maybe (ExpBase Info VName)
_, Maybe (ExpBase Info VName)
Nothing, Maybe (ExpBase Info VName)
Nothing) ->
                  [SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt32 SubExp
i', ErrorMsgPart SubExp
":"]
  (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
I.DimSlice SubExp
i' SubExp
n SubExp
s', SubExp
ok_or_empty, [ErrorMsgPart SubExp]
parts)
  where zero :: SubExp
zero = Int32 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int32
0::Int32)
        negone :: SubExp
negone = Int32 -> SubExp
forall v. IsValue v => v -> SubExp
constant (-Int32
1::Int32)
        one :: SubExp
one = Int32 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int32
1::Int32)

internaliseScanOrReduce :: String -> String
                        -> (SubExp -> I.Lambda -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS))
                        -> (E.Exp, E.Exp, E.Exp, SrcLoc)
                        -> InternaliseM [SubExp]
internaliseScanOrReduce :: String
-> String
-> (SubExp
    -> Lambda -> [SubExp] -> Closure -> InternaliseM (SOAC SOACS))
-> (ExpBase Info VName, ExpBase Info VName, ExpBase Info VName,
    SrcLoc)
-> InternaliseM [SubExp]
internaliseScanOrReduce String
desc String
what SubExp
-> Lambda -> [SubExp] -> Closure -> InternaliseM (SOAC SOACS)
f (ExpBase Info VName
lam, ExpBase Info VName
ne, ExpBase Info VName
arr, SrcLoc
loc) = do
  Closure
arrs <- String -> ExpBase Info VName -> InternaliseM Closure
internaliseExpToVars (String
whatString -> String -> String
forall a. [a] -> [a] -> [a]
++String
"_arr") ExpBase Info VName
arr
  [SubExp]
nes <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp (String
whatString -> String -> String
forall a. [a] -> [a] -> [a]
++String
"_ne") ExpBase Info VName
ne
  [SubExp]
nes' <- [(SubExp, VName)]
-> ((SubExp, VName) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([SubExp] -> Closure -> [(SubExp, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
nes Closure
arrs) (((SubExp, VName) -> InternaliseM SubExp) -> InternaliseM [SubExp])
-> ((SubExp, VName) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \(SubExp
ne', VName
arr') -> do
    Type
rowtype <- Int -> Type -> Type
forall shape u.
ArrayShape shape =>
Int -> TypeBase shape u -> TypeBase shape u
I.stripArray Int
1 (Type -> Type) -> InternaliseM Type -> InternaliseM Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
arr'
    ErrorMsg SubExp
-> SrcLoc -> Type -> String -> SubExp -> InternaliseM SubExp
ensureShape
      ErrorMsg SubExp
"Row shape of input array does not match shape of neutral element"
      SrcLoc
loc Type
rowtype (String
whatString -> String -> String
forall a. [a] -> [a] -> [a]
++String
"_ne_right_shape") SubExp
ne'
  [Type]
nests <- (SubExp -> InternaliseM Type) -> [SubExp] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
I.subExpType [SubExp]
nes'
  [Type]
arrts <- (VName -> InternaliseM Type) -> Closure -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType Closure
arrs
  Lambda
lam' <- InternaliseLambda
-> ExpBase Info VName -> [Type] -> [Type] -> InternaliseM Lambda
internaliseFoldLambda InternaliseLambda
internaliseLambda ExpBase Info VName
lam [Type]
nests [Type]
arrts
  SubExp
w <- Int -> [Type] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 ([Type] -> SubExp) -> InternaliseM [Type] -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM Type) -> Closure -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType Closure
arrs
  String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (ExpT SOACS -> InternaliseM [SubExp])
-> (SOAC SOACS -> ExpT SOACS)
-> SOAC SOACS
-> InternaliseM [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SOAC SOACS -> ExpT SOACS
forall lore. Op lore -> ExpT lore
I.Op (SOAC SOACS -> InternaliseM [SubExp])
-> InternaliseM (SOAC SOACS) -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SubExp
-> Lambda -> [SubExp] -> Closure -> InternaliseM (SOAC SOACS)
f SubExp
w Lambda
lam' [SubExp]
nes' Closure
arrs

internaliseHist :: String
                -> E.Exp -> E.Exp -> E.Exp -> E.Exp -> E.Exp -> E.Exp -> SrcLoc
                -> InternaliseM [SubExp]
internaliseHist :: String
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> SrcLoc
-> InternaliseM [SubExp]
internaliseHist String
desc ExpBase Info VName
rf ExpBase Info VName
hist ExpBase Info VName
op ExpBase Info VName
ne ExpBase Info VName
buckets ExpBase Info VName
img SrcLoc
loc = do
  SubExp
rf' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"hist_rf" ExpBase Info VName
rf
  [SubExp]
ne' <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"hist_ne" ExpBase Info VName
ne
  Closure
hist' <- String -> ExpBase Info VName -> InternaliseM Closure
internaliseExpToVars String
"hist_hist" ExpBase Info VName
hist
  VName
buckets' <- String -> Exp (Lore InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m VName
letExp String
"hist_buckets" (ExpT SOACS -> InternaliseM VName)
-> (SubExp -> ExpT SOACS) -> SubExp -> InternaliseM VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS)
-> (SubExp -> BasicOp) -> SubExp -> ExpT SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
SubExp (SubExp -> InternaliseM VName)
-> InternaliseM SubExp -> InternaliseM VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<
              String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"hist_buckets" ExpBase Info VName
buckets
  Closure
img' <- String -> ExpBase Info VName -> InternaliseM Closure
internaliseExpToVars String
"hist_img" ExpBase Info VName
img

  -- reshape neutral element to have same size as the destination array
  [SubExp]
ne_shp <- [(SubExp, VName)]
-> ((SubExp, VName) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([SubExp] -> Closure -> [(SubExp, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
ne' Closure
hist') (((SubExp, VName) -> InternaliseM SubExp) -> InternaliseM [SubExp])
-> ((SubExp, VName) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \(SubExp
n, VName
h) -> do
    Type
rowtype <- Int -> Type -> Type
forall shape u.
ArrayShape shape =>
Int -> TypeBase shape u -> TypeBase shape u
I.stripArray Int
1 (Type -> Type) -> InternaliseM Type -> InternaliseM Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
h
    ErrorMsg SubExp
-> SrcLoc -> Type -> String -> SubExp -> InternaliseM SubExp
ensureShape
      ErrorMsg SubExp
"Row shape of destination array does not match shape of neutral element"
      SrcLoc
loc Type
rowtype String
"hist_ne_right_shape" SubExp
n
  [Type]
ne_ts <- (SubExp -> InternaliseM Type) -> [SubExp] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
I.subExpType [SubExp]
ne_shp
  [Type]
his_ts <- (VName -> InternaliseM Type) -> Closure -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType Closure
hist'
  Lambda
op' <- InternaliseLambda
-> ExpBase Info VName -> [Type] -> [Type] -> InternaliseM Lambda
internaliseFoldLambda InternaliseLambda
internaliseLambda ExpBase Info VName
op [Type]
ne_ts [Type]
his_ts

  -- reshape return type of bucket function to have same size as neutral element
  -- (modulo the index)
  Param Type
bucket_param <- String -> Type -> InternaliseM (Param Type)
forall (m :: * -> *) dec.
MonadFreshNames m =>
String -> dec -> m (Param dec)
newParam String
"bucket_p" (Type -> InternaliseM (Param Type))
-> Type -> InternaliseM (Param Type)
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int32
  [Param Type]
img_params <- (Type -> InternaliseM (Param Type))
-> [Type] -> InternaliseM [Param Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> Type -> InternaliseM (Param Type)
forall (m :: * -> *) dec.
MonadFreshNames m =>
String -> dec -> m (Param dec)
newParam String
"img_p" (Type -> InternaliseM (Param Type))
-> (Type -> Type) -> Type -> InternaliseM (Param Type)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
rowType) ([Type] -> InternaliseM [Param Type])
-> InternaliseM [Type] -> InternaliseM [Param Type]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (VName -> InternaliseM Type) -> Closure -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType Closure
img'
  let params :: [Param Type]
params = Param Type
bucket_param Param Type -> [Param Type] -> [Param Type]
forall a. a -> [a] -> [a]
: [Param Type]
img_params
      rettype :: [Type]
rettype = PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int32 Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
: [Type]
ne_ts
      body :: Body
body = Stms SOACS -> [SubExp] -> Body
forall lore. Bindable lore => Stms lore -> [SubExp] -> Body lore
mkBody Stms SOACS
forall a. Monoid a => a
mempty ([SubExp] -> Body) -> [SubExp] -> Body
forall a b. (a -> b) -> a -> b
$ (Param Type -> SubExp) -> [Param Type] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
I.Var (VName -> SubExp) -> (Param Type -> VName) -> Param Type -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> VName
forall dec. Param dec -> VName
paramName) [Param Type]
params
  Body
body' <- Scope SOACS -> InternaliseM Body -> InternaliseM Body
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param Type] -> Scope SOACS
forall lore dec.
(LParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
scopeOfLParams [Param Type]
params) (InternaliseM Body -> InternaliseM Body)
-> InternaliseM Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$
           ErrorMsg SubExp -> SrcLoc -> [Type] -> Body -> InternaliseM Body
ensureResultShape
           ErrorMsg SubExp
"Row shape of value array does not match row shape of hist target"
           (ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
img) [Type]
rettype Body
body

  -- get sizes of histogram and image arrays
  SubExp
w_hist <- Int -> [Type] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 ([Type] -> SubExp) -> InternaliseM [Type] -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM Type) -> Closure -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType Closure
hist'
  SubExp
w_img <- Int -> [Type] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 ([Type] -> SubExp) -> InternaliseM [Type] -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM Type) -> Closure -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType Closure
img'

  -- Generate an assertion and reshapes to ensure that buckets' and
  -- img' are the same size.
  Shape
b_shape <- Type -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape (Type -> Shape) -> InternaliseM Type -> InternaliseM Shape
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
buckets'
  let b_w :: SubExp
b_w = Int -> Shape -> SubExp
shapeSize Int
0 Shape
b_shape
  SubExp
cmp <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"bucket_cmp" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
I.int32) SubExp
b_w SubExp
w_img
  Certificates
c <- String
-> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certificates
assert String
"bucket_cert" SubExp
cmp
       ErrorMsg SubExp
"length of index and value array does not match" SrcLoc
loc
  VName
buckets'' <- Certificates -> InternaliseM VName -> InternaliseM VName
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
c (InternaliseM VName -> InternaliseM VName)
-> InternaliseM VName -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$ String -> Exp (Lore InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m VName
letExp (VName -> String
baseString VName
buckets') (Exp (Lore InternaliseM) -> InternaliseM VName)
-> Exp (Lore InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$
    BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
I.Reshape (ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp
reshapeOuter [SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimCoercion SubExp
w_img] Int
1 Shape
b_shape) VName
buckets'

  String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Op SOACS -> ExpT SOACS
forall lore. Op lore -> ExpT lore
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
    SubExp -> [HistOp SOACS] -> Lambda -> Closure -> SOAC SOACS
forall lore.
SubExp -> [HistOp lore] -> Lambda lore -> Closure -> SOAC lore
I.Hist SubExp
w_img [SubExp -> SubExp -> Closure -> [SubExp] -> Lambda -> HistOp SOACS
forall lore.
SubExp
-> SubExp -> Closure -> [SubExp] -> Lambda lore -> HistOp lore
HistOp SubExp
w_hist SubExp
rf' Closure
hist' [SubExp]
ne_shp Lambda
op'] ([LParam] -> Body -> [Type] -> Lambda
forall lore. [LParam lore] -> BodyT lore -> [Type] -> LambdaT lore
I.Lambda [Param Type]
[LParam]
params Body
body' [Type]
rettype) (Closure -> SOAC SOACS) -> Closure -> SOAC SOACS
forall a b. (a -> b) -> a -> b
$ VName
buckets'' VName -> Closure -> Closure
forall a. a -> [a] -> [a]
: Closure
img'

internaliseStreamMap :: String -> StreamOrd -> E.Exp -> E.Exp
                     -> InternaliseM [SubExp]
internaliseStreamMap :: String
-> StreamOrd
-> ExpBase Info VName
-> ExpBase Info VName
-> InternaliseM [SubExp]
internaliseStreamMap String
desc StreamOrd
o ExpBase Info VName
lam ExpBase Info VName
arr = do
  Closure
arrs <- String -> ExpBase Info VName -> InternaliseM Closure
internaliseExpToVars String
"stream_input" ExpBase Info VName
arr
  Lambda
lam' <- InternaliseLambda
-> ExpBase Info VName -> [SubExp] -> InternaliseM Lambda
internaliseStreamMapLambda InternaliseLambda
internaliseLambda ExpBase Info VName
lam ([SubExp] -> InternaliseM Lambda)
-> [SubExp] -> InternaliseM Lambda
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> Closure -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var Closure
arrs
  SubExp
w <- Int -> [Type] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 ([Type] -> SubExp) -> InternaliseM [Type] -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM Type) -> Closure -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType Closure
arrs
  let form :: StreamForm SOACS
form = StreamOrd
-> Commutativity -> Lambda -> [SubExp] -> StreamForm SOACS
forall lore.
StreamOrd
-> Commutativity -> Lambda lore -> [SubExp] -> StreamForm lore
I.Parallel StreamOrd
o Commutativity
Commutative ([LParam] -> Body -> [Type] -> Lambda
forall lore. [LParam lore] -> BodyT lore -> [Type] -> LambdaT lore
I.Lambda [] (Stms SOACS -> [SubExp] -> Body
forall lore. Bindable lore => Stms lore -> [SubExp] -> Body lore
mkBody Stms SOACS
forall a. Monoid a => a
mempty []) []) []
  String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Op SOACS -> ExpT SOACS
forall lore. Op lore -> ExpT lore
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> StreamForm SOACS -> Lambda -> Closure -> SOAC SOACS
forall lore.
SubExp -> StreamForm lore -> Lambda lore -> Closure -> SOAC lore
I.Stream SubExp
w StreamForm SOACS
form Lambda
lam' Closure
arrs

internaliseStreamRed :: String -> StreamOrd -> Commutativity
                     -> E.Exp -> E.Exp -> E.Exp
                     -> InternaliseM [SubExp]
internaliseStreamRed :: String
-> StreamOrd
-> Commutativity
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> InternaliseM [SubExp]
internaliseStreamRed String
desc StreamOrd
o Commutativity
comm ExpBase Info VName
lam0 ExpBase Info VName
lam ExpBase Info VName
arr = do
  Closure
arrs <- String -> ExpBase Info VName -> InternaliseM Closure
internaliseExpToVars String
"stream_input" ExpBase Info VName
arr
  [Type]
rowts <- (VName -> InternaliseM Type) -> Closure -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Type -> Type) -> InternaliseM Type -> InternaliseM Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
I.rowType (InternaliseM Type -> InternaliseM Type)
-> (VName -> InternaliseM Type) -> VName -> InternaliseM Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType) Closure
arrs
  ([Param Type]
lam_params, Body
lam_body) <-
    InternaliseLambda
-> ExpBase Info VName -> [Type] -> InternaliseM ([LParam], Body)
internaliseStreamLambda InternaliseLambda
internaliseLambda ExpBase Info VName
lam [Type]
rowts
  let (Param Type
chunk_param, [Param Type]
_, [Param Type]
lam_val_params) =
        Int -> [Param Type] -> (Param Type, [Param Type], [Param Type])
forall dec.
Int -> [Param dec] -> (Param dec, [Param dec], [Param dec])
partitionChunkedFoldParameters Int
0 [Param Type]
lam_params

  -- Synthesize neutral elements by applying the fold function
  -- to an empty chunk.
  Closure -> Exp (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
Closure -> Exp (Lore m) -> m ()
letBindNames [Param Type -> VName
forall dec. Param dec -> VName
I.paramName Param Type
chunk_param] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
    BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ Int32 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int32
0::Int32)
  [Param Type] -> (Param Type -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Param Type]
lam_val_params ((Param Type -> InternaliseM ()) -> InternaliseM ())
-> (Param Type -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \Param Type
p ->
    Closure -> Exp (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
Closure -> Exp (Lore m) -> m ()
letBindNames [Param Type -> VName
forall dec. Param dec -> VName
I.paramName Param Type
p] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
    BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ PrimType -> [SubExp] -> BasicOp
I.Scratch (Type -> PrimType
forall shape u. TypeBase shape u -> PrimType
I.elemType (Type -> PrimType) -> Type -> PrimType
forall a b. (a -> b) -> a -> b
$ Param Type -> Type
forall dec. Typed dec => Param dec -> Type
I.paramType Param Type
p) ([SubExp] -> BasicOp) -> [SubExp] -> BasicOp
forall a b. (a -> b) -> a -> b
$
    Type -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
I.arrayDims (Type -> [SubExp]) -> Type -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Param Type -> Type
forall dec. Typed dec => Param dec -> Type
I.paramType Param Type
p
  [SubExp]
nes <- Body -> InternaliseM [SubExp]
forall (m :: * -> *). MonadBinder m => Body (Lore m) -> m [SubExp]
bodyBind (Body -> InternaliseM [SubExp])
-> InternaliseM Body -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Body -> InternaliseM Body
forall lore (m :: * -> *).
(Renameable lore, MonadFreshNames m) =>
Body lore -> m (Body lore)
renameBody Body
lam_body

  [Type]
nes_ts <- (SubExp -> InternaliseM Type) -> [SubExp] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
I.subExpType [SubExp]
nes
  SubExp
outsz  <- Int -> [Type] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 ([Type] -> SubExp) -> InternaliseM [Type] -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM Type) -> Closure -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType Closure
arrs
  let acc_arr_tps :: [Type]
acc_arr_tps = [ Type -> Shape -> NoUniqueness -> Type
forall shape u_unused u.
ArrayShape shape =>
TypeBase shape u_unused -> shape -> u -> TypeBase shape u
I.arrayOf Type
t ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
I.Shape [SubExp
outsz]) NoUniqueness
NoUniqueness | Type
t <- [Type]
nes_ts ]
  Lambda
lam0' <- InternaliseLambda
-> ExpBase Info VName -> [Type] -> [Type] -> InternaliseM Lambda
internaliseFoldLambda InternaliseLambda
internaliseLambda ExpBase Info VName
lam0 [Type]
nes_ts [Type]
acc_arr_tps

  let lam0_acc_params :: [Param Type]
lam0_acc_params = Int -> [Param Type] -> [Param Type]
forall a. Int -> [a] -> [a]
take ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
nes) ([Param Type] -> [Param Type]) -> [Param Type] -> [Param Type]
forall a b. (a -> b) -> a -> b
$ Lambda -> [LParam]
forall lore. LambdaT lore -> [LParam lore]
I.lambdaParams Lambda
lam0'
  [Param Type]
lam_acc_params <- [Param Type]
-> (Param Type -> InternaliseM (Param Type))
-> InternaliseM [Param Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Param Type]
lam0_acc_params ((Param Type -> InternaliseM (Param Type))
 -> InternaliseM [Param Type])
-> (Param Type -> InternaliseM (Param Type))
-> InternaliseM [Param Type]
forall a b. (a -> b) -> a -> b
$ \Param Type
p -> do
    VName
name <- String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName (String -> InternaliseM VName) -> String -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$ VName -> String
baseString (VName -> String) -> VName -> String
forall a b. (a -> b) -> a -> b
$ Param Type -> VName
forall dec. Param dec -> VName
I.paramName Param Type
p
    Param Type -> InternaliseM (Param Type)
forall (m :: * -> *) a. Monad m => a -> m a
return Param Type
p { paramName :: VName
I.paramName = VName
name }

  -- Make sure the chunk size parameter comes first.
  let lam_params' :: [Param Type]
lam_params' = Param Type
chunk_param Param Type -> [Param Type] -> [Param Type]
forall a. a -> [a] -> [a]
: [Param Type]
lam_acc_params [Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
++ [Param Type]
lam_val_params

  Body
body_with_lam0 <-
    ErrorMsg SubExp -> SrcLoc -> [Type] -> Body -> InternaliseM Body
ensureResultShape ErrorMsg SubExp
"shape of result does not match shape of initial value"
    (ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
lam0) [Type]
nes_ts (Body -> InternaliseM Body)
-> (InternaliseM Body -> InternaliseM Body)
-> InternaliseM Body
-> InternaliseM Body
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< InternaliseM Body -> InternaliseM Body
forall (m :: * -> *).
MonadBinder m =>
m (Body (Lore m)) -> m (Body (Lore m))
insertStmsM (InternaliseM Body -> InternaliseM Body)
-> InternaliseM Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$ Scope SOACS -> InternaliseM Body -> InternaliseM Body
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param Type] -> Scope SOACS
forall lore dec.
(LParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
scopeOfLParams [Param Type]
lam_params') (InternaliseM Body -> InternaliseM Body)
-> InternaliseM Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$ do
      [SubExp]
lam_res <- Body (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *). MonadBinder m => Body (Lore m) -> m [SubExp]
bodyBind Body (Lore InternaliseM)
Body
lam_body
      [SubExp]
lam_res' <- ErrorMsg SubExp
-> SrcLoc -> Closure -> [Type] -> [SubExp] -> InternaliseM [SubExp]
forall u.
Typed (TypeBase Shape u) =>
ErrorMsg SubExp
-> SrcLoc
-> Closure
-> [TypeBase Shape u]
-> [SubExp]
-> InternaliseM [SubExp]
ensureArgShapes
                  ErrorMsg SubExp
"shape of chunk function result does not match shape of initial value"
                  (ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
lam) [] ((Param Type -> Type) -> [Param Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> Type
forall t. Typed t => t -> Type
I.typeOf ([Param Type] -> [Type]) -> [Param Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ Lambda -> [LParam]
forall lore. LambdaT lore -> [LParam lore]
I.lambdaParams Lambda
lam0') [SubExp]
lam_res
      [SubExp]
new_lam_res <- Lambda (Lore InternaliseM)
-> [InternaliseM (Exp (Lore InternaliseM))]
-> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
Lambda (Lore m) -> [m (Exp (Lore m))] -> m [SubExp]
eLambda Lambda (Lore InternaliseM)
Lambda
lam0' ([InternaliseM (Exp (Lore InternaliseM))] -> InternaliseM [SubExp])
-> [InternaliseM (Exp (Lore InternaliseM))]
-> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ (SubExp -> InternaliseM (ExpT SOACS))
-> [SubExp] -> [InternaliseM (ExpT SOACS)]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> InternaliseM (ExpT SOACS)
forall (m :: * -> *). MonadBinder m => SubExp -> m (Exp (Lore m))
eSubExp ([SubExp] -> [InternaliseM (ExpT SOACS)])
-> [SubExp] -> [InternaliseM (ExpT SOACS)]
forall a b. (a -> b) -> a -> b
$
                     (Param Type -> SubExp) -> [Param Type] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
I.Var (VName -> SubExp) -> (Param Type -> VName) -> Param Type -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> VName
forall dec. Param dec -> VName
paramName) [Param Type]
lam_acc_params [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
lam_res'
      Body -> InternaliseM Body
forall (m :: * -> *) a. Monad m => a -> m a
return (Body -> InternaliseM Body) -> Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [SubExp]
new_lam_res

  let form :: StreamForm SOACS
form = StreamOrd
-> Commutativity -> Lambda -> [SubExp] -> StreamForm SOACS
forall lore.
StreamOrd
-> Commutativity -> Lambda lore -> [SubExp] -> StreamForm lore
I.Parallel StreamOrd
o Commutativity
comm Lambda
lam0' [SubExp]
nes
      lam' :: Lambda
lam' = Lambda :: forall lore. [LParam lore] -> BodyT lore -> [Type] -> LambdaT lore
I.Lambda { lambdaParams :: [LParam]
lambdaParams = [Param Type]
[LParam]
lam_params'
                      , lambdaBody :: Body
lambdaBody = Body
body_with_lam0
                      , lambdaReturnType :: [Type]
lambdaReturnType = [Type]
nes_ts }
  SubExp
w <- Int -> [Type] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 ([Type] -> SubExp) -> InternaliseM [Type] -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM Type) -> Closure -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType Closure
arrs
  String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Op SOACS -> ExpT SOACS
forall lore. Op lore -> ExpT lore
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> StreamForm SOACS -> Lambda -> Closure -> SOAC SOACS
forall lore.
SubExp -> StreamForm lore -> Lambda lore -> Closure -> SOAC lore
I.Stream SubExp
w StreamForm SOACS
form Lambda
lam' Closure
arrs

internaliseExp1 :: String -> E.Exp -> InternaliseM I.SubExp
internaliseExp1 :: String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
desc ExpBase Info VName
e = do
  [SubExp]
vs <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
e
  case [SubExp]
vs of [SubExp
se] -> SubExp -> InternaliseM SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return SubExp
se
             [SubExp]
_ -> String -> InternaliseM SubExp
forall a. HasCallStack => String -> a
error String
"Internalise.internaliseExp1: was passed not just a single subexpression"

-- | Promote to dimension type as appropriate for the original type.
-- Also return original type.
internaliseDimExp :: String -> E.Exp -> InternaliseM (I.SubExp, IntType)
internaliseDimExp :: String -> ExpBase Info VName -> InternaliseM (SubExp, IntType)
internaliseDimExp String
s ExpBase Info VName
e = do
  SubExp
e' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
s ExpBase Info VName
e
  case ExpBase Info VName -> PatternType
E.typeOf ExpBase Info VName
e of
    E.Scalar (E.Prim (Signed IntType
it))   -> (,IntType
it) (SubExp -> (SubExp, IntType))
-> InternaliseM SubExp -> InternaliseM (SubExp, IntType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int32 SubExp
e'
    E.Scalar (E.Prim (Unsigned IntType
it)) -> (,IntType
it) (SubExp -> (SubExp, IntType))
-> InternaliseM SubExp -> InternaliseM (SubExp, IntType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
IntType -> SubExp -> m SubExp
asIntZ IntType
Int32 SubExp
e'
    PatternType
_                               -> String -> InternaliseM (SubExp, IntType)
forall a. HasCallStack => String -> a
error String
"internaliseDimExp: bad type"

internaliseExpToVars :: String -> E.Exp -> InternaliseM [I.VName]
internaliseExpToVars :: String -> ExpBase Info VName -> InternaliseM Closure
internaliseExpToVars String
desc ExpBase Info VName
e =
  (SubExp -> InternaliseM VName) -> [SubExp] -> InternaliseM Closure
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM VName
asIdent ([SubExp] -> InternaliseM Closure)
-> InternaliseM [SubExp] -> InternaliseM Closure
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
e
  where asIdent :: SubExp -> InternaliseM VName
asIdent (I.Var VName
v) = VName -> InternaliseM VName
forall (m :: * -> *) a. Monad m => a -> m a
return VName
v
        asIdent SubExp
se        = String -> Exp (Lore InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m VName
letExp String
desc (Exp (Lore InternaliseM) -> InternaliseM VName)
-> Exp (Lore InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
se

internaliseOperation :: String
                     -> E.Exp
                     -> (I.VName -> InternaliseM I.BasicOp)
                     -> InternaliseM [I.SubExp]
internaliseOperation :: String
-> ExpBase Info VName
-> (VName -> InternaliseM BasicOp)
-> InternaliseM [SubExp]
internaliseOperation String
s ExpBase Info VName
e VName -> InternaliseM BasicOp
op = do
  Closure
vs <- String -> ExpBase Info VName -> InternaliseM Closure
internaliseExpToVars String
s ExpBase Info VName
e
  String -> [Exp (Lore InternaliseM)] -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> [Exp (Lore m)] -> m [SubExp]
letSubExps String
s ([ExpT SOACS] -> InternaliseM [SubExp])
-> InternaliseM [ExpT SOACS] -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (VName -> InternaliseM (ExpT SOACS))
-> Closure -> InternaliseM [ExpT SOACS]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((BasicOp -> ExpT SOACS)
-> InternaliseM BasicOp -> InternaliseM (ExpT SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (InternaliseM BasicOp -> InternaliseM (ExpT SOACS))
-> (VName -> InternaliseM BasicOp)
-> VName
-> InternaliseM (ExpT SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> InternaliseM BasicOp
op) Closure
vs

certifyingNonzero :: SrcLoc -> IntType -> SubExp
                  -> InternaliseM a
                  -> InternaliseM a
certifyingNonzero :: SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
x InternaliseM a
m = do
  SubExp
zero <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"zero" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
          CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp (PrimType -> CmpOp
CmpEq (IntType -> PrimType
IntType IntType
t)) SubExp
x (IntType -> Integer -> SubExp
intConst IntType
t Integer
0)
  SubExp
nonzero <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"nonzero" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
UnOp UnOp
Not SubExp
zero
  Certificates
c <- String
-> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certificates
assert String
"nonzero_cert" SubExp
nonzero ErrorMsg SubExp
"division by zero" SrcLoc
loc
  Certificates -> InternaliseM a -> InternaliseM a
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
c InternaliseM a
m

certifyingNonnegative :: SrcLoc -> IntType -> SubExp
                      -> InternaliseM a
                      -> InternaliseM a
certifyingNonnegative :: SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonnegative SrcLoc
loc IntType
t SubExp
x InternaliseM a
m = do
  SubExp
nonnegative <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"nonnegative" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
                 CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp (IntType -> CmpOp
CmpSle IntType
t) (IntType -> Integer -> SubExp
intConst IntType
t Integer
0) SubExp
x
  Certificates
c <- String
-> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certificates
assert String
"nonzero_cert" SubExp
nonnegative ErrorMsg SubExp
"negative exponent" SrcLoc
loc
  Certificates -> InternaliseM a -> InternaliseM a
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
c InternaliseM a
m

internaliseBinOp :: SrcLoc -> String
                 -> E.BinOp
                 -> I.SubExp -> I.SubExp
                 -> E.PrimType
                 -> E.PrimType
                 -> InternaliseM [I.SubExp]
internaliseBinOp :: SrcLoc
-> String
-> BinOp
-> SubExp
-> SubExp
-> PrimType
-> PrimType
-> InternaliseM [SubExp]
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Plus SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> Overflow -> BinOp
I.Add IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Plus SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> Overflow -> BinOp
I.Add IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Plus SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (FloatType -> BinOp
I.FAdd FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Minus SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> Overflow -> BinOp
I.Sub IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Minus SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> Overflow -> BinOp
I.Sub IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Minus SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (FloatType -> BinOp
I.FSub FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Times SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> Overflow -> BinOp
I.Mul IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Times SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> Overflow -> BinOp
I.Mul IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Times SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (FloatType -> BinOp
I.FMul FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc String
desc BinOp
E.Divide SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  SrcLoc
-> IntType
-> SubExp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> Safety -> BinOp
I.SDiv IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc String
desc BinOp
E.Divide SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  SrcLoc
-> IntType
-> SubExp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> Safety -> BinOp
I.UDiv IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Divide SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (FloatType -> BinOp
I.FDiv FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Pow SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (FloatType -> BinOp
I.FPow FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc String
desc BinOp
E.Pow SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  SrcLoc
-> IntType
-> SubExp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonnegative SrcLoc
loc IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> BinOp
I.Pow IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Pow SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> BinOp
I.Pow IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc String
desc BinOp
E.Mod SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  SrcLoc
-> IntType
-> SubExp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> Safety -> BinOp
I.SMod IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc String
desc BinOp
E.Mod SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  SrcLoc
-> IntType
-> SubExp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> Safety -> BinOp
I.UMod IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Mod SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (FloatType -> BinOp
I.FMod FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc String
desc BinOp
E.Quot SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  SrcLoc
-> IntType
-> SubExp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> Safety -> BinOp
I.SQuot IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc String
desc BinOp
E.Quot SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  SrcLoc
-> IntType
-> SubExp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> Safety -> BinOp
I.UDiv IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc String
desc BinOp
E.Rem SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  SrcLoc
-> IntType
-> SubExp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> Safety -> BinOp
I.SRem IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc String
desc BinOp
E.Rem SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  SrcLoc
-> IntType
-> SubExp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> Safety -> BinOp
I.UMod IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.ShiftR SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> BinOp
I.AShr IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.ShiftR SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> BinOp
I.LShr IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.ShiftL SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> BinOp
I.Shl IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.ShiftL SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> BinOp
I.Shl IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Band SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> BinOp
I.And IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Band SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> BinOp
I.And IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Xor SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> BinOp
I.Xor IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Xor SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> BinOp
I.Xor IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Bor SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> BinOp
I.Or IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Bor SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc (IntType -> BinOp
I.Or IntType
t) SubExp
x SubExp
y

internaliseBinOp SrcLoc
_ String
desc BinOp
E.Equal SubExp
x SubExp
y PrimType
t PrimType
_ =
  String -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp String
desc (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimType
internalisePrimType PrimType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.NotEqual SubExp
x SubExp
y PrimType
t PrimType
_ = do
  SubExp
eq <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp (String
descString -> String -> String
forall a. [a] -> [a] -> [a]
++String
"true") (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimType
internalisePrimType PrimType
t) SubExp
x SubExp
y
  (SubExp -> [SubExp])
-> InternaliseM SubExp -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> [SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InternaliseM SubExp -> InternaliseM [SubExp])
-> InternaliseM SubExp -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
desc (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
I.Not SubExp
eq
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Less SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  String -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp String
desc (IntType -> CmpOp
I.CmpSlt IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Less SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  String -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp String
desc (IntType -> CmpOp
I.CmpUlt IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Leq SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  String -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp String
desc (IntType -> CmpOp
I.CmpSle IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Leq SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  String -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp String
desc (IntType -> CmpOp
I.CmpUle IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Greater SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  String -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp String
desc (IntType -> CmpOp
I.CmpSlt IntType
t) SubExp
y SubExp
x -- Note the swapped x and y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Greater SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  String -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp String
desc (IntType -> CmpOp
I.CmpUlt IntType
t) SubExp
y SubExp
x -- Note the swapped x and y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Geq SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  String -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp String
desc (IntType -> CmpOp
I.CmpSle IntType
t) SubExp
y SubExp
x -- Note the swapped x and y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Geq SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  String -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp String
desc (IntType -> CmpOp
I.CmpUle IntType
t) SubExp
y SubExp
x -- Note the swapped x and y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Less SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  String -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp String
desc (FloatType -> CmpOp
I.FCmpLt FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Leq SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  String -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp String
desc (FloatType -> CmpOp
I.FCmpLe FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Greater SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  String -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp String
desc (FloatType -> CmpOp
I.FCmpLt FloatType
t) SubExp
y SubExp
x -- Note the swapped x and y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Geq SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  String -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp String
desc (FloatType -> CmpOp
I.FCmpLe FloatType
t) SubExp
y SubExp
x -- Note the swapped x and y

-- Relational operators for booleans.
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Less SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
  String -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp String
desc CmpOp
I.CmpLlt SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Leq SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
  String -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp String
desc CmpOp
I.CmpLle SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Greater SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
  String -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp String
desc CmpOp
I.CmpLlt SubExp
y SubExp
x -- Note the swapped x and y
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Geq SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
  String -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp String
desc CmpOp
I.CmpLle SubExp
y SubExp
x -- Note the swapped x and y

internaliseBinOp SrcLoc
_ String
_ BinOp
op SubExp
_ SubExp
_ PrimType
t1 PrimType
t2 =
  String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error (String -> InternaliseM [SubExp])
-> String -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String
"Invalid binary operator " String -> String -> String
forall a. [a] -> [a] -> [a]
++ BinOp -> String
forall a. Pretty a => a -> String
pretty BinOp
op String -> String -> String
forall a. [a] -> [a] -> [a]
++
  String
" with operand types " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PrimType -> String
forall a. Pretty a => a -> String
pretty PrimType
t1 String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
", " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PrimType -> String
forall a. Pretty a => a -> String
pretty PrimType
t2

simpleBinOp :: String
            -> I.BinOp
            -> I.SubExp -> I.SubExp
            -> InternaliseM [I.SubExp]
simpleBinOp :: String -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp String
desc BinOp
bop SubExp
x SubExp
y =
  String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
bop SubExp
x SubExp
y

simpleCmpOp :: String
            -> I.CmpOp
            -> I.SubExp -> I.SubExp
            -> InternaliseM [I.SubExp]
simpleCmpOp :: String -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp String
desc CmpOp
op SubExp
x SubExp
y =
  String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp CmpOp
op SubExp
x SubExp
y

findFuncall :: E.Exp -> InternaliseM (E.QualName VName,
                                      [(E.Exp, Maybe VName)],
                                      E.StructType,
                                      [VName])
findFuncall :: ExpBase Info VName
-> InternaliseM
     (QualName VName, [(ExpBase Info VName, Maybe VName)], StructType,
      Closure)
findFuncall (E.Var QualName VName
fname (Info PatternType
t) SrcLoc
_) =
  (QualName VName, [(ExpBase Info VName, Maybe VName)], StructType,
 Closure)
-> InternaliseM
     (QualName VName, [(ExpBase Info VName, Maybe VName)], StructType,
      Closure)
forall (m :: * -> *) a. Monad m => a -> m a
return (QualName VName
fname, [], PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
t, [])
findFuncall (E.Apply ExpBase Info VName
f ExpBase Info VName
arg (Info (Diet
_, Maybe VName
argext)) (Info PatternType
ret, Info Closure
retext) SrcLoc
_) = do
  (QualName VName
fname, [(ExpBase Info VName, Maybe VName)]
args, StructType
_, Closure
_) <- ExpBase Info VName
-> InternaliseM
     (QualName VName, [(ExpBase Info VName, Maybe VName)], StructType,
      Closure)
findFuncall ExpBase Info VName
f
  (QualName VName, [(ExpBase Info VName, Maybe VName)], StructType,
 Closure)
-> InternaliseM
     (QualName VName, [(ExpBase Info VName, Maybe VName)], StructType,
      Closure)
forall (m :: * -> *) a. Monad m => a -> m a
return (QualName VName
fname, [(ExpBase Info VName, Maybe VName)]
args [(ExpBase Info VName, Maybe VName)]
-> [(ExpBase Info VName, Maybe VName)]
-> [(ExpBase Info VName, Maybe VName)]
forall a. [a] -> [a] -> [a]
++ [(ExpBase Info VName
arg, Maybe VName
argext)], PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret, Closure
retext)
findFuncall ExpBase Info VName
e =
  String
-> InternaliseM
     (QualName VName, [(ExpBase Info VName, Maybe VName)], StructType,
      Closure)
forall a. HasCallStack => String -> a
error (String
 -> InternaliseM
      (QualName VName, [(ExpBase Info VName, Maybe VName)], StructType,
       Closure))
-> String
-> InternaliseM
     (QualName VName, [(ExpBase Info VName, Maybe VName)], StructType,
      Closure)
forall a b. (a -> b) -> a -> b
$ String
"Invalid function expression in application: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ ExpBase Info VName -> String
forall a. Pretty a => a -> String
pretty ExpBase Info VName
e

internaliseLambda :: InternaliseLambda

internaliseLambda :: InternaliseLambda
internaliseLambda (E.Parens ExpBase Info VName
e SrcLoc
_) [Type]
rowtypes =
  InternaliseLambda
internaliseLambda ExpBase Info VName
e [Type]
rowtypes

internaliseLambda (E.Lambda [Pattern]
params ExpBase Info VName
body Maybe (TypeExp VName)
_ (Info (Aliasing
_, StructType
rettype)) SrcLoc
_) [Type]
rowtypes =
  [Pattern]
-> [Type]
-> ([LParam] -> InternaliseM ([Param Type], Body, [ExtType]))
-> InternaliseM ([Param Type], Body, [ExtType])
forall a.
[Pattern]
-> [Type] -> ([LParam] -> InternaliseM a) -> InternaliseM a
bindingLambdaParams [Pattern]
params [Type]
rowtypes (([LParam] -> InternaliseM ([Param Type], Body, [ExtType]))
 -> InternaliseM ([Param Type], Body, [ExtType]))
-> ([LParam] -> InternaliseM ([Param Type], Body, [ExtType]))
-> InternaliseM ([Param Type], Body, [ExtType])
forall a b. (a -> b) -> a -> b
$ \[LParam]
params' -> do
    [DeclExtType]
rettype' <- StructType -> InternaliseM [DeclExtType]
internaliseReturnType StructType
rettype
    Body
body' <- ExpBase Info VName -> InternaliseM Body
internaliseBody ExpBase Info VName
body
    ([Param Type], Body, [ExtType])
-> InternaliseM ([Param Type], Body, [ExtType])
forall (m :: * -> *) a. Monad m => a -> m a
return ([Param Type]
[LParam]
params', Body
body', (DeclExtType -> ExtType) -> [DeclExtType] -> [ExtType]
forall a b. (a -> b) -> [a] -> [b]
map DeclExtType -> ExtType
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl [DeclExtType]
rettype')

internaliseLambda ExpBase Info VName
e [Type]
_ = String -> InternaliseM ([Param Type], Body, [ExtType])
forall a. HasCallStack => String -> a
error (String -> InternaliseM ([Param Type], Body, [ExtType]))
-> String -> InternaliseM ([Param Type], Body, [ExtType])
forall a b. (a -> b) -> a -> b
$ String
"internaliseLambda: unexpected expression:\n" String -> String -> String
forall a. [a] -> [a] -> [a]
++ ExpBase Info VName -> String
forall a. Pretty a => a -> String
pretty ExpBase Info VName
e

-- | Some operators and functions are overloaded or otherwise special
-- - we detect and treat them here.
isOverloadedFunction :: E.QualName VName -> [E.Exp] -> SrcLoc
                     -> Maybe (String -> InternaliseM [SubExp])
isOverloadedFunction :: QualName VName
-> [ExpBase Info VName]
-> SrcLoc
-> Maybe (String -> InternaliseM [SubExp])
isOverloadedFunction QualName VName
qname [ExpBase Info VName]
args SrcLoc
loc = do
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ VName -> Int
baseTag (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qname) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag
  let handlers :: [[ExpBase Info VName]
 -> String -> Maybe (String -> InternaliseM [SubExp])]
handlers = [[ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM [SubExp])
forall a.
(Eq a, IsString a) =>
[ExpBase Info VName]
-> a -> Maybe (String -> InternaliseM [SubExp])
handleSign,
                  [ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM [SubExp])
forall (f :: * -> *).
Applicative f =>
[ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM (f SubExp))
handleIntrinsicOps,
                  [ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM [SubExp])
handleOps,
                  [ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM [SubExp])
handleSOACs,
                  [ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM [SubExp])
handleRest]
  [Maybe (String -> InternaliseM [SubExp])]
-> Maybe (String -> InternaliseM [SubExp])
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, MonadPlus m) =>
t (m a) -> m a
msum [[ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM [SubExp])
h [ExpBase Info VName]
args (String -> Maybe (String -> InternaliseM [SubExp]))
-> String -> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ VName -> String
baseString (VName -> String) -> VName -> String
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qname | [ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM [SubExp])
h <- [[ExpBase Info VName]
 -> String -> Maybe (String -> InternaliseM [SubExp])]
handlers]
  where
    handleSign :: [ExpBase Info VName]
-> a -> Maybe (String -> InternaliseM [SubExp])
handleSign [ExpBase Info VName
x] a
"sign_i8"  = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> ExpBase Info VName -> String -> InternaliseM [SubExp]
toSigned IntType
I.Int8 ExpBase Info VName
x
    handleSign [ExpBase Info VName
x] a
"sign_i16" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> ExpBase Info VName -> String -> InternaliseM [SubExp]
toSigned IntType
I.Int16 ExpBase Info VName
x
    handleSign [ExpBase Info VName
x] a
"sign_i32" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> ExpBase Info VName -> String -> InternaliseM [SubExp]
toSigned IntType
I.Int32 ExpBase Info VName
x
    handleSign [ExpBase Info VName
x] a
"sign_i64" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> ExpBase Info VName -> String -> InternaliseM [SubExp]
toSigned IntType
I.Int64 ExpBase Info VName
x

    handleSign [ExpBase Info VName
x] a
"unsign_i8"  = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> ExpBase Info VName -> String -> InternaliseM [SubExp]
toUnsigned IntType
I.Int8 ExpBase Info VName
x
    handleSign [ExpBase Info VName
x] a
"unsign_i16" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> ExpBase Info VName -> String -> InternaliseM [SubExp]
toUnsigned IntType
I.Int16 ExpBase Info VName
x
    handleSign [ExpBase Info VName
x] a
"unsign_i32" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> ExpBase Info VName -> String -> InternaliseM [SubExp]
toUnsigned IntType
I.Int32 ExpBase Info VName
x
    handleSign [ExpBase Info VName
x] a
"unsign_i64" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> ExpBase Info VName -> String -> InternaliseM [SubExp]
toUnsigned IntType
I.Int64 ExpBase Info VName
x

    handleSign [ExpBase Info VName]
_ a
_ = Maybe (String -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

    handleIntrinsicOps :: [ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM (f SubExp))
handleIntrinsicOps [ExpBase Info VName
x] String
s
      | Just UnOp
unop <- (UnOp -> Bool) -> [UnOp] -> Maybe UnOp
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((String -> String -> Bool
forall a. Eq a => a -> a -> Bool
==String
s) (String -> Bool) -> (UnOp -> String) -> UnOp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UnOp -> String
forall a. Pretty a => a -> String
pretty) [UnOp]
allUnOps = (String -> InternaliseM (f SubExp))
-> Maybe (String -> InternaliseM (f SubExp))
forall a. a -> Maybe a
Just ((String -> InternaliseM (f SubExp))
 -> Maybe (String -> InternaliseM (f SubExp)))
-> (String -> InternaliseM (f SubExp))
-> Maybe (String -> InternaliseM (f SubExp))
forall a b. (a -> b) -> a -> b
$ \String
desc -> do
          SubExp
x' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"x" ExpBase Info VName
x
          (SubExp -> f SubExp)
-> InternaliseM SubExp -> InternaliseM (f SubExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> f SubExp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InternaliseM SubExp -> InternaliseM (f SubExp))
-> InternaliseM SubExp -> InternaliseM (f SubExp)
forall a b. (a -> b) -> a -> b
$ String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
desc (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
unop SubExp
x'
    handleIntrinsicOps [TupLit [ExpBase Info VName
x,ExpBase Info VName
y] SrcLoc
_] String
s
      | Just BinOp
bop <- (BinOp -> Bool) -> [BinOp] -> Maybe BinOp
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((String -> String -> Bool
forall a. Eq a => a -> a -> Bool
==String
s) (String -> Bool) -> (BinOp -> String) -> BinOp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BinOp -> String
forall a. Pretty a => a -> String
pretty) [BinOp]
allBinOps = (String -> InternaliseM (f SubExp))
-> Maybe (String -> InternaliseM (f SubExp))
forall a. a -> Maybe a
Just ((String -> InternaliseM (f SubExp))
 -> Maybe (String -> InternaliseM (f SubExp)))
-> (String -> InternaliseM (f SubExp))
-> Maybe (String -> InternaliseM (f SubExp))
forall a b. (a -> b) -> a -> b
$ \String
desc -> do
          SubExp
x' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"x" ExpBase Info VName
x
          SubExp
y' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"y" ExpBase Info VName
y
          (SubExp -> f SubExp)
-> InternaliseM SubExp -> InternaliseM (f SubExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> f SubExp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InternaliseM SubExp -> InternaliseM (f SubExp))
-> InternaliseM SubExp -> InternaliseM (f SubExp)
forall a b. (a -> b) -> a -> b
$ String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
desc (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
bop SubExp
x' SubExp
y'
      | Just CmpOp
cmp <- (CmpOp -> Bool) -> [CmpOp] -> Maybe CmpOp
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((String -> String -> Bool
forall a. Eq a => a -> a -> Bool
==String
s) (String -> Bool) -> (CmpOp -> String) -> CmpOp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CmpOp -> String
forall a. Pretty a => a -> String
pretty) [CmpOp]
allCmpOps = (String -> InternaliseM (f SubExp))
-> Maybe (String -> InternaliseM (f SubExp))
forall a. a -> Maybe a
Just ((String -> InternaliseM (f SubExp))
 -> Maybe (String -> InternaliseM (f SubExp)))
-> (String -> InternaliseM (f SubExp))
-> Maybe (String -> InternaliseM (f SubExp))
forall a b. (a -> b) -> a -> b
$ \String
desc -> do
          SubExp
x' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"x" ExpBase Info VName
x
          SubExp
y' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"y" ExpBase Info VName
y
          (SubExp -> f SubExp)
-> InternaliseM SubExp -> InternaliseM (f SubExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> f SubExp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InternaliseM SubExp -> InternaliseM (f SubExp))
-> InternaliseM SubExp -> InternaliseM (f SubExp)
forall a b. (a -> b) -> a -> b
$ String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
desc (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp CmpOp
cmp SubExp
x' SubExp
y'
    handleIntrinsicOps [ExpBase Info VName
x] String
s
      | Just ConvOp
conv <- (ConvOp -> Bool) -> [ConvOp] -> Maybe ConvOp
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((String -> String -> Bool
forall a. Eq a => a -> a -> Bool
==String
s) (String -> Bool) -> (ConvOp -> String) -> ConvOp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ConvOp -> String
forall a. Pretty a => a -> String
pretty) [ConvOp]
allConvOps = (String -> InternaliseM (f SubExp))
-> Maybe (String -> InternaliseM (f SubExp))
forall a. a -> Maybe a
Just ((String -> InternaliseM (f SubExp))
 -> Maybe (String -> InternaliseM (f SubExp)))
-> (String -> InternaliseM (f SubExp))
-> Maybe (String -> InternaliseM (f SubExp))
forall a b. (a -> b) -> a -> b
$ \String
desc -> do
          SubExp
x' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"x" ExpBase Info VName
x
          (SubExp -> f SubExp)
-> InternaliseM SubExp -> InternaliseM (f SubExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> f SubExp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InternaliseM SubExp -> InternaliseM (f SubExp))
-> InternaliseM SubExp -> InternaliseM (f SubExp)
forall a b. (a -> b) -> a -> b
$ String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
desc (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp ConvOp
conv SubExp
x'
    handleIntrinsicOps [ExpBase Info VName]
_ String
_ = Maybe (String -> InternaliseM (f SubExp))
forall a. Maybe a
Nothing

    -- Short-circuiting operators are magical.
    handleOps :: [ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM [SubExp])
handleOps [ExpBase Info VName
x,ExpBase Info VName
y] String
"&&" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc (ExpBase Info VName -> InternaliseM [SubExp])
-> ExpBase Info VName -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
      ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> (Info PatternType, Info Closure)
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> (f PatternType, f Closure)
-> SrcLoc
-> ExpBase f vn
E.If ExpBase Info VName
x ExpBase Info VName
y (PrimValue -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
E.Literal (Bool -> PrimValue
E.BoolValue Bool
False) SrcLoc
forall a. Monoid a => a
mempty) (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) Aliasing -> PatternType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar (ScalarTypeBase (DimDecl VName) Aliasing -> PatternType)
-> ScalarTypeBase (DimDecl VName) Aliasing -> PatternType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) Aliasing
forall dim as. PrimType -> ScalarTypeBase dim as
E.Prim PrimType
E.Bool, Closure -> Info Closure
forall a. a -> Info a
Info []) SrcLoc
forall a. Monoid a => a
mempty
    handleOps [ExpBase Info VName
x,ExpBase Info VName
y] String
"||" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
        String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc (ExpBase Info VName -> InternaliseM [SubExp])
-> ExpBase Info VName -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
        ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> (Info PatternType, Info Closure)
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> (f PatternType, f Closure)
-> SrcLoc
-> ExpBase f vn
E.If ExpBase Info VName
x (PrimValue -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
E.Literal (Bool -> PrimValue
E.BoolValue Bool
True) SrcLoc
forall a. Monoid a => a
mempty) ExpBase Info VName
y (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) Aliasing -> PatternType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar (ScalarTypeBase (DimDecl VName) Aliasing -> PatternType)
-> ScalarTypeBase (DimDecl VName) Aliasing -> PatternType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) Aliasing
forall dim as. PrimType -> ScalarTypeBase dim as
E.Prim PrimType
E.Bool, Closure -> Info Closure
forall a. a -> Info a
Info []) SrcLoc
forall a. Monoid a => a
mempty

    -- Handle equality and inequality specially, to treat the case of
    -- arrays.
    handleOps [ExpBase Info VName
xe,ExpBase Info VName
ye] String
op
      | Just String -> SubExp -> InternaliseM [SubExp]
cmp_f <- String -> Maybe (String -> SubExp -> InternaliseM [SubExp])
forall a (m :: * -> *).
(IsString a, MonadBinder m, Eq a) =>
a -> Maybe (String -> SubExp -> m [SubExp])
isEqlOp String
op = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc -> do
          [SubExp]
xe' <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"x" ExpBase Info VName
xe
          [SubExp]
ye' <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"y" ExpBase Info VName
ye
          [SubExp]
rs <- (SubExp -> SubExp -> InternaliseM SubExp)
-> [SubExp] -> [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM (String -> SubExp -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
(MonadBinder m, Bindable (Lore m), BinderOps (Lore m),
 Op (Lore m) ~ SOAC (Lore m)) =>
String -> SubExp -> SubExp -> m SubExp
doComparison String
desc) [SubExp]
xe' [SubExp]
ye'
          String -> SubExp -> InternaliseM [SubExp]
cmp_f String
desc (SubExp -> InternaliseM [SubExp])
-> InternaliseM SubExp -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"eq" (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [SubExp] -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => [SubExp] -> m (Exp (Lore m))
eAll [SubExp]
rs
        where isEqlOp :: a -> Maybe (String -> SubExp -> m [SubExp])
isEqlOp a
"!=" = (String -> SubExp -> m [SubExp])
-> Maybe (String -> SubExp -> m [SubExp])
forall a. a -> Maybe a
Just ((String -> SubExp -> m [SubExp])
 -> Maybe (String -> SubExp -> m [SubExp]))
-> (String -> SubExp -> m [SubExp])
-> Maybe (String -> SubExp -> m [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc SubExp
eq ->
                String -> Exp (Lore m) -> m [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore m) -> m [SubExp]) -> Exp (Lore m) -> m [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
I.Not SubExp
eq
              isEqlOp a
"==" = (String -> SubExp -> m [SubExp])
-> Maybe (String -> SubExp -> m [SubExp])
forall a. a -> Maybe a
Just ((String -> SubExp -> m [SubExp])
 -> Maybe (String -> SubExp -> m [SubExp]))
-> (String -> SubExp -> m [SubExp])
-> Maybe (String -> SubExp -> m [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
_ SubExp
eq ->
                [SubExp] -> m [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [SubExp
eq]
              isEqlOp a
_ = Maybe (String -> SubExp -> m [SubExp])
forall a. Maybe a
Nothing

              doComparison :: String -> SubExp -> SubExp -> m SubExp
doComparison String
desc SubExp
x SubExp
y = do
                Type
x_t <- SubExp -> m Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
I.subExpType SubExp
x
                Type
y_t <- SubExp -> m Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
I.subExpType SubExp
y
                case Type
x_t of
                  I.Prim PrimType
t -> String -> Exp (Lore m) -> m SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
desc (Exp (Lore m) -> m SubExp) -> Exp (Lore m) -> m SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
t) SubExp
x SubExp
y
                  Type
_ -> do
                    let x_dims :: [SubExp]
x_dims = Type -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
I.arrayDims Type
x_t
                        y_dims :: [SubExp]
y_dims = Type -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
I.arrayDims Type
y_t
                    [SubExp]
dims_match <- [(SubExp, SubExp)] -> ((SubExp, SubExp) -> m SubExp) -> m [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([SubExp] -> [SubExp] -> [(SubExp, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
x_dims [SubExp]
y_dims) (((SubExp, SubExp) -> m SubExp) -> m [SubExp])
-> ((SubExp, SubExp) -> m SubExp) -> m [SubExp]
forall a b. (a -> b) -> a -> b
$ \(SubExp
x_dim, SubExp
y_dim) ->
                      String -> Exp (Lore m) -> m SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"dim_eq" (Exp (Lore m) -> m SubExp) -> Exp (Lore m) -> m SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
int32) SubExp
x_dim SubExp
y_dim
                    SubExp
shapes_match <- String -> Exp (Lore m) -> m SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"shapes_match" (Exp (Lore m) -> m SubExp) -> m (Exp (Lore m)) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [SubExp] -> m (Exp (Lore m))
forall (m :: * -> *). MonadBinder m => [SubExp] -> m (Exp (Lore m))
eAll [SubExp]
dims_match
                    Body (Lore m)
compare_elems_body <- Binder (Lore m) (Body (Lore m)) -> m (Body (Lore m))
forall lore (m :: * -> *) somelore.
(Bindable lore, MonadFreshNames m, HasScope somelore m,
 SameScope somelore lore) =>
Binder lore (Body lore) -> m (Body lore)
runBodyBinder (Binder (Lore m) (Body (Lore m)) -> m (Body (Lore m)))
-> Binder (Lore m) (Body (Lore m)) -> m (Body (Lore m))
forall a b. (a -> b) -> a -> b
$ do
                      -- Flatten both x and y.
                      SubExp
x_num_elems <- String
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"x_num_elems" (Exp (Lore m) -> BinderT (Lore m) (State VNameSource) SubExp)
-> BinderT (Lore m) (State VNameSource) (Exp (Lore m))
-> BinderT (Lore m) (State VNameSource) SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<
                                     BinOp
-> SubExp
-> [SubExp]
-> BinderT
     (Lore m)
     (State VNameSource)
     (Exp (Lore (BinderT (Lore m) (State VNameSource))))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> SubExp -> [SubExp] -> m (Exp (Lore m))
foldBinOp (IntType -> Overflow -> BinOp
I.Mul IntType
Int32 Overflow
I.OverflowUndef) (Int32 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int32
1::Int32)) [SubExp]
x_dims
                      VName
x' <- String
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m VName
letExp String
"x" (Exp (Lore (BinderT (Lore m) (State VNameSource)))
 -> BinderT (Lore m) (State VNameSource) VName)
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
x
                      VName
y' <- String
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m VName
letExp String
"x" (Exp (Lore (BinderT (Lore m) (State VNameSource)))
 -> BinderT (Lore m) (State VNameSource) VName)
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
y
                      VName
x_flat <- String
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m VName
letExp String
"x_flat" (Exp (Lore (BinderT (Lore m) (State VNameSource)))
 -> BinderT (Lore m) (State VNameSource) VName)
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
I.Reshape [SubExp -> DimChange SubExp
forall d. d -> DimChange d
I.DimNew SubExp
x_num_elems] VName
x'
                      VName
y_flat <- String
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m VName
letExp String
"y_flat" (Exp (Lore (BinderT (Lore m) (State VNameSource)))
 -> BinderT (Lore m) (State VNameSource) VName)
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
I.Reshape [SubExp -> DimChange SubExp
forall d. d -> DimChange d
I.DimNew SubExp
x_num_elems] VName
y'

                      -- Compare the elements.
                      Lambda (Lore m)
cmp_lam <- CmpOp
-> PrimType
-> BinderT
     (Lore m)
     (State VNameSource)
     (Lambda (Lore (BinderT (Lore m) (State VNameSource))))
forall (m :: * -> *).
(MonadBinder m, Bindable (Lore m)) =>
CmpOp -> PrimType -> m (Lambda (Lore m))
cmpOpLambda (PrimType -> CmpOp
I.CmpEq (Type -> PrimType
forall shape u. TypeBase shape u -> PrimType
elemType Type
x_t)) (Type -> PrimType
forall shape u. TypeBase shape u -> PrimType
elemType Type
x_t)
                      VName
cmps <- String
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m VName
letExp String
"cmps" (Exp (Lore (BinderT (Lore m) (State VNameSource)))
 -> BinderT (Lore m) (State VNameSource) VName)
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$ Op (Lore m) -> Exp (Lore m)
forall lore. Op lore -> ExpT lore
I.Op (Op (Lore m) -> Exp (Lore m)) -> Op (Lore m) -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$
                              SubExp -> ScremaForm (Lore m) -> Closure -> SOAC (Lore m)
forall lore. SubExp -> ScremaForm lore -> Closure -> SOAC lore
I.Screma SubExp
x_num_elems (Lambda (Lore m) -> ScremaForm (Lore m)
forall lore. Lambda lore -> ScremaForm lore
I.mapSOAC Lambda (Lore m)
cmp_lam) [VName
x_flat, VName
y_flat]

                      -- Check that all were equal.
                      Lambda (Lore m)
and_lam <- BinOp
-> PrimType
-> BinderT
     (Lore m)
     (State VNameSource)
     (Lambda (Lore (BinderT (Lore m) (State VNameSource))))
forall (m :: * -> *).
(MonadBinder m, Bindable (Lore m)) =>
BinOp -> PrimType -> m (Lambda (Lore m))
binOpLambda BinOp
I.LogAnd PrimType
I.Bool
                      ScremaForm (Lore m)
reduce <- [Reduce (Lore m)]
-> BinderT (Lore m) (State VNameSource) (ScremaForm (Lore m))
forall lore (m :: * -> *).
(Bindable lore, MonadFreshNames m) =>
[Reduce lore] -> m (ScremaForm lore)
I.reduceSOAC [Commutativity -> Lambda (Lore m) -> [SubExp] -> Reduce (Lore m)
forall lore.
Commutativity -> Lambda lore -> [SubExp] -> Reduce lore
Reduce Commutativity
Commutative Lambda (Lore m)
and_lam [Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
True]]
                      SubExp
all_equal <- String
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"all_equal" (Exp (Lore (BinderT (Lore m) (State VNameSource)))
 -> BinderT (Lore m) (State VNameSource) SubExp)
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) SubExp
forall a b. (a -> b) -> a -> b
$ Op (Lore m) -> Exp (Lore m)
forall lore. Op lore -> ExpT lore
I.Op (Op (Lore m) -> Exp (Lore m)) -> Op (Lore m) -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ SubExp -> ScremaForm (Lore m) -> Closure -> SOAC (Lore m)
forall lore. SubExp -> ScremaForm lore -> Closure -> SOAC lore
I.Screma SubExp
x_num_elems ScremaForm (Lore m)
reduce [VName
cmps]
                      Body (Lore m) -> Binder (Lore m) (Body (Lore m))
forall (m :: * -> *) a. Monad m => a -> m a
return (Body (Lore m) -> Binder (Lore m) (Body (Lore m)))
-> Body (Lore m) -> Binder (Lore m) (Body (Lore m))
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Body (Lore m)
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [SubExp
all_equal]

                    String -> Exp (Lore m) -> m SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"arrays_equal" (Exp (Lore m) -> m SubExp) -> Exp (Lore m) -> m SubExp
forall a b. (a -> b) -> a -> b
$
                      SubExp
-> Body (Lore m)
-> Body (Lore m)
-> IfDec (BranchType (Lore m))
-> Exp (Lore m)
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If SubExp
shapes_match Body (Lore m)
compare_elems_body ([SubExp] -> Body (Lore m)
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
False]) (IfDec (BranchType (Lore m)) -> Exp (Lore m))
-> IfDec (BranchType (Lore m)) -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$
                      [Type] -> IfDec ExtType
ifCommon [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.Bool]

    handleOps [ExpBase Info VName
x,ExpBase Info VName
y] String
name
      | Just BinOp
bop <- (BinOp -> Bool) -> [BinOp] -> Maybe BinOp
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((String
nameString -> String -> Bool
forall a. Eq a => a -> a -> Bool
==) (String -> Bool) -> (BinOp -> String) -> BinOp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BinOp -> String
forall a. Pretty a => a -> String
pretty) [BinOp
forall a. Bounded a => a
minBound..BinOp
forall a. Bounded a => a
maxBound::E.BinOp] =
      (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc -> do
        SubExp
x' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"x" ExpBase Info VName
x
        SubExp
y' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"y" ExpBase Info VName
y
        case (ExpBase Info VName -> PatternType
E.typeOf ExpBase Info VName
x, ExpBase Info VName -> PatternType
E.typeOf ExpBase Info VName
y) of
          (E.Scalar (E.Prim PrimType
t1), E.Scalar (E.Prim PrimType
t2)) ->
            SrcLoc
-> String
-> BinOp
-> SubExp
-> SubExp
-> PrimType
-> PrimType
-> InternaliseM [SubExp]
internaliseBinOp SrcLoc
loc String
desc BinOp
bop SubExp
x' SubExp
y' PrimType
t1 PrimType
t2
          (PatternType, PatternType)
_ -> String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error String
"Futhark.Internalise.internaliseExp: non-primitive type in BinOp."

    handleOps [ExpBase Info VName]
_ String
_ = Maybe (String -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

    handleSOACs :: [ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM [SubExp])
handleSOACs [TupLit [ExpBase Info VName
lam, ExpBase Info VName
arr] SrcLoc
_] String
"map" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc -> do
      Closure
arr' <- String -> ExpBase Info VName -> InternaliseM Closure
internaliseExpToVars String
"map_arr" ExpBase Info VName
arr
      Lambda
lam' <- InternaliseLambda
-> ExpBase Info VName -> [SubExp] -> InternaliseM Lambda
internaliseMapLambda InternaliseLambda
internaliseLambda ExpBase Info VName
lam ([SubExp] -> InternaliseM Lambda)
-> [SubExp] -> InternaliseM Lambda
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> Closure -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var Closure
arr'
      SubExp
w <- Int -> [Type] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 ([Type] -> SubExp) -> InternaliseM [Type] -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM Type) -> Closure -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType Closure
arr'
      String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Op SOACS -> ExpT SOACS
forall lore. Op lore -> ExpT lore
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
        SubExp -> ScremaForm SOACS -> Closure -> SOAC SOACS
forall lore. SubExp -> ScremaForm lore -> Closure -> SOAC lore
I.Screma SubExp
w (Lambda -> ScremaForm SOACS
forall lore. Lambda lore -> ScremaForm lore
I.mapSOAC Lambda
lam') Closure
arr'

    handleSOACs [TupLit [ExpBase Info VName
k, ExpBase Info VName
lam, ExpBase Info VName
arr] SrcLoc
_] String
"partition" = do
      Int
k' <- Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int32 -> Int) -> Maybe Int32 -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> Maybe Int32
forall vn. ExpBase Info vn -> Maybe Int32
isInt32 ExpBase Info VName
k
      (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
_desc -> do
        Closure
arrs <- String -> ExpBase Info VName -> InternaliseM Closure
internaliseExpToVars String
"partition_input" ExpBase Info VName
arr
        Lambda
lam' <- InternaliseLambda
-> Int -> ExpBase Info VName -> [SubExp] -> InternaliseM Lambda
internalisePartitionLambda InternaliseLambda
internaliseLambda Int
k' ExpBase Info VName
lam ([SubExp] -> InternaliseM Lambda)
-> [SubExp] -> InternaliseM Lambda
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> Closure -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var Closure
arrs
        ([SubExp] -> [SubExp] -> [SubExp])
-> ([SubExp], [SubExp]) -> [SubExp]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
(++) (([SubExp], [SubExp]) -> [SubExp])
-> InternaliseM ([SubExp], [SubExp]) -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Lambda -> Closure -> InternaliseM ([SubExp], [SubExp])
partitionWithSOACS Int
k' Lambda
lam' Closure
arrs
        where isInt32 :: ExpBase Info vn -> Maybe Int32
isInt32 (Literal (SignedValue (Int32Value Int32
k')) SrcLoc
_) = Int32 -> Maybe Int32
forall a. a -> Maybe a
Just Int32
k'
              isInt32 (IntLit Integer
k' (Info (E.Scalar (E.Prim (Signed IntType
Int32)))) SrcLoc
_) = Int32 -> Maybe Int32
forall a. a -> Maybe a
Just (Int32 -> Maybe Int32) -> Int32 -> Maybe Int32
forall a b. (a -> b) -> a -> b
$ Integer -> Int32
forall a. Num a => Integer -> a
fromInteger Integer
k'
              isInt32 ExpBase Info vn
_ = Maybe Int32
forall a. Maybe a
Nothing

    handleSOACs [TupLit [ExpBase Info VName
lam, ExpBase Info VName
ne, ExpBase Info VName
arr] SrcLoc
_] String
"reduce" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> String
-> (SubExp
    -> Lambda -> [SubExp] -> Closure -> InternaliseM (SOAC SOACS))
-> (ExpBase Info VName, ExpBase Info VName, ExpBase Info VName,
    SrcLoc)
-> InternaliseM [SubExp]
internaliseScanOrReduce String
desc String
"reduce" SubExp
-> Lambda -> [SubExp] -> Closure -> InternaliseM (SOAC SOACS)
forall (f :: * -> *) lore.
(Bindable lore, MonadFreshNames f) =>
SubExp -> Lambda lore -> [SubExp] -> Closure -> f (SOAC lore)
reduce (ExpBase Info VName
lam, ExpBase Info VName
ne, ExpBase Info VName
arr, SrcLoc
loc)
      where reduce :: SubExp -> Lambda lore -> [SubExp] -> Closure -> f (SOAC lore)
reduce SubExp
w Lambda lore
red_lam [SubExp]
nes Closure
arrs =
              SubExp -> ScremaForm lore -> Closure -> SOAC lore
forall lore. SubExp -> ScremaForm lore -> Closure -> SOAC lore
I.Screma SubExp
w (ScremaForm lore -> Closure -> SOAC lore)
-> f (ScremaForm lore) -> f (Closure -> SOAC lore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
              [Reduce lore] -> f (ScremaForm lore)
forall lore (m :: * -> *).
(Bindable lore, MonadFreshNames m) =>
[Reduce lore] -> m (ScremaForm lore)
I.reduceSOAC [Commutativity -> Lambda lore -> [SubExp] -> Reduce lore
forall lore.
Commutativity -> Lambda lore -> [SubExp] -> Reduce lore
Reduce Commutativity
Noncommutative Lambda lore
red_lam [SubExp]
nes] f (Closure -> SOAC lore) -> f Closure -> f (SOAC lore)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Closure -> f Closure
forall (f :: * -> *) a. Applicative f => a -> f a
pure Closure
arrs

    handleSOACs [TupLit [ExpBase Info VName
lam, ExpBase Info VName
ne, ExpBase Info VName
arr] SrcLoc
_] String
"reduce_comm" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> String
-> (SubExp
    -> Lambda -> [SubExp] -> Closure -> InternaliseM (SOAC SOACS))
-> (ExpBase Info VName, ExpBase Info VName, ExpBase Info VName,
    SrcLoc)
-> InternaliseM [SubExp]
internaliseScanOrReduce String
desc String
"reduce" SubExp
-> Lambda -> [SubExp] -> Closure -> InternaliseM (SOAC SOACS)
forall (f :: * -> *) lore.
(Bindable lore, MonadFreshNames f) =>
SubExp -> Lambda lore -> [SubExp] -> Closure -> f (SOAC lore)
reduce (ExpBase Info VName
lam, ExpBase Info VName
ne, ExpBase Info VName
arr, SrcLoc
loc)
      where reduce :: SubExp -> Lambda lore -> [SubExp] -> Closure -> f (SOAC lore)
reduce SubExp
w Lambda lore
red_lam [SubExp]
nes Closure
arrs =
              SubExp -> ScremaForm lore -> Closure -> SOAC lore
forall lore. SubExp -> ScremaForm lore -> Closure -> SOAC lore
I.Screma SubExp
w (ScremaForm lore -> Closure -> SOAC lore)
-> f (ScremaForm lore) -> f (Closure -> SOAC lore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
              [Reduce lore] -> f (ScremaForm lore)
forall lore (m :: * -> *).
(Bindable lore, MonadFreshNames m) =>
[Reduce lore] -> m (ScremaForm lore)
I.reduceSOAC [Commutativity -> Lambda lore -> [SubExp] -> Reduce lore
forall lore.
Commutativity -> Lambda lore -> [SubExp] -> Reduce lore
Reduce Commutativity
Commutative Lambda lore
red_lam [SubExp]
nes] f (Closure -> SOAC lore) -> f Closure -> f (SOAC lore)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Closure -> f Closure
forall (f :: * -> *) a. Applicative f => a -> f a
pure Closure
arrs

    handleSOACs [TupLit [ExpBase Info VName
lam, ExpBase Info VName
ne, ExpBase Info VName
arr] SrcLoc
_] String
"scan" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> String
-> (SubExp
    -> Lambda -> [SubExp] -> Closure -> InternaliseM (SOAC SOACS))
-> (ExpBase Info VName, ExpBase Info VName, ExpBase Info VName,
    SrcLoc)
-> InternaliseM [SubExp]
internaliseScanOrReduce String
desc String
"scan" SubExp
-> Lambda -> [SubExp] -> Closure -> InternaliseM (SOAC SOACS)
forall (f :: * -> *) lore.
(Bindable lore, MonadFreshNames f) =>
SubExp -> Lambda lore -> [SubExp] -> Closure -> f (SOAC lore)
reduce (ExpBase Info VName
lam, ExpBase Info VName
ne, ExpBase Info VName
arr, SrcLoc
loc)
      where reduce :: SubExp -> Lambda lore -> [SubExp] -> Closure -> f (SOAC lore)
reduce SubExp
w Lambda lore
scan_lam [SubExp]
nes Closure
arrs =
              SubExp -> ScremaForm lore -> Closure -> SOAC lore
forall lore. SubExp -> ScremaForm lore -> Closure -> SOAC lore
I.Screma SubExp
w (ScremaForm lore -> Closure -> SOAC lore)
-> f (ScremaForm lore) -> f (Closure -> SOAC lore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Scan lore] -> f (ScremaForm lore)
forall lore (m :: * -> *).
(Bindable lore, MonadFreshNames m) =>
[Scan lore] -> m (ScremaForm lore)
I.scanSOAC [Lambda lore -> [SubExp] -> Scan lore
forall lore. Lambda lore -> [SubExp] -> Scan lore
Scan Lambda lore
scan_lam [SubExp]
nes] f (Closure -> SOAC lore) -> f Closure -> f (SOAC lore)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Closure -> f Closure
forall (f :: * -> *) a. Applicative f => a -> f a
pure Closure
arrs

    handleSOACs [TupLit [ExpBase Info VName
op, ExpBase Info VName
f, ExpBase Info VName
arr] SrcLoc
_] String
"reduce_stream" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> StreamOrd
-> Commutativity
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> InternaliseM [SubExp]
internaliseStreamRed String
desc StreamOrd
InOrder Commutativity
Noncommutative ExpBase Info VName
op ExpBase Info VName
f ExpBase Info VName
arr

    handleSOACs [TupLit [ExpBase Info VName
op, ExpBase Info VName
f, ExpBase Info VName
arr] SrcLoc
_] String
"reduce_stream_per" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> StreamOrd
-> Commutativity
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> InternaliseM [SubExp]
internaliseStreamRed String
desc StreamOrd
Disorder Commutativity
Commutative ExpBase Info VName
op ExpBase Info VName
f ExpBase Info VName
arr

    handleSOACs [TupLit [ExpBase Info VName
f, ExpBase Info VName
arr] SrcLoc
_] String
"map_stream" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> StreamOrd
-> ExpBase Info VName
-> ExpBase Info VName
-> InternaliseM [SubExp]
internaliseStreamMap String
desc StreamOrd
InOrder ExpBase Info VName
f ExpBase Info VName
arr

    handleSOACs [TupLit [ExpBase Info VName
f, ExpBase Info VName
arr] SrcLoc
_] String
"map_stream_per" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> StreamOrd
-> ExpBase Info VName
-> ExpBase Info VName
-> InternaliseM [SubExp]
internaliseStreamMap String
desc StreamOrd
Disorder ExpBase Info VName
f ExpBase Info VName
arr

    handleSOACs [TupLit [ExpBase Info VName
rf, ExpBase Info VName
dest, ExpBase Info VName
op, ExpBase Info VName
ne, ExpBase Info VName
buckets, ExpBase Info VName
img] SrcLoc
_] String
"hist" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> SrcLoc
-> InternaliseM [SubExp]
internaliseHist String
desc ExpBase Info VName
rf ExpBase Info VName
dest ExpBase Info VName
op ExpBase Info VName
ne ExpBase Info VName
buckets ExpBase Info VName
img SrcLoc
loc

    handleSOACs [ExpBase Info VName]
_ String
_ = Maybe (String -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

    handleRest :: [ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM [SubExp])
handleRest [ExpBase Info VName
x] String
"!" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> String -> InternaliseM [SubExp]
complementF ExpBase Info VName
x

    handleRest [ExpBase Info VName
x] String
"opaque" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      (SubExp -> InternaliseM SubExp)
-> [SubExp] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
desc (ExpT SOACS -> InternaliseM SubExp)
-> (SubExp -> ExpT SOACS) -> SubExp -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS)
-> (SubExp -> BasicOp) -> SubExp -> ExpT SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
Opaque) ([SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"opaque_arg" ExpBase Info VName
x

    handleRest [E.TupLit [ExpBase Info VName
a, ExpBase Info VName
si, ExpBase Info VName
v] SrcLoc
_] String
"scatter" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> String
-> InternaliseM [SubExp]
scatterF ExpBase Info VName
a ExpBase Info VName
si ExpBase Info VName
v

    handleRest [E.TupLit [ExpBase Info VName
n, ExpBase Info VName
m, ExpBase Info VName
arr] SrcLoc
_] String
"unflatten" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc -> do
      Closure
arrs <- String -> ExpBase Info VName -> InternaliseM Closure
internaliseExpToVars String
"unflatten_arr" ExpBase Info VName
arr
      SubExp
n' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"n" ExpBase Info VName
n
      SubExp
m' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"m" ExpBase Info VName
m
      -- The unflattened dimension needs to have the same number of elements
      -- as the original dimension.
      SubExp
old_dim <- Int -> [Type] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
I.arraysSize Int
0 ([Type] -> SubExp) -> InternaliseM [Type] -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM Type) -> Closure -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType Closure
arrs
      SubExp
dim_ok <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"dim_ok" (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<
                CmpOp
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
CmpOp -> m (Exp (Lore m)) -> m (Exp (Lore m)) -> m (Exp (Lore m))
eCmpOp (PrimType -> CmpOp
I.CmpEq PrimType
I.int32)
                (BinOp
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> m (Exp (Lore m)) -> m (Exp (Lore m)) -> m (Exp (Lore m))
eBinOp (IntType -> Overflow -> BinOp
I.Mul IntType
Int32 Overflow
I.OverflowUndef) (SubExp -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => SubExp -> m (Exp (Lore m))
eSubExp SubExp
n') (SubExp -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => SubExp -> m (Exp (Lore m))
eSubExp SubExp
m'))
                (SubExp -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => SubExp -> m (Exp (Lore m))
eSubExp SubExp
old_dim)
      Certificates
dim_ok_cert <- String
-> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certificates
assert String
"dim_ok_cert" SubExp
dim_ok
                     ErrorMsg SubExp
"new shape has different number of elements than old shape" SrcLoc
loc
      Certificates -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
dim_ok_cert (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Closure -> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM Closure
arrs ((VName -> InternaliseM SubExp) -> InternaliseM [SubExp])
-> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
arr' -> do
        Type
arr_t <- VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
arr'
        String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
desc (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
          ShapeChange SubExp -> VName -> BasicOp
I.Reshape (ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp
reshapeOuter [SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimNew SubExp
n', SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimNew SubExp
m'] Int
1 (Shape -> ShapeChange SubExp) -> Shape -> ShapeChange SubExp
forall a b. (a -> b) -> a -> b
$ Type -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape Type
arr_t) VName
arr'

    handleRest [ExpBase Info VName
arr] String
"flatten" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc -> do
      Closure
arrs <- String -> ExpBase Info VName -> InternaliseM Closure
internaliseExpToVars String
"flatten_arr" ExpBase Info VName
arr
      Closure -> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM Closure
arrs ((VName -> InternaliseM SubExp) -> InternaliseM [SubExp])
-> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
arr' -> do
        Type
arr_t <- VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
arr'
        let n :: SubExp
n = Int -> Type -> SubExp
forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
0 Type
arr_t
            m :: SubExp
m = Int -> Type -> SubExp
forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
1 Type
arr_t
        SubExp
k <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"flat_dim" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Mul IntType
Int32 Overflow
I.OverflowUndef) SubExp
n SubExp
m
        String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
desc (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
          ShapeChange SubExp -> VName -> BasicOp
I.Reshape (ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp
reshapeOuter [SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimNew SubExp
k] Int
2 (Shape -> ShapeChange SubExp) -> Shape -> ShapeChange SubExp
forall a b. (a -> b) -> a -> b
$ Type -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape Type
arr_t) VName
arr'

    handleRest [TupLit [ExpBase Info VName
x, ExpBase Info VName
y] SrcLoc
_] String
"concat" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc -> do
      Closure
xs <- String -> ExpBase Info VName -> InternaliseM Closure
internaliseExpToVars String
"concat_x" ExpBase Info VName
x
      Closure
ys <- String -> ExpBase Info VName -> InternaliseM Closure
internaliseExpToVars String
"concat_y" ExpBase Info VName
y
      SubExp
outer_size <- Int -> [Type] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 ([Type] -> SubExp) -> InternaliseM [Type] -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM Type) -> Closure -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType Closure
xs
      let sumdims :: SubExp -> SubExp -> m SubExp
sumdims SubExp
xsize SubExp
ysize = String -> Exp (Lore m) -> m SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"conc_tmp" (Exp (Lore m) -> m SubExp) -> Exp (Lore m) -> m SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$
                                BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Add IntType
I.Int32 Overflow
I.OverflowUndef) SubExp
xsize SubExp
ysize
      SubExp
ressize <- (SubExp -> SubExp -> InternaliseM SubExp)
-> SubExp -> [SubExp] -> InternaliseM SubExp
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM SubExp -> SubExp -> InternaliseM SubExp
forall (m :: * -> *). MonadBinder m => SubExp -> SubExp -> m SubExp
sumdims SubExp
outer_size ([SubExp] -> InternaliseM SubExp)
-> InternaliseM [SubExp] -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<
                 (Closure -> InternaliseM SubExp)
-> [Closure] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (([Type] -> SubExp) -> InternaliseM [Type] -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> [Type] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0) (InternaliseM [Type] -> InternaliseM SubExp)
-> (Closure -> InternaliseM [Type])
-> Closure
-> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> InternaliseM Type) -> Closure -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType) [Closure
ys]

      let conc :: VName -> VName -> ExpT SOACS
conc VName
xarr VName
yarr =
            BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ Int -> VName -> Closure -> SubExp -> BasicOp
I.Concat Int
0 VName
xarr [VName
yarr] SubExp
ressize
      String -> [Exp (Lore InternaliseM)] -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> [Exp (Lore m)] -> m [SubExp]
letSubExps String
desc ([Exp (Lore InternaliseM)] -> InternaliseM [SubExp])
-> [Exp (Lore InternaliseM)] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ (VName -> VName -> ExpT SOACS)
-> Closure -> Closure -> [ExpT SOACS]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith VName -> VName -> ExpT SOACS
conc Closure
xs Closure
ys

    handleRest [TupLit [ExpBase Info VName
offset, ExpBase Info VName
e] SrcLoc
_] String
"rotate" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc -> do
      SubExp
offset' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"rotation_offset" ExpBase Info VName
offset
      String
-> ExpBase Info VName
-> (VName -> InternaliseM BasicOp)
-> InternaliseM [SubExp]
internaliseOperation String
desc ExpBase Info VName
e ((VName -> InternaliseM BasicOp) -> InternaliseM [SubExp])
-> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
v -> do
        Int
r <- Type -> Int
forall shape u. ArrayShape shape => TypeBase shape u -> Int
I.arrayRank (Type -> Int) -> InternaliseM Type -> InternaliseM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
v
        let zero :: SubExp
zero = IntType -> Integer -> SubExp
intConst IntType
Int32 Integer
0
            offsets :: [SubExp]
offsets = SubExp
offset' SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: Int -> SubExp -> [SubExp]
forall a. Int -> a -> [a]
replicate (Int
rInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) SubExp
zero
        BasicOp -> InternaliseM BasicOp
forall (m :: * -> *) a. Monad m => a -> m a
return (BasicOp -> InternaliseM BasicOp)
-> BasicOp -> InternaliseM BasicOp
forall a b. (a -> b) -> a -> b
$ [SubExp] -> VName -> BasicOp
I.Rotate [SubExp]
offsets VName
v

    handleRest [ExpBase Info VName
e] String
"transpose" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> ExpBase Info VName
-> (VName -> InternaliseM BasicOp)
-> InternaliseM [SubExp]
internaliseOperation String
desc ExpBase Info VName
e ((VName -> InternaliseM BasicOp) -> InternaliseM [SubExp])
-> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
v -> do
        Int
r <- Type -> Int
forall shape u. ArrayShape shape => TypeBase shape u -> Int
I.arrayRank (Type -> Int) -> InternaliseM Type -> InternaliseM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
v
        BasicOp -> InternaliseM BasicOp
forall (m :: * -> *) a. Monad m => a -> m a
return (BasicOp -> InternaliseM BasicOp)
-> BasicOp -> InternaliseM BasicOp
forall a b. (a -> b) -> a -> b
$ [Int] -> VName -> BasicOp
I.Rearrange ([Int
1,Int
0] [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
2..Int
rInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1]) VName
v

    handleRest [TupLit [ExpBase Info VName
x, ExpBase Info VName
y] SrcLoc
_] String
"zip" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
(++) ([SubExp] -> [SubExp] -> [SubExp])
-> InternaliseM [SubExp] -> InternaliseM ([SubExp] -> [SubExp])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp (String
desc String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_zip_x") ExpBase Info VName
x
           InternaliseM ([SubExp] -> [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp (String
desc String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_zip_y") ExpBase Info VName
y

    handleRest [ExpBase Info VName
x] String
"unzip" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ (String -> ExpBase Info VName -> InternaliseM [SubExp])
-> ExpBase Info VName -> String -> InternaliseM [SubExp]
forall a b c. (a -> b -> c) -> b -> a -> c
flip String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp ExpBase Info VName
x
    handleRest [ExpBase Info VName
x] String
"trace" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ (String -> ExpBase Info VName -> InternaliseM [SubExp])
-> ExpBase Info VName -> String -> InternaliseM [SubExp]
forall a b c. (a -> b -> c) -> b -> a -> c
flip String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp ExpBase Info VName
x
    handleRest [ExpBase Info VName
x] String
"break" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ (String -> ExpBase Info VName -> InternaliseM [SubExp])
-> ExpBase Info VName -> String -> InternaliseM [SubExp]
forall a b c. (a -> b -> c) -> b -> a -> c
flip String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp ExpBase Info VName
x

    handleRest [ExpBase Info VName]
_ String
_ = Maybe (String -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

    toSigned :: IntType -> ExpBase Info VName -> String -> InternaliseM [SubExp]
toSigned IntType
int_to ExpBase Info VName
e String
desc = do
      SubExp
e' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"trunc_arg" ExpBase Info VName
e
      case ExpBase Info VName -> PatternType
E.typeOf ExpBase Info VName
e of
        E.Scalar (E.Prim PrimType
E.Bool) ->
          String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ SubExp -> Body -> Body -> IfDec (BranchType SOACS) -> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If SubExp
e' ([SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
1])
                                    ([SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
0]) (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
                                    [Type] -> IfDec ExtType
ifCommon [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim (PrimType -> Type) -> PrimType -> Type
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
I.IntType IntType
int_to]
        E.Scalar (E.Prim (E.Signed IntType
int_from)) ->
          String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (IntType -> IntType -> ConvOp
I.SExt IntType
int_from IntType
int_to) SubExp
e'
        E.Scalar (E.Prim (E.Unsigned IntType
int_from)) ->
          String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (IntType -> IntType -> ConvOp
I.ZExt IntType
int_from IntType
int_to) SubExp
e'
        E.Scalar (E.Prim (E.FloatType FloatType
float_from)) ->
          String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (FloatType -> IntType -> ConvOp
I.FPToSI FloatType
float_from IntType
int_to) SubExp
e'
        PatternType
_ -> String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error String
"Futhark.Internalise: non-numeric type in ToSigned"

    toUnsigned :: IntType -> ExpBase Info VName -> String -> InternaliseM [SubExp]
toUnsigned IntType
int_to ExpBase Info VName
e String
desc = do
      SubExp
e' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"trunc_arg" ExpBase Info VName
e
      case ExpBase Info VName -> PatternType
E.typeOf ExpBase Info VName
e of
        E.Scalar (E.Prim PrimType
E.Bool) ->
          String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ SubExp -> Body -> Body -> IfDec (BranchType SOACS) -> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If SubExp
e' ([SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
1])
                                    ([SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
0]) (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
                                    [Type] -> IfDec ExtType
ifCommon [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim (PrimType -> Type) -> PrimType -> Type
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
I.IntType IntType
int_to]
        E.Scalar (E.Prim (E.Signed IntType
int_from)) ->
          String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (IntType -> IntType -> ConvOp
I.ZExt IntType
int_from IntType
int_to) SubExp
e'
        E.Scalar (E.Prim (E.Unsigned IntType
int_from)) ->
          String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (IntType -> IntType -> ConvOp
I.ZExt IntType
int_from IntType
int_to) SubExp
e'
        E.Scalar (E.Prim (E.FloatType FloatType
float_from)) ->
          String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (FloatType -> IntType -> ConvOp
I.FPToUI FloatType
float_from IntType
int_to) SubExp
e'
        PatternType
_ -> String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error String
"Futhark.Internalise.internaliseExp: non-numeric type in ToUnsigned"

    complementF :: ExpBase Info VName -> String -> InternaliseM [SubExp]
complementF ExpBase Info VName
e String
desc = do
      SubExp
e' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"complement_arg" ExpBase Info VName
e
      Type
et <- SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
e'
      case Type
et of I.Prim (I.IntType IntType
t) ->
                   String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (IntType -> UnOp
I.Complement IntType
t) SubExp
e'
                 I.Prim PrimType
I.Bool ->
                   String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
I.Not SubExp
e'
                 Type
_ ->
                   String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error String
"Futhark.Internalise.internaliseExp: non-int/bool type in Complement"

    scatterF :: ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> String
-> InternaliseM [SubExp]
scatterF ExpBase Info VName
a ExpBase Info VName
si ExpBase Info VName
v String
desc = do
      VName
si' <- String -> Exp (Lore InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m VName
letExp String
"write_si" (ExpT SOACS -> InternaliseM VName)
-> (SubExp -> ExpT SOACS) -> SubExp -> InternaliseM VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS)
-> (SubExp -> BasicOp) -> SubExp -> ExpT SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
SubExp (SubExp -> InternaliseM VName)
-> InternaliseM SubExp -> InternaliseM VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"write_arg_i" ExpBase Info VName
si
      Closure
svs <- String -> ExpBase Info VName -> InternaliseM Closure
internaliseExpToVars String
"write_arg_v" ExpBase Info VName
v
      Closure
sas <- String -> ExpBase Info VName -> InternaliseM Closure
internaliseExpToVars String
"write_arg_a" ExpBase Info VName
a

      Shape
si_shape <- Type -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape (Type -> Shape) -> InternaliseM Type -> InternaliseM Shape
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
si'
      let si_w :: SubExp
si_w = Int -> Shape -> SubExp
shapeSize Int
0 Shape
si_shape
      [Type]
sv_ts <- (VName -> InternaliseM Type) -> Closure -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType Closure
svs

      Closure
svs' <- [(VName, Type)]
-> ((VName, Type) -> InternaliseM VName) -> InternaliseM Closure
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (Closure -> [Type] -> [(VName, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip Closure
svs [Type]
sv_ts) (((VName, Type) -> InternaliseM VName) -> InternaliseM Closure)
-> ((VName, Type) -> InternaliseM VName) -> InternaliseM Closure
forall a b. (a -> b) -> a -> b
$ \(VName
sv,Type
sv_t) -> do
        let sv_shape :: Shape
sv_shape = Type -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape Type
sv_t
            sv_w :: SubExp
sv_w = Int -> Type -> SubExp
forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
0 Type
sv_t

        -- Generate an assertion and reshapes to ensure that sv and si' are the same
        -- size.
        SubExp
cmp <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"write_cmp" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
          CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
I.int32) SubExp
si_w SubExp
sv_w
        Certificates
c <- String
-> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certificates
assert String
"write_cert" SubExp
cmp
             ErrorMsg SubExp
"length of index and value array does not match" SrcLoc
loc
        Certificates -> InternaliseM VName -> InternaliseM VName
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
c (InternaliseM VName -> InternaliseM VName)
-> InternaliseM VName -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$ String -> Exp (Lore InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m VName
letExp (VName -> String
baseString VName
sv String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_write_sv") (Exp (Lore InternaliseM) -> InternaliseM VName)
-> Exp (Lore InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$
          BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
I.Reshape (ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp
reshapeOuter [SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimCoercion SubExp
si_w] Int
1 Shape
sv_shape) VName
sv

      Type
indexType <- Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
rowType (Type -> Type) -> InternaliseM Type -> InternaliseM Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
si'
      VName
indexName <- String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"write_index"
      Closure
valueNames <- Int -> InternaliseM VName -> InternaliseM Closure
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM ([Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
sv_ts) (InternaliseM VName -> InternaliseM Closure)
-> InternaliseM VName -> InternaliseM Closure
forall a b. (a -> b) -> a -> b
$ String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"write_value"

      [Type]
sa_ts <- (VName -> InternaliseM Type) -> Closure -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType Closure
sas
      let bodyTypes :: [Type]
bodyTypes = Int -> Type -> [Type]
forall a. Int -> a -> [a]
replicate ([Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
sv_ts) Type
indexType [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
rowType [Type]
sa_ts
          paramTypes :: [Type]
paramTypes = Type
indexType Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
: (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
rowType [Type]
sv_ts
          bodyNames :: Closure
bodyNames = VName
indexName VName -> Closure -> Closure
forall a. a -> [a] -> [a]
: Closure
valueNames
          bodyParams :: [Param Type]
bodyParams = (VName -> Type -> Param Type) -> Closure -> [Type] -> [Param Type]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith VName -> Type -> Param Type
forall dec. VName -> dec -> Param dec
I.Param Closure
bodyNames [Type]
paramTypes

      -- This body is pretty boring right now, as every input is exactly the output.
      -- But it can get funky later on if fused with something else.
      Body
body <- Scope SOACS -> InternaliseM Body -> InternaliseM Body
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param Type] -> Scope SOACS
forall lore dec.
(LParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
scopeOfLParams [Param Type]
bodyParams) (InternaliseM Body -> InternaliseM Body)
-> InternaliseM Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$ InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
m (Body (Lore m)) -> m (Body (Lore m))
insertStmsM (InternaliseM (Body (Lore InternaliseM))
 -> InternaliseM (Body (Lore InternaliseM)))
-> InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ do
        let outs :: Closure
outs = Int -> VName -> Closure
forall a. Int -> a -> [a]
replicate (Closure -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Closure
valueNames) VName
indexName Closure -> Closure -> Closure
forall a. [a] -> [a] -> [a]
++ Closure
valueNames
        [SubExp]
results <- Closure -> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM Closure
outs ((VName -> InternaliseM SubExp) -> InternaliseM [SubExp])
-> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
name ->
          String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"write_res" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
name
        ErrorMsg SubExp -> SrcLoc -> [Type] -> Body -> InternaliseM Body
ensureResultShape ErrorMsg SubExp
"scatter value has wrong size" SrcLoc
loc
          [Type]
bodyTypes (Body -> InternaliseM Body) -> Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [SubExp]
results

      let lam :: Lambda
lam = Lambda :: forall lore. [LParam lore] -> BodyT lore -> [Type] -> LambdaT lore
I.Lambda { lambdaParams :: [LParam]
I.lambdaParams = [Param Type]
[LParam]
bodyParams
                         , lambdaReturnType :: [Type]
I.lambdaReturnType = [Type]
bodyTypes
                         , lambdaBody :: Body
I.lambdaBody = Body
body
                         }
          sivs :: Closure
sivs = VName
si' VName -> Closure -> Closure
forall a. a -> [a] -> [a]
: Closure
svs'

      let sa_ws :: [SubExp]
sa_ws = (Type -> SubExp) -> [Type] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Type -> SubExp
forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
0) [Type]
sa_ts
      String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Op SOACS -> ExpT SOACS
forall lore. Op lore -> ExpT lore
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> Lambda -> Closure -> [(SubExp, Int, VName)] -> SOAC SOACS
forall lore.
SubExp
-> Lambda lore -> Closure -> [(SubExp, Int, VName)] -> SOAC lore
I.Scatter SubExp
si_w Lambda
lam Closure
sivs ([(SubExp, Int, VName)] -> SOAC SOACS)
-> [(SubExp, Int, VName)] -> SOAC SOACS
forall a b. (a -> b) -> a -> b
$ [SubExp] -> [Int] -> Closure -> [(SubExp, Int, VName)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [SubExp]
sa_ws (Int -> [Int]
forall a. a -> [a]
repeat Int
1) Closure
sas

funcall :: String -> QualName VName -> [SubExp] -> SrcLoc
        -> InternaliseM ([SubExp], [I.ExtType])
funcall :: String
-> QualName VName
-> [SubExp]
-> SrcLoc
-> InternaliseM ([SubExp], [ExtType])
funcall String
desc (QualName Closure
_ VName
fname) [SubExp]
args SrcLoc
loc = do
  (Name
fname', Closure
closure, Closure
shapes, [DeclType]
value_paramts, [Param DeclType]
fun_params, [(SubExp, Type)] -> Maybe [DeclExtType]
rettype_fun) <-
    VName -> InternaliseM FunInfo
lookupFunction VName
fname
  [Type]
argts <- (SubExp -> InternaliseM Type) -> [SubExp] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType [SubExp]
args
  [Type]
closure_ts <- (VName -> InternaliseM Type) -> Closure -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType Closure
closure
  let shapeargs :: [SubExp]
shapeargs = Closure -> [DeclType] -> [Type] -> [SubExp]
forall u0 u1.
Closure -> [TypeBase Shape u0] -> [TypeBase Shape u1] -> [SubExp]
argShapes Closure
shapes [DeclType]
value_paramts [Type]
argts
      diets :: [Diet]
diets = Int -> Diet -> [Diet]
forall a. Int -> a -> [a]
replicate (Closure -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Closure
closure Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
shapeargs) Diet
I.ObservePrim [Diet] -> [Diet] -> [Diet]
forall a. [a] -> [a] -> [a]
++
              (DeclType -> Diet) -> [DeclType] -> [Diet]
forall a b. (a -> b) -> [a] -> [b]
map DeclType -> Diet
forall shape. TypeBase shape Uniqueness -> Diet
I.diet [DeclType]
value_paramts
      constOrShape :: b -> TypeBase shape u
constOrShape = TypeBase shape u -> b -> TypeBase shape u
forall a b. a -> b -> a
const (TypeBase shape u -> b -> TypeBase shape u)
-> TypeBase shape u -> b -> TypeBase shape u
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase shape u
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int32
      paramts :: [Type]
paramts = [Type]
closure_ts [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++
                (SubExp -> Type) -> [SubExp] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> Type
forall b shape u. b -> TypeBase shape u
constOrShape [SubExp]
shapeargs [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ (DeclType -> Type) -> [DeclType] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map DeclType -> Type
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl [DeclType]
value_paramts
  [SubExp]
args' <- ErrorMsg SubExp
-> SrcLoc -> Closure -> [Type] -> [SubExp] -> InternaliseM [SubExp]
forall u.
Typed (TypeBase Shape u) =>
ErrorMsg SubExp
-> SrcLoc
-> Closure
-> [TypeBase Shape u]
-> [SubExp]
-> InternaliseM [SubExp]
ensureArgShapes ErrorMsg SubExp
"function arguments of wrong shape"
           SrcLoc
loc ((Param DeclType -> VName) -> [Param DeclType] -> Closure
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
fun_params)
           [Type]
paramts ((VName -> SubExp) -> Closure -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var Closure
closure [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
shapeargs [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
args)
  [Type]
argts' <- (SubExp -> InternaliseM Type) -> [SubExp] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType [SubExp]
args'
  case [(SubExp, Type)] -> Maybe [DeclExtType]
rettype_fun ([(SubExp, Type)] -> Maybe [DeclExtType])
-> [(SubExp, Type)] -> Maybe [DeclExtType]
forall a b. (a -> b) -> a -> b
$ [SubExp] -> [Type] -> [(SubExp, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
args' [Type]
argts' of
    Maybe [DeclExtType]
Nothing -> String -> InternaliseM ([SubExp], [ExtType])
forall a. HasCallStack => String -> a
error (String -> InternaliseM ([SubExp], [ExtType]))
-> String -> InternaliseM ([SubExp], [ExtType])
forall a b. (a -> b) -> a -> b
$ String
"Cannot apply " String -> String -> String
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Pretty a => a -> String
pretty VName
fname String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" to arguments\n " String -> String -> String
forall a. [a] -> [a] -> [a]
++
               [SubExp] -> String
forall a. Pretty a => a -> String
pretty [SubExp]
args' String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"\nof types\n " String -> String -> String
forall a. [a] -> [a] -> [a]
++
               [Type] -> String
forall a. Pretty a => a -> String
pretty [Type]
argts' String -> String -> String
forall a. [a] -> [a] -> [a]
++
               String
"\nFunction has parameters\n " String -> String -> String
forall a. [a] -> [a] -> [a]
++ [Param DeclType] -> String
forall a. Pretty a => a -> String
pretty [Param DeclType]
fun_params
    Just [DeclExtType]
ts -> do
      Safety
safety <- InternaliseM Safety
askSafety
      Attrs
attrs <- (InternaliseEnv -> Attrs) -> InternaliseM Attrs
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks InternaliseEnv -> Attrs
envAttrs
      [SubExp]
ses <- Attrs -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. MonadBinder m => Attrs -> m a -> m a
attributing Attrs
attrs (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [SubExp]
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
             Name
-> [(SubExp, Diet)]
-> [RetType SOACS]
-> (Safety, SrcLoc, [SrcLoc])
-> ExpT SOACS
forall lore.
Name
-> [(SubExp, Diet)]
-> [RetType lore]
-> (Safety, SrcLoc, [SrcLoc])
-> ExpT lore
I.Apply Name
fname' ([SubExp] -> [Diet] -> [(SubExp, Diet)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
args' [Diet]
diets) [DeclExtType]
[RetType SOACS]
ts (Safety
safety, SrcLoc
loc, [SrcLoc]
forall a. Monoid a => a
mempty)
      ([SubExp], [ExtType]) -> InternaliseM ([SubExp], [ExtType])
forall (m :: * -> *) a. Monad m => a -> m a
return ([SubExp]
ses, (DeclExtType -> ExtType) -> [DeclExtType] -> [ExtType]
forall a b. (a -> b) -> [a] -> [b]
map DeclExtType -> ExtType
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl [DeclExtType]
ts)

-- Bind existential names defined by an expression, based on the
-- concrete values that expression evaluated to.  This most
-- importantly should be done after function calls, but also
-- everything else that can produce existentials in the source
-- language.
bindExtSizes :: E.StructType -> [VName] -> [SubExp] -> InternaliseM ()
bindExtSizes :: StructType -> Closure -> [SubExp] -> InternaliseM ()
bindExtSizes StructType
ret Closure
retext [SubExp]
ses = do
  [DeclExtType]
ts <- [[DeclExtType]] -> [DeclExtType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[DeclExtType]] -> [DeclExtType])
-> InternaliseM [[DeclExtType]] -> InternaliseM [DeclExtType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
        BoundInTypes
-> Map VName VName -> [StructType] -> InternaliseM [[DeclExtType]]
internaliseParamTypes BoundInTypes
forall a. Monoid a => a
mempty ([(VName, VName)] -> Map VName VName
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, VName)] -> Map VName VName)
-> [(VName, VName)] -> Map VName VName
forall a b. (a -> b) -> a -> b
$ Closure -> Closure -> [(VName, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip Closure
retext Closure
retext) [StructType
ret]
  [Type]
ses_ts <- (SubExp -> InternaliseM Type) -> [SubExp] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType [SubExp]
ses

  let combine :: DeclExtType -> Type -> Map VName SubExp
combine DeclExtType
t1 Type
t2 =
        [Map VName SubExp] -> Map VName SubExp
forall a. Monoid a => [a] -> a
mconcat ([Map VName SubExp] -> Map VName SubExp)
-> [Map VName SubExp] -> Map VName SubExp
forall a b. (a -> b) -> a -> b
$ (Ext SubExp -> SubExp -> Map VName SubExp)
-> [Ext SubExp] -> [SubExp] -> [Map VName SubExp]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Ext SubExp -> SubExp -> Map VName SubExp
combine' (DeclExtType -> [Ext SubExp]
forall u. TypeBase ExtShape u -> [Ext SubExp]
arrayExtDims DeclExtType
t1) (Type -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
arrayDims Type
t2)
      combine' :: Ext SubExp -> SubExp -> Map VName SubExp
combine' (I.Free (I.Var VName
v)) SubExp
se
        | VName
v VName -> Closure -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` Closure
retext = VName -> SubExp -> Map VName SubExp
forall k a. k -> a -> Map k a
M.singleton VName
v SubExp
se
      combine' Ext SubExp
_ SubExp
_ = Map VName SubExp
forall a. Monoid a => a
mempty

  [(VName, SubExp)]
-> ((VName, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Map VName SubExp -> [(VName, SubExp)]
forall k a. Map k a -> [(k, a)]
M.toList (Map VName SubExp -> [(VName, SubExp)])
-> Map VName SubExp -> [(VName, SubExp)]
forall a b. (a -> b) -> a -> b
$ [Map VName SubExp] -> Map VName SubExp
forall a. Monoid a => [a] -> a
mconcat ([Map VName SubExp] -> Map VName SubExp)
-> [Map VName SubExp] -> Map VName SubExp
forall a b. (a -> b) -> a -> b
$ (DeclExtType -> Type -> Map VName SubExp)
-> [DeclExtType] -> [Type] -> [Map VName SubExp]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith DeclExtType -> Type -> Map VName SubExp
combine [DeclExtType]
ts [Type]
ses_ts) (((VName, SubExp) -> InternaliseM ()) -> InternaliseM ())
-> ((VName, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \(VName
v, SubExp
se) ->
    Closure -> Exp (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
Closure -> Exp (Lore m) -> m ()
letBindNames [VName
v] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se

askSafety :: InternaliseM Safety
askSafety :: InternaliseM Safety
askSafety = do Bool
check <- (InternaliseEnv -> Bool) -> InternaliseM Bool
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks InternaliseEnv -> Bool
envDoBoundsChecks
               Safety -> InternaliseM Safety
forall (m :: * -> *) a. Monad m => a -> m a
return (Safety -> InternaliseM Safety) -> Safety -> InternaliseM Safety
forall a b. (a -> b) -> a -> b
$ if Bool
check then Safety
I.Safe else Safety
I.Unsafe

-- Implement partitioning using maps, scans and writes.
partitionWithSOACS :: Int -> I.Lambda -> [I.VName] -> InternaliseM ([I.SubExp], [I.SubExp])
partitionWithSOACS :: Int -> Lambda -> Closure -> InternaliseM ([SubExp], [SubExp])
partitionWithSOACS Int
k Lambda
lam Closure
arrs = do
  [Type]
arr_ts <- (VName -> InternaliseM Type) -> Closure -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType Closure
arrs
  let w :: SubExp
w = Int -> [Type] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [Type]
arr_ts
  Closure
classes_and_increments <- String -> Exp (Lore InternaliseM) -> InternaliseM Closure
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Closure
letTupExp String
"increments" (Exp (Lore InternaliseM) -> InternaliseM Closure)
-> Exp (Lore InternaliseM) -> InternaliseM Closure
forall a b. (a -> b) -> a -> b
$ Op SOACS -> ExpT SOACS
forall lore. Op lore -> ExpT lore
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> ScremaForm SOACS -> Closure -> SOAC SOACS
forall lore. SubExp -> ScremaForm lore -> Closure -> SOAC lore
I.Screma SubExp
w (Lambda -> ScremaForm SOACS
forall lore. Lambda lore -> ScremaForm lore
mapSOAC Lambda
lam) Closure
arrs
  (VName
classes, Closure
increments) <- case Closure
classes_and_increments of
                             VName
classes : Closure
increments -> (VName, Closure) -> InternaliseM (VName, Closure)
forall (m :: * -> *) a. Monad m => a -> m a
return (VName
classes, Int -> Closure -> Closure
forall a. Int -> [a] -> [a]
take Int
k Closure
increments)
                             Closure
_                    -> String -> InternaliseM (VName, Closure)
forall a. HasCallStack => String -> a
error String
"partitionWithSOACS"

  [Param Type]
add_lam_x_params <-
    Int -> InternaliseM (Param Type) -> InternaliseM [Param Type]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
k (InternaliseM (Param Type) -> InternaliseM [Param Type])
-> InternaliseM (Param Type) -> InternaliseM [Param Type]
forall a b. (a -> b) -> a -> b
$ VName -> Type -> Param Type
forall dec. VName -> dec -> Param dec
I.Param (VName -> Type -> Param Type)
-> InternaliseM VName -> InternaliseM (Type -> Param Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"x" InternaliseM (Type -> Param Type)
-> InternaliseM Type -> InternaliseM (Param Type)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> InternaliseM Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int32)
  [Param Type]
add_lam_y_params <-
    Int -> InternaliseM (Param Type) -> InternaliseM [Param Type]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
k (InternaliseM (Param Type) -> InternaliseM [Param Type])
-> InternaliseM (Param Type) -> InternaliseM [Param Type]
forall a b. (a -> b) -> a -> b
$ VName -> Type -> Param Type
forall dec. VName -> dec -> Param dec
I.Param (VName -> Type -> Param Type)
-> InternaliseM VName -> InternaliseM (Type -> Param Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"y" InternaliseM (Type -> Param Type)
-> InternaliseM Type -> InternaliseM (Param Type)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> InternaliseM Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int32)
  Body
add_lam_body <- Binder SOACS Body -> InternaliseM Body
forall lore (m :: * -> *) somelore.
(Bindable lore, MonadFreshNames m, HasScope somelore m,
 SameScope somelore lore) =>
Binder lore (Body lore) -> m (Body lore)
runBodyBinder (Binder SOACS Body -> InternaliseM Body)
-> Binder SOACS Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$
                  Scope SOACS -> Binder SOACS Body -> Binder SOACS Body
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param Type] -> Scope SOACS
forall lore dec.
(LParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
scopeOfLParams ([Param Type] -> Scope SOACS) -> [Param Type] -> Scope SOACS
forall a b. (a -> b) -> a -> b
$ [Param Type]
add_lam_x_params[Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
++[Param Type]
add_lam_y_params) (Binder SOACS Body -> Binder SOACS Body)
-> Binder SOACS Body -> Binder SOACS Body
forall a b. (a -> b) -> a -> b
$
    ([SubExp] -> Body)
-> BinderT SOACS (State VNameSource) [SubExp] -> Binder SOACS Body
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody (BinderT SOACS (State VNameSource) [SubExp] -> Binder SOACS Body)
-> BinderT SOACS (State VNameSource) [SubExp] -> Binder SOACS Body
forall a b. (a -> b) -> a -> b
$ [(Param Type, Param Type)]
-> ((Param Type, Param Type)
    -> BinderT SOACS (State VNameSource) SubExp)
-> BinderT SOACS (State VNameSource) [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([Param Type] -> [Param Type] -> [(Param Type, Param Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param Type]
add_lam_x_params [Param Type]
add_lam_y_params) (((Param Type, Param Type)
  -> BinderT SOACS (State VNameSource) SubExp)
 -> BinderT SOACS (State VNameSource) [SubExp])
-> ((Param Type, Param Type)
    -> BinderT SOACS (State VNameSource) SubExp)
-> BinderT SOACS (State VNameSource) [SubExp]
forall a b. (a -> b) -> a -> b
$ \(Param Type
x,Param Type
y) ->
      String
-> Exp (Lore (BinderT SOACS (State VNameSource)))
-> BinderT SOACS (State VNameSource) SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"z" (Exp (Lore (BinderT SOACS (State VNameSource)))
 -> BinderT SOACS (State VNameSource) SubExp)
-> Exp (Lore (BinderT SOACS (State VNameSource)))
-> BinderT SOACS (State VNameSource) SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Add IntType
Int32 Overflow
I.OverflowUndef)
      (VName -> SubExp
I.Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param Type -> VName
forall dec. Param dec -> VName
I.paramName Param Type
x) (VName -> SubExp
I.Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param Type -> VName
forall dec. Param dec -> VName
I.paramName Param Type
y)
  let add_lam :: Lambda
add_lam = Lambda :: forall lore. [LParam lore] -> BodyT lore -> [Type] -> LambdaT lore
I.Lambda { lambdaBody :: Body
I.lambdaBody = Body
add_lam_body
                         , lambdaParams :: [LParam]
I.lambdaParams = [Param Type]
add_lam_x_params [Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
++ [Param Type]
add_lam_y_params
                         , lambdaReturnType :: [Type]
I.lambdaReturnType = Int -> Type -> [Type]
forall a. Int -> a -> [a]
replicate Int
k (Type -> [Type]) -> Type -> [Type]
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int32
                         }
      nes :: [SubExp]
nes = Int -> SubExp -> [SubExp]
forall a. Int -> a -> [a]
replicate (Closure -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Closure
increments) (SubExp -> [SubExp]) -> SubExp -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Int32 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int32
0::Int32)

  ScremaForm SOACS
scan <- [Scan SOACS] -> InternaliseM (ScremaForm SOACS)
forall lore (m :: * -> *).
(Bindable lore, MonadFreshNames m) =>
[Scan lore] -> m (ScremaForm lore)
I.scanSOAC [Lambda -> [SubExp] -> Scan SOACS
forall lore. Lambda lore -> [SubExp] -> Scan lore
I.Scan Lambda
add_lam [SubExp]
nes]
  Closure
all_offsets <- String -> Exp (Lore InternaliseM) -> InternaliseM Closure
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Closure
letTupExp String
"offsets" (Exp (Lore InternaliseM) -> InternaliseM Closure)
-> Exp (Lore InternaliseM) -> InternaliseM Closure
forall a b. (a -> b) -> a -> b
$ Op SOACS -> ExpT SOACS
forall lore. Op lore -> ExpT lore
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> ScremaForm SOACS -> Closure -> SOAC SOACS
forall lore. SubExp -> ScremaForm lore -> Closure -> SOAC lore
I.Screma SubExp
w ScremaForm SOACS
scan Closure
increments

  -- We have the offsets for each of the partitions, but we also need
  -- the total sizes, which are the last elements in the offests.  We
  -- just have to be careful in case the array is empty.
  SubExp
last_index <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"last_index" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
Int32 Overflow
OverflowUndef) SubExp
w (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ Int32 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int32
1::Int32)
  Body
nonempty_body <- Binder SOACS Body -> InternaliseM Body
forall lore (m :: * -> *) somelore.
(Bindable lore, MonadFreshNames m, HasScope somelore m,
 SameScope somelore lore) =>
Binder lore (Body lore) -> m (Body lore)
runBodyBinder (Binder SOACS Body -> InternaliseM Body)
-> Binder SOACS Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$ ([SubExp] -> Body)
-> BinderT SOACS (State VNameSource) [SubExp] -> Binder SOACS Body
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody (BinderT SOACS (State VNameSource) [SubExp] -> Binder SOACS Body)
-> BinderT SOACS (State VNameSource) [SubExp] -> Binder SOACS Body
forall a b. (a -> b) -> a -> b
$ Closure
-> (VName -> BinderT SOACS (State VNameSource) SubExp)
-> BinderT SOACS (State VNameSource) [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM Closure
all_offsets ((VName -> BinderT SOACS (State VNameSource) SubExp)
 -> BinderT SOACS (State VNameSource) [SubExp])
-> (VName -> BinderT SOACS (State VNameSource) SubExp)
-> BinderT SOACS (State VNameSource) [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
offset_array ->
    String
-> Exp (Lore (BinderT SOACS (State VNameSource)))
-> BinderT SOACS (State VNameSource) SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"last_offset" (Exp (Lore (BinderT SOACS (State VNameSource)))
 -> BinderT SOACS (State VNameSource) SubExp)
-> Exp (Lore (BinderT SOACS (State VNameSource)))
-> BinderT SOACS (State VNameSource) SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ VName -> [DimIndex SubExp] -> BasicOp
I.Index VName
offset_array [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
I.DimFix SubExp
last_index]
  let empty_body :: Body
empty_body = [SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody ([SubExp] -> Body) -> [SubExp] -> Body
forall a b. (a -> b) -> a -> b
$ Int -> SubExp -> [SubExp]
forall a. Int -> a -> [a]
replicate Int
k (SubExp -> [SubExp]) -> SubExp -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Int32 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int32
0::Int32)
  SubExp
is_empty <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"is_empty" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
CmpEq PrimType
int32) SubExp
w (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ Int32 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int32
0::Int32)
  Closure
sizes <- String -> Exp (Lore InternaliseM) -> InternaliseM Closure
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Closure
letTupExp String
"partition_size" (Exp (Lore InternaliseM) -> InternaliseM Closure)
-> Exp (Lore InternaliseM) -> InternaliseM Closure
forall a b. (a -> b) -> a -> b
$
           SubExp -> Body -> Body -> IfDec (BranchType SOACS) -> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If SubExp
is_empty Body
empty_body Body
nonempty_body (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
           [Type] -> IfDec ExtType
ifCommon ([Type] -> IfDec ExtType) -> [Type] -> IfDec ExtType
forall a b. (a -> b) -> a -> b
$ Int -> Type -> [Type]
forall a. Int -> a -> [a]
replicate Int
k (Type -> [Type]) -> Type -> [Type]
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int32

  -- The total size of all partitions must necessarily be equal to the
  -- size of the input array.

  -- Create scratch arrays for the result.
  Closure
blanks <- [Type] -> (Type -> InternaliseM VName) -> InternaliseM Closure
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Type]
arr_ts ((Type -> InternaliseM VName) -> InternaliseM Closure)
-> (Type -> InternaliseM VName) -> InternaliseM Closure
forall a b. (a -> b) -> a -> b
$ \Type
arr_t ->
    String -> Exp (Lore InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m VName
letExp String
"partition_dest" (Exp (Lore InternaliseM) -> InternaliseM VName)
-> Exp (Lore InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
    PrimType -> [SubExp] -> BasicOp
Scratch (Type -> PrimType
forall shape u. TypeBase shape u -> PrimType
elemType Type
arr_t) (SubExp
w SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
drop Int
1 (Type -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
I.arrayDims Type
arr_t))

  -- Now write into the result.
  Lambda
write_lam <- do
    Param Type
c_param <- VName -> Type -> Param Type
forall dec. VName -> dec -> Param dec
I.Param (VName -> Type -> Param Type)
-> InternaliseM VName -> InternaliseM (Type -> Param Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"c" InternaliseM (Type -> Param Type)
-> InternaliseM Type -> InternaliseM (Param Type)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> InternaliseM Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int32)
    [Param Type]
offset_params <- Int -> InternaliseM (Param Type) -> InternaliseM [Param Type]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
k (InternaliseM (Param Type) -> InternaliseM [Param Type])
-> InternaliseM (Param Type) -> InternaliseM [Param Type]
forall a b. (a -> b) -> a -> b
$ VName -> Type -> Param Type
forall dec. VName -> dec -> Param dec
I.Param (VName -> Type -> Param Type)
-> InternaliseM VName -> InternaliseM (Type -> Param Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"offset" InternaliseM (Type -> Param Type)
-> InternaliseM Type -> InternaliseM (Param Type)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> InternaliseM Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int32)
    [Param Type]
value_params <- [Type]
-> (Type -> InternaliseM (Param Type)) -> InternaliseM [Param Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Type]
arr_ts ((Type -> InternaliseM (Param Type)) -> InternaliseM [Param Type])
-> (Type -> InternaliseM (Param Type)) -> InternaliseM [Param Type]
forall a b. (a -> b) -> a -> b
$ \Type
arr_t ->
      VName -> Type -> Param Type
forall dec. VName -> dec -> Param dec
I.Param (VName -> Type -> Param Type)
-> InternaliseM VName -> InternaliseM (Type -> Param Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"v" InternaliseM (Type -> Param Type)
-> InternaliseM Type -> InternaliseM (Param Type)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> InternaliseM Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
I.rowType Type
arr_t)
    (SubExp
offset, Stms SOACS
offset_stms) <- InternaliseM SubExp
-> InternaliseM (SubExp, Stms (Lore InternaliseM))
forall (m :: * -> *) a.
MonadBinder m =>
m a -> m (a, Stms (Lore m))
collectStms (InternaliseM SubExp
 -> InternaliseM (SubExp, Stms (Lore InternaliseM)))
-> InternaliseM SubExp
-> InternaliseM (SubExp, Stms (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ [SubExp] -> SubExp -> Int -> [LParam] -> InternaliseM SubExp
mkOffsetLambdaBody ((VName -> SubExp) -> Closure -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var Closure
sizes)
                             (VName -> SubExp
I.Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param Type -> VName
forall dec. Param dec -> VName
I.paramName Param Type
c_param) Int
0 [Param Type]
[LParam]
offset_params
    Lambda -> InternaliseM Lambda
forall (m :: * -> *) a. Monad m => a -> m a
return Lambda :: forall lore. [LParam lore] -> BodyT lore -> [Type] -> LambdaT lore
I.Lambda { lambdaParams :: [LParam]
I.lambdaParams = Param Type
c_param Param Type -> [Param Type] -> [Param Type]
forall a. a -> [a] -> [a]
: [Param Type]
offset_params [Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
++ [Param Type]
value_params
                    , lambdaReturnType :: [Type]
I.lambdaReturnType = Int -> Type -> [Type]
forall a. Int -> a -> [a]
replicate ([Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
arr_ts) (PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int32) [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++
                                           (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
I.rowType [Type]
arr_ts
                    , lambdaBody :: Body
I.lambdaBody = Stms SOACS -> [SubExp] -> Body
forall lore. Bindable lore => Stms lore -> [SubExp] -> Body lore
mkBody Stms SOACS
offset_stms ([SubExp] -> Body) -> [SubExp] -> Body
forall a b. (a -> b) -> a -> b
$
                                     Int -> SubExp -> [SubExp]
forall a. Int -> a -> [a]
replicate ([Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
arr_ts) SubExp
offset [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++
                                     (Param Type -> SubExp) -> [Param Type] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
I.Var (VName -> SubExp) -> (Param Type -> VName) -> Param Type -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> VName
forall dec. Param dec -> VName
I.paramName) [Param Type]
value_params
                    }
  Closure
results <- String -> Exp (Lore InternaliseM) -> InternaliseM Closure
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Closure
letTupExp String
"partition_res" (Exp (Lore InternaliseM) -> InternaliseM Closure)
-> Exp (Lore InternaliseM) -> InternaliseM Closure
forall a b. (a -> b) -> a -> b
$ Op SOACS -> ExpT SOACS
forall lore. Op lore -> ExpT lore
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> Lambda -> Closure -> [(SubExp, Int, VName)] -> SOAC SOACS
forall lore.
SubExp
-> Lambda lore -> Closure -> [(SubExp, Int, VName)] -> SOAC lore
I.Scatter SubExp
w
             Lambda
write_lam (VName
classes VName -> Closure -> Closure
forall a. a -> [a] -> [a]
: Closure
all_offsets Closure -> Closure -> Closure
forall a. [a] -> [a] -> [a]
++ Closure
arrs) ([(SubExp, Int, VName)] -> SOAC SOACS)
-> [(SubExp, Int, VName)] -> SOAC SOACS
forall a b. (a -> b) -> a -> b
$
             [SubExp] -> [Int] -> Closure -> [(SubExp, Int, VName)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 (SubExp -> [SubExp]
forall a. a -> [a]
repeat SubExp
w) (Int -> [Int]
forall a. a -> [a]
repeat Int
1) Closure
blanks
  SubExp
sizes' <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"partition_sizes" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
            [SubExp] -> Type -> BasicOp
I.ArrayLit ((VName -> SubExp) -> Closure -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var Closure
sizes) (Type -> BasicOp) -> Type -> BasicOp
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int32
  ([SubExp], [SubExp]) -> InternaliseM ([SubExp], [SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return ((VName -> SubExp) -> Closure -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var Closure
results, [SubExp
sizes'])
  where
    mkOffsetLambdaBody :: [SubExp]
                       -> SubExp
                       -> Int
                       -> [I.LParam]
                       -> InternaliseM SubExp
    mkOffsetLambdaBody :: [SubExp] -> SubExp -> Int -> [LParam] -> InternaliseM SubExp
mkOffsetLambdaBody [SubExp]
_ SubExp
_ Int
_ [] =
      SubExp -> InternaliseM SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> InternaliseM SubExp) -> SubExp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ Int32 -> SubExp
forall v. IsValue v => v -> SubExp
constant (-Int32
1::Int32)
    mkOffsetLambdaBody [SubExp]
sizes SubExp
c Int
i (LParam
p:[LParam]
ps) = do
      SubExp
is_this_one <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"is_this_one" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
CmpEq PrimType
int32) SubExp
c (Int -> SubExp
forall v. IsValue v => v -> SubExp
constant Int
i)
      SubExp
next_one <- [SubExp] -> SubExp -> Int -> [LParam] -> InternaliseM SubExp
mkOffsetLambdaBody [SubExp]
sizes SubExp
c (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [LParam]
ps
      SubExp
this_one <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"this_offset" (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<
                  BinOp
-> SubExp -> [SubExp] -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> SubExp -> [SubExp] -> m (Exp (Lore m))
foldBinOp (IntType -> Overflow -> BinOp
Add IntType
Int32 Overflow
OverflowUndef) (Int32 -> SubExp
forall v. IsValue v => v -> SubExp
constant (-Int32
1::Int32))
                  (VName -> SubExp
I.Var (Param Type -> VName
forall dec. Param dec -> VName
I.paramName Param Type
LParam
p) SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
take Int
i [SubExp]
sizes)
      String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"total_res" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ SubExp -> Body -> Body -> IfDec (BranchType SOACS) -> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If SubExp
is_this_one
        ([SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [SubExp
this_one]) ([SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [SubExp
next_one]) (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ [Type] -> IfDec ExtType
ifCommon [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int32]

typeExpForError :: E.TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError :: TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError (E.TEVar QualName VName
qn SrcLoc
_) =
  [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [String -> ErrorMsgPart SubExp
forall a. String -> ErrorMsgPart a
ErrorString (String -> ErrorMsgPart SubExp) -> String -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ QualName VName -> String
forall a. Pretty a => a -> String
pretty QualName VName
qn]
typeExpForError (E.TEUnique TypeExp VName
te SrcLoc
_) =
  (ErrorMsgPart SubExp
"*"ErrorMsgPart SubExp
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. a -> [a] -> [a]
:) ([ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp])
-> InternaliseM [ErrorMsgPart SubExp]
-> InternaliseM [ErrorMsgPart SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
te
typeExpForError (E.TEArray TypeExp VName
te DimExp VName
d SrcLoc
_) = do
  ErrorMsgPart SubExp
d' <- DimExp VName -> InternaliseM (ErrorMsgPart SubExp)
dimExpForError DimExp VName
d
  [ErrorMsgPart SubExp]
te' <- TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
te
  [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return ([ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp
"[", ErrorMsgPart SubExp
d', ErrorMsgPart SubExp
"]"] [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
te'
typeExpForError (E.TETuple [TypeExp VName]
tes SrcLoc
_) = do
  [[ErrorMsgPart SubExp]]
tes' <- (TypeExp VName -> InternaliseM [ErrorMsgPart SubExp])
-> [TypeExp VName] -> InternaliseM [[ErrorMsgPart SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError [TypeExp VName]
tes
  [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return ([ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp
"("] [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
-> [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
", "] [[ErrorMsgPart SubExp]]
tes' [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
")"]
typeExpForError (E.TERecord [(Name, TypeExp VName)]
fields SrcLoc
_) = do
  [[ErrorMsgPart SubExp]]
fields' <- ((Name, TypeExp VName) -> InternaliseM [ErrorMsgPart SubExp])
-> [(Name, TypeExp VName)] -> InternaliseM [[ErrorMsgPart SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Name, TypeExp VName) -> InternaliseM [ErrorMsgPart SubExp]
forall a.
Pretty a =>
(a, TypeExp VName) -> InternaliseM [ErrorMsgPart SubExp]
onField [(Name, TypeExp VName)]
fields
  [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return ([ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp
"{"] [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
-> [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
", "] [[ErrorMsgPart SubExp]]
fields' [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"}"]
  where onField :: (a, TypeExp VName) -> InternaliseM [ErrorMsgPart SubExp]
onField (a
k, TypeExp VName
te) =
          (String -> ErrorMsgPart SubExp
forall a. String -> ErrorMsgPart a
ErrorString (a -> String
forall a. Pretty a => a -> String
pretty a
k String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
": ")ErrorMsgPart SubExp
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. a -> [a] -> [a]
:) ([ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp])
-> InternaliseM [ErrorMsgPart SubExp]
-> InternaliseM [ErrorMsgPart SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
te
typeExpForError (E.TEArrow Maybe VName
_ TypeExp VName
t1 TypeExp VName
t2 SrcLoc
_) = do
  [ErrorMsgPart SubExp]
t1' <- TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
t1
  [ErrorMsgPart SubExp]
t2' <- TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
t2
  [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return ([ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp]
t1' [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
" -> "] [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
t2'
typeExpForError (E.TEApply TypeExp VName
t TypeArgExp VName
arg SrcLoc
_) = do
  [ErrorMsgPart SubExp]
t' <- TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
t
  [ErrorMsgPart SubExp]
arg' <- case TypeArgExp VName
arg of TypeArgExpType TypeExp VName
argt -> TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
argt
                      TypeArgExpDim DimExp VName
d SrcLoc
_   -> ErrorMsgPart SubExp -> [ErrorMsgPart SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ErrorMsgPart SubExp -> [ErrorMsgPart SubExp])
-> InternaliseM (ErrorMsgPart SubExp)
-> InternaliseM [ErrorMsgPart SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DimExp VName -> InternaliseM (ErrorMsgPart SubExp)
dimExpForError DimExp VName
d
  [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return ([ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp]
t' [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
" "] [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
arg'
typeExpForError (E.TESum [(Name, [TypeExp VName])]
cs SrcLoc
_) = do
  [[ErrorMsgPart SubExp]]
cs' <- ((Name, [TypeExp VName]) -> InternaliseM [ErrorMsgPart SubExp])
-> [(Name, [TypeExp VName])]
-> InternaliseM [[ErrorMsgPart SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([TypeExp VName] -> InternaliseM [ErrorMsgPart SubExp]
onClause ([TypeExp VName] -> InternaliseM [ErrorMsgPart SubExp])
-> ((Name, [TypeExp VName]) -> [TypeExp VName])
-> (Name, [TypeExp VName])
-> InternaliseM [ErrorMsgPart SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, [TypeExp VName]) -> [TypeExp VName]
forall a b. (a, b) -> b
snd) [(Name, [TypeExp VName])]
cs
  [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return ([ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp]
-> [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
" | "] [[ErrorMsgPart SubExp]]
cs'
  where onClause :: [TypeExp VName] -> InternaliseM [ErrorMsgPart SubExp]
onClause [TypeExp VName]
c = do
          [[ErrorMsgPart SubExp]]
c' <- (TypeExp VName -> InternaliseM [ErrorMsgPart SubExp])
-> [TypeExp VName] -> InternaliseM [[ErrorMsgPart SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError [TypeExp VName]
c
          [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return ([ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp]
-> [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
" "] [[ErrorMsgPart SubExp]]
c'

dimExpForError :: E.DimExp VName -> InternaliseM (ErrorMsgPart SubExp)
dimExpForError :: DimExp VName -> InternaliseM (ErrorMsgPart SubExp)
dimExpForError (DimExpNamed QualName VName
d SrcLoc
_) = do
  Maybe [SubExp]
substs <- VName -> InternaliseM (Maybe [SubExp])
lookupSubst (VName -> InternaliseM (Maybe [SubExp]))
-> VName -> InternaliseM (Maybe [SubExp])
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
E.qualLeaf QualName VName
d
  SubExp
d' <- case Maybe [SubExp]
substs of
          Just [SubExp
v] -> SubExp -> InternaliseM SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return SubExp
v
          Maybe [SubExp]
_        -> SubExp -> InternaliseM SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> InternaliseM SubExp) -> SubExp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
E.qualLeaf QualName VName
d
  ErrorMsgPart SubExp -> InternaliseM (ErrorMsgPart SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (ErrorMsgPart SubExp -> InternaliseM (ErrorMsgPart SubExp))
-> ErrorMsgPart SubExp -> InternaliseM (ErrorMsgPart SubExp)
forall a b. (a -> b) -> a -> b
$ SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt32 SubExp
d'
dimExpForError (DimExpConst Int
d SrcLoc
_) =
  ErrorMsgPart SubExp -> InternaliseM (ErrorMsgPart SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (ErrorMsgPart SubExp -> InternaliseM (ErrorMsgPart SubExp))
-> ErrorMsgPart SubExp -> InternaliseM (ErrorMsgPart SubExp)
forall a b. (a -> b) -> a -> b
$ String -> ErrorMsgPart SubExp
forall a. String -> ErrorMsgPart a
ErrorString (String -> ErrorMsgPart SubExp) -> String -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ Int -> String
forall a. Pretty a => a -> String
pretty Int
d
dimExpForError DimExp VName
DimExpAny = ErrorMsgPart SubExp -> InternaliseM (ErrorMsgPart SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return ErrorMsgPart SubExp
""

-- A smart constructor that compacts neighbouring literals for easier
-- reading in the IR.
errorMsg :: [ErrorMsgPart a] -> ErrorMsg a
errorMsg :: [ErrorMsgPart a] -> ErrorMsg a
errorMsg = [ErrorMsgPart a] -> ErrorMsg a
forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg ([ErrorMsgPart a] -> ErrorMsg a)
-> ([ErrorMsgPart a] -> [ErrorMsgPart a])
-> [ErrorMsgPart a]
-> ErrorMsg a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ErrorMsgPart a] -> [ErrorMsgPart a]
forall a. [ErrorMsgPart a] -> [ErrorMsgPart a]
compact
  where compact :: [ErrorMsgPart a] -> [ErrorMsgPart a]
compact [] = []
        compact (ErrorString String
x : ErrorString String
y : [ErrorMsgPart a]
parts) =
          [ErrorMsgPart a] -> [ErrorMsgPart a]
compact (String -> ErrorMsgPart a
forall a. String -> ErrorMsgPart a
ErrorString (String
xString -> String -> String
forall a. [a] -> [a] -> [a]
++String
y) ErrorMsgPart a -> [ErrorMsgPart a] -> [ErrorMsgPart a]
forall a. a -> [a] -> [a]
: [ErrorMsgPart a]
parts)
        compact (ErrorMsgPart a
x:[ErrorMsgPart a]
y) = ErrorMsgPart a
x ErrorMsgPart a -> [ErrorMsgPart a] -> [ErrorMsgPart a]
forall a. a -> [a] -> [a]
: [ErrorMsgPart a] -> [ErrorMsgPart a]
compact [ErrorMsgPart a]
y