{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Safe #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}

-- |
--
-- This module implements a transformation from source to core
-- Futhark.
module Futhark.Internalise (internaliseProg) where

import Control.Monad.Reader
import Control.Monad.State
import Data.Bitraversable
import Data.List (find, intercalate, intersperse, nub, transpose)
import qualified Data.List.NonEmpty as NE
import qualified Data.Map.Strict as M
import qualified Data.Set as S
import Futhark.IR.SOACS as I hiding (stmPattern)
import Futhark.Internalise.AccurateSizes
import Futhark.Internalise.Bindings
import Futhark.Internalise.Defunctionalise as Defunctionalise
import Futhark.Internalise.Defunctorise as Defunctorise
import Futhark.Internalise.Lambdas
import Futhark.Internalise.Monad as I
import Futhark.Internalise.Monomorphise as Monomorphise
import Futhark.Internalise.TypesValues
import Futhark.Transform.Rename as I
import Futhark.Util (splitAt3)
import Language.Futhark as E hiding (TypeArg)
import Language.Futhark.Semantic (Imports)

-- | 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
      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

      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."]

      (([DeclExtType]
rettype', Result
body_res), Stms SOACS
body_stms) <- InternaliseM ([DeclExtType], Result)
-> InternaliseM (([DeclExtType], Result), Stms (Lore InternaliseM))
forall (m :: * -> *) a.
MonadBinder m =>
m a -> m (a, Stms (Lore m))
collectStms (InternaliseM ([DeclExtType], Result)
 -> InternaliseM
      (([DeclExtType], Result), Stms (Lore InternaliseM)))
-> InternaliseM ([DeclExtType], Result)
-> InternaliseM (([DeclExtType], Result), Stms (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ do
        Result
body_res <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
"res" ExpBase Info VName
body
        [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
        ([DeclExtType], Result) -> InternaliseM ([DeclExtType], Result)
forall (m :: * -> *) a. Monad m => a -> m a
return ([DeclExtType]
rettype', Result
body_res)
      Body
body' <-
        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') (Body -> InternaliseM Body) -> Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$
          Stms SOACS -> Result -> Body
forall lore. Bindable lore => Stms lore -> Result -> Body lore
mkBody Stms SOACS
body_stms Result
body_res

      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
    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
_ =
      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 String
"size"

-- | 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

data EntryTrust
  = -- | This parameter or return value is an opaque type.  When a
    -- parameter, this implies that it must have been returned by a
    -- previous call to Futhark, and hence we can preserve (constant)
    -- size constraints.
    EntryTrusted
  | -- | The type is directly exposed.  Any size constraint cannot be
    -- trusted.
    EntryUntrusted

entryTrust :: EntryType -> EntryTrust
entryTrust :: EntryType -> EntryTrust
entryTrust EntryType
t
  | E.Scalar (E.Prim E.Unsigned {}) <- EntryType -> StructType
E.entryType EntryType
t =
    EntryTrust
EntryUntrusted
  | E.Array ()
_ Uniqueness
_ (E.Prim E.Unsigned {}) ShapeDecl (DimDecl VName)
_ <- EntryType -> StructType
E.entryType EntryType
t =
    EntryTrust
EntryUntrusted
  | E.Scalar E.Prim {} <- EntryType -> StructType
E.entryType EntryType
t =
    EntryTrust
EntryUntrusted
  | E.Array ()
_ Uniqueness
_ E.Prim {} ShapeDecl (DimDecl VName)
_ <- EntryType -> StructType
E.entryType EntryType
t =
    EntryTrust
EntryUntrusted
  | Bool
otherwise =
    EntryTrust
EntryTrusted

fixEntryParamSizes :: MonadFreshNames m => E.Pattern -> EntryTrust -> m E.Pattern
fixEntryParamSizes :: Pattern -> EntryTrust -> m Pattern
fixEntryParamSizes Pattern
p EntryTrust
EntryTrusted = Pattern -> m Pattern
forall (f :: * -> *) a. Applicative f => a -> f a
pure Pattern
p
fixEntryParamSizes Pattern
p EntryTrust
EntryUntrusted = Pattern -> m Pattern
forall (m :: * -> *). MonadFreshNames m => Pattern -> m Pattern
allDimsFreshInPat Pattern
p

-- When we are returning a value from the entry point, we fully
-- existentialise the return type.  This is because it might otherwise
-- refer to sizes that are not in scope, because the generated entry
-- point function does not keep the size parameters of the original
-- entry point.
fullyExistential ::
  [[I.TypeBase ExtShape u]] ->
  [[I.TypeBase ExtShape u]]
fullyExistential :: [[TypeBase ExtShape u]] -> [[TypeBase ExtShape u]]
fullyExistential [[TypeBase ExtShape u]]
tss =
  State Int [[TypeBase ExtShape u]] -> Int -> [[TypeBase ExtShape u]]
forall s a. State s a -> s -> a
evalState (([TypeBase ExtShape u]
 -> StateT Int Identity [TypeBase ExtShape u])
-> [[TypeBase ExtShape u]] -> State Int [[TypeBase ExtShape u]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((TypeBase ExtShape u -> StateT Int Identity (TypeBase ExtShape u))
-> [TypeBase ExtShape u]
-> StateT Int Identity [TypeBase ExtShape u]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((ExtShape -> StateT Int Identity ExtShape)
-> (u -> StateT Int Identity u)
-> TypeBase ExtShape u
-> StateT Int Identity (TypeBase ExtShape u)
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 ((ExtSize -> StateT Int Identity ExtSize)
-> ExtShape -> StateT Int Identity ExtShape
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ExtSize -> StateT Int Identity ExtSize
forall (m :: * -> *) p a. MonadState Int m => p -> m (Ext a)
onDim) u -> StateT Int Identity u
forall (f :: * -> *) a. Applicative f => a -> f a
pure)) [[TypeBase ExtShape u]]
tss) Int
0
  where
    onDim :: p -> m (Ext a)
onDim p
_ = do
      Int
i <- m Int
forall s (m :: * -> *). MonadState s m => m s
get
      (Int -> Int) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
      Ext a -> m (Ext a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Ext a -> m (Ext a)) -> Ext a -> m (Ext a)
forall a b. (a -> b) -> a -> b
$ Int -> Ext a
forall a. Int -> Ext a
Ext Int
i

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 -> EntryTrust -> InternaliseM Pattern)
-> [Pattern] -> [EntryTrust] -> InternaliseM [Pattern]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM Pattern -> EntryTrust -> InternaliseM Pattern
forall (m :: * -> *).
MonadFreshNames m =>
Pattern -> EntryTrust -> m Pattern
fixEntryParamSizes [Pattern]
params ([EntryTrust] -> InternaliseM [Pattern])
-> [EntryTrust] -> InternaliseM [Pattern]
forall a b. (a -> b) -> a -> b
$ (EntryType -> EntryTrust) -> [EntryType] -> [EntryTrust]
forall a b. (a -> b) -> [a] -> [b]
map EntryType -> EntryTrust
entryTrust [EntryType]
e_paramts
  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 <- [[DeclExtType]] -> [[DeclExtType]]
forall u. [[TypeBase ExtShape u]] -> [[TypeBase ExtShape u]]
fullyExistential ([[DeclExtType]] -> [[DeclExtType]])
-> InternaliseM [[DeclExtType]] -> InternaliseM [[DeclExtType]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StructType -> InternaliseM [[DeclExtType]]
internaliseEntryReturnType 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 :: Result
args = (Param DeclType -> SubExp) -> [Param DeclType] -> Result
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] -> Result) -> [Param DeclType] -> Result
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 Result
maybe_const <- VName -> InternaliseM (Maybe Result)
lookupConst VName
ofname
      Result
vals <- case Maybe Result
maybe_const of
        Just Result
ses ->
          Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
ses
        Maybe Result
Nothing ->
          (Result, [ExtType]) -> Result
forall a b. (a, b) -> a
fst ((Result, [ExtType]) -> Result)
-> InternaliseM (Result, [ExtType]) -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String
-> QualName VName
-> Result
-> SrcLoc
-> InternaliseM (Result, [ExtType])
funcall String
"entry_result" (VName -> QualName VName
forall v. v -> QualName v
E.qualName VName
ofname) Result
args SrcLoc
loc
      Result
ctx <-
        [DeclExtType] -> [Result] -> Result
forall u a. [TypeBase ExtShape u] -> [[a]] -> [a]
extractShapeContext ([[DeclExtType]] -> [DeclExtType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[DeclExtType]]
entry_rettype)
          ([Result] -> Result)
-> InternaliseM [Result] -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> InternaliseM Result) -> Result -> InternaliseM [Result]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Type -> Result) -> InternaliseM Type -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Type -> Result
forall u. TypeBase Shape u -> Result
I.arrayDims (InternaliseM Type -> InternaliseM Result)
-> (SubExp -> InternaliseM Type) -> SubExp -> InternaliseM Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType) Result
vals
      Result -> InternaliseM (Body (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => Result -> m (Body (Lore m))
resultBodyM (Result
ctx Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ Result
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
    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
d Int -> 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
$ Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody (Result -> Body) -> InternaliseM Result -> InternaliseM Body
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
"res" ExpBase Info VName
e

bodyFromStms ::
  InternaliseM (Result, a) ->
  InternaliseM (Body, a)
bodyFromStms :: InternaliseM (Result, a) -> InternaliseM (Body, a)
bodyFromStms InternaliseM (Result, a)
m = do
  ((Result
res, a
a), Stms SOACS
stms) <- InternaliseM (Result, a)
-> InternaliseM ((Result, a), Stms (Lore InternaliseM))
forall (m :: * -> *) a.
MonadBinder m =>
m a -> m (a, Stms (Lore m))
collectStms InternaliseM (Result, 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)
-> Result -> InternaliseM (Body (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
Stms (Lore m) -> Result -> m (Body (Lore m))
mkBodyM Stms (Lore InternaliseM)
Stms SOACS
stms Result
res

internaliseExp :: String -> E.Exp -> InternaliseM [I.SubExp]
internaliseExp :: String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc (E.Parens ExpBase Info VName
e SrcLoc
_) =
  String -> ExpBase Info VName -> InternaliseM Result
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 Result
internaliseExp String
desc ExpBase Info VName
e
internaliseExp String
desc (E.StringLit [Word8]
vs SrcLoc
_) =
  (SubExp -> Result) -> InternaliseM SubExp -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> Result
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InternaliseM SubExp -> InternaliseM Result)
-> InternaliseM SubExp -> InternaliseM Result
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
$ Result -> Type -> BasicOp
I.ArrayLit ((Word8 -> SubExp) -> [Word8] -> Result
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 Result
subst <- VName -> InternaliseM (Maybe Result)
lookupSubst VName
name
  case Maybe Result
subst of
    Just Result
substs -> Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
substs
    Maybe Result
Nothing -> do
      -- If this identifier is the name of a constant, we have to turn it
      -- into a call to the corresponding function.
      Maybe Result
is_const <- VName -> InternaliseM (Maybe Result)
lookupConst VName
name
      case Maybe Result
is_const of
        Just Result
ses -> Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
ses
        Maybe Result
Nothing -> (SubExp -> Result -> Result
forall a. a -> [a] -> [a]
: []) (SubExp -> Result) -> (VName -> SubExp) -> VName -> Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
I.Var (VName -> Result) -> InternaliseM VName -> InternaliseM Result
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
  Result
dims <- case Closure
vs of
    [] -> Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return [] -- Will this happen?
    VName
v : Closure
_ -> Type -> Result
forall u. TypeBase Shape u -> Result
I.arrayDims (Type -> Result) -> InternaliseM Type -> InternaliseM Result
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
-> Result
-> [DimIndexBase Info VName]
-> InternaliseM ([DimIndex SubExp], Certificates)
internaliseSlice SrcLoc
loc Result
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'
  Result
ses <- Certificates -> InternaliseM Result -> InternaliseM Result
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
cs (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ String -> [Exp (Lore InternaliseM)] -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
String -> [Exp (Lore m)] -> m Result
letSubExps String
desc ([ExpT SOACS] -> InternaliseM Result)
-> InternaliseM [ExpT SOACS] -> InternaliseM Result
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 -> Result -> InternaliseM ()
bindExtSizes (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret) Closure
retext Result
ses
  Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
ses

-- XXX: we map empty records and tuples to bools, because otherwise
-- arrays of unit will lose their sizes.
internaliseExp String
_ (E.TupLit [] SrcLoc
_) =
  Result -> InternaliseM Result
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
_) =
  Result -> InternaliseM Result
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
_) = [Result] -> Result
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([Result] -> Result)
-> InternaliseM [Result] -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ExpBase Info VName -> InternaliseM Result)
-> [ExpBase Info VName] -> InternaliseM [Result]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc) [ExpBase Info VName]
es
internaliseExp String
desc (E.RecordLit [FieldBase Info VName]
orig_fields SrcLoc
_) =
  ((Name, Result) -> Result) -> [(Name, Result)] -> Result
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Name, Result) -> Result
forall a b. (a, b) -> b
snd ([(Name, Result)] -> Result)
-> ([Map Name Result] -> [(Name, Result)])
-> [Map Name Result]
-> Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name Result -> [(Name, Result)]
forall a. Map Name a -> [(Name, a)]
sortFields (Map Name Result -> [(Name, Result)])
-> ([Map Name Result] -> Map Name Result)
-> [Map Name Result]
-> [(Name, Result)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Map Name Result] -> Map Name Result
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
M.unions ([Map Name Result] -> Result)
-> InternaliseM [Map Name Result] -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (FieldBase Info VName -> InternaliseM (Map Name Result))
-> [FieldBase Info VName] -> InternaliseM [Map Name Result]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM FieldBase Info VName -> InternaliseM (Map Name Result)
internaliseField [FieldBase Info VName]
orig_fields
  where
    internaliseField :: FieldBase Info VName -> InternaliseM (Map Name Result)
internaliseField (E.RecordFieldExplicit Name
name ExpBase Info VName
e SrcLoc
_) =
      Name -> Result -> Map Name Result
forall k a. k -> a -> Map k a
M.singleton Name
name (Result -> Map Name Result)
-> InternaliseM Result -> InternaliseM (Map Name Result)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc ExpBase Info VName
e
    internaliseField (E.RecordFieldImplicit VName
name Info PatternType
t SrcLoc
loc) =
      FieldBase Info VName -> InternaliseM (Map Name Result)
internaliseField (FieldBase Info VName -> InternaliseM (Map Name Result))
-> FieldBase Info VName -> InternaliseM (Map Name Result)
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]
es Int -> [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 Result
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 Result)
-> (VName -> InternaliseM SubExp) -> InternaliseM Result
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
. IntType -> Integer -> SubExp
intConst IntType
Int32 (Integer -> SubExp) -> (Int -> Integer) -> Int -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Integer
forall a. Integral a => a -> Integer
toInteger) [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
    [Result]
es' <- (ExpBase Info VName -> InternaliseM Result)
-> [ExpBase Info VName] -> InternaliseM [Result]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
"arr_elem") [ExpBase Info VName]
es
    [DeclExtType]
arr_t_ext <- StructType -> InternaliseM [DeclExtType]
internaliseReturnType (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
arr_t)

    [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
forall u. TypeBase ExtShape u -> Maybe (TypeBase Shape u)
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 [Result]
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
            Result
e' : [Result]
_ -> (SubExp -> InternaliseM Type) -> Result -> 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 Result
e'

    let arraylit :: Result -> Type -> InternaliseM (ExpT SOACS)
arraylit Result
ks Type
rt = do
          Result
ks' <-
            (SubExp -> InternaliseM SubExp) -> Result -> InternaliseM Result
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"
              )
              Result
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
$ Result -> Type -> BasicOp
I.ArrayLit Result
ks' Type
rt

    String -> [Exp (Lore InternaliseM)] -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
String -> [Exp (Lore m)] -> m Result
letSubExps String
desc
      ([ExpT SOACS] -> InternaliseM Result)
-> InternaliseM [ExpT SOACS] -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< if [Result] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Result]
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 (Result -> Type -> InternaliseM (ExpT SOACS)
arraylit []) [Type]
rowtypes
        else (Result -> Type -> InternaliseM (ExpT SOACS))
-> [Result] -> [Type] -> InternaliseM [ExpT SOACS]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM Result -> Type -> InternaliseM (ExpT SOACS)
arraylit ([Result] -> [Result]
forall a. [[a]] -> [[a]]
transpose [Result]
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_es Int -> [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
asIntS 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
asIntS 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
asIntS 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
            (Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody [SubExp
bounds_invalid_downwards])
            (Result -> Body
forall lore. Bindable lore => Result -> 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
            (Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody [SubExp
distance_downwards_exclusive])
            (Result -> Body
forall lore. Bindable lore => Result -> 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
asIntS 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 -> Result -> InternaliseM ()
bindExtSizes (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret) Closure
retext [SubExp
se]
  Result -> InternaliseM Result
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 Result
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
  Result
ses <- String -> ExpBase Info VName -> InternaliseM Result
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 -> Result -> InternaliseM ()
bindExtSizes (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret) Closure
retext Result
ses
  [(SubExp, DeclExtType)]
-> ((SubExp, DeclExtType) -> InternaliseM SubExp)
-> InternaliseM Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (Result -> [DeclExtType] -> [(SubExp, DeclExtType)]
forall a b. [a] -> [b] -> [(a, b)]
zip Result
ses [DeclExtType]
ts) (((SubExp, DeclExtType) -> InternaliseM SubExp)
 -> InternaliseM Result)
-> ((SubExp, DeclExtType) -> InternaliseM SubExp)
-> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ \(SubExp
e', DeclExtType
t') -> do
    Result
dims <- Type -> Result
forall u. TypeBase Shape u -> Result
arrayDims (Type -> Result) -> InternaliseM Type -> InternaliseM Result
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) -> Result -> [ErrorMsgPart SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt32 Result
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 Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
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 Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
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 Result
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.
  Result
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 Result
internalise <- QualName VName
-> [ExpBase Info VName]
-> SrcLoc
-> Maybe (String -> InternaliseM Result)
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 Result
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]
          [Result]
args' <- [Result] -> [Result]
forall a. [a] -> [a]
reverse ([Result] -> [Result])
-> InternaliseM [Result] -> InternaliseM [Result]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((ExpBase Info VName, Maybe VName) -> InternaliseM Result)
-> [(ExpBase Info VName, Maybe VName)] -> InternaliseM [Result]
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 Result
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'' = (Result -> [(SubExp, Diet)]) -> [Result] -> [(SubExp, Diet)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Result -> [(SubExp, Diet)]
forall a. [a] -> [(a, Diet)]
tag [Result]
args'
          String -> Exp (Lore InternaliseM) -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
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
          Result
args' <- [Result] -> Result
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([Result] -> Result)
-> ([Result] -> [Result]) -> [Result] -> Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Result] -> [Result]
forall a. [a] -> [a]
reverse ([Result] -> Result)
-> InternaliseM [Result] -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((ExpBase Info VName, Maybe VName) -> InternaliseM Result)
-> [(ExpBase Info VName, Maybe VName)] -> InternaliseM [Result]
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 Result
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)
          (Result, [ExtType]) -> Result
forall a b. (a, b) -> a
fst ((Result, [ExtType]) -> Result)
-> InternaliseM (Result, [ExtType]) -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String
-> QualName VName
-> Result
-> SrcLoc
-> InternaliseM (Result, [ExtType])
funcall String
desc QualName VName
qfname Result
args' SrcLoc
loc

  StructType -> Closure -> Result -> InternaliseM ()
bindExtSizes StructType
ret Closure
retext Result
ses
  Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
ses
internaliseExp String
desc (E.LetPat Pattern
pat ExpBase Info VName
e ExpBase Info VName
body (Info PatternType
ret, Info Closure
retext) SrcLoc
_) = do
  Result
ses <- String
-> Pattern
-> ExpBase Info VName
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM Result)
-> InternaliseM Result
forall a.
String
-> Pattern
-> ExpBase Info VName
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat String
desc Pattern
pat ExpBase Info VName
e ExpBase Info VName
body (String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc)
  StructType -> Closure -> Result -> InternaliseM ()
bindExtSizes (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret) Closure
retext Result
ses
  Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
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 Result
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
  Result
ses <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
"loop_init" ExpBase Info VName
mergeexp
  ((Body
loopbody', (LoopForm SOACS
form', [Param DeclType]
shapepat, [Param DeclType]
mergepat', Result
mergeinit')), Stms SOACS
initstms) <-
    InternaliseM
  (Body,
   (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
-> InternaliseM
     ((Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)),
      Stms (Lore InternaliseM))
forall (m :: * -> *) a.
MonadBinder m =>
m a -> m (a, Stms (Lore m))
collectStms (InternaliseM
   (Body,
    (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
 -> InternaliseM
      ((Body,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)),
       Stms (Lore InternaliseM)))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
-> InternaliseM
     ((Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)),
      Stms (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ Result
-> LoopFormBase Info VName
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
handleForm Result
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) -> Result -> 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 Result
mergeinit'

  Result
ctxinit <- Closure -> [FParam] -> [Type] -> InternaliseM Result
forall (m :: * -> *).
(HasScope SOACS m, Monad m) =>
Closure -> [FParam] -> [Type] -> m Result
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]
[FParam]
mergepat' [Type]
mergeinit_ts'

  let ctxmerge :: [(Param DeclType, SubExp)]
ctxmerge = [Param DeclType] -> Result -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
shapepat Result
ctxinit
      valmerge :: [(Param DeclType, SubExp)]
valmerge = [Param DeclType] -> Result -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
mergepat' Result
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
$
          Result -> InternaliseM Body
forall (m :: * -> *). MonadBinder m => Result -> m (Body (Lore m))
resultBodyM
            (Result -> InternaliseM Body)
-> InternaliseM Result -> InternaliseM Body
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ErrorMsg SubExp
-> SrcLoc -> Closure -> [Type] -> Result -> InternaliseM Result
forall u.
Typed (TypeBase Shape u) =>
ErrorMsg SubExp
-> SrcLoc
-> Closure
-> [TypeBase Shape u]
-> Result
-> InternaliseM Result
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
            (Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Body (Lore InternaliseM) -> InternaliseM Result
forall (m :: * -> *). MonadBinder m => Body (Lore m) -> m Result
bodyBind Body (Lore InternaliseM)
Body
loopbody'

  Attrs
attrs <- (InternaliseEnv -> Attrs) -> InternaliseM Attrs
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks InternaliseEnv -> Attrs
envAttrs
  Result
loop_res <-
    (VName -> SubExp) -> Closure -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var (Closure -> Result) -> (Closure -> Closure) -> Closure -> Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Closure -> Closure
dropCond
      (Closure -> Result) -> InternaliseM Closure -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Attrs -> InternaliseM Closure -> InternaliseM Closure
forall (m :: * -> *) a. MonadBinder m => Attrs -> m a -> m a
attributing
        Attrs
attrs
        (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 -> Result -> InternaliseM ()
bindExtSizes (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret) Closure
retext Result
loop_res
  Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
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]
-> Result
-> LoopForm SOACS
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forLoop [Param DeclType]
mergepat' [Param DeclType]
shapepat Result
mergeinit LoopForm SOACS
form' =
      InternaliseM
  (Result,
   (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a. InternaliseM (Result, a) -> InternaliseM (Body, a)
bodyFromStms (InternaliseM
   (Result,
    (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a b. (a -> b) -> a -> b
$
        LoopForm SOACS
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall lore a (m :: * -> *) b.
(Scoped lore a, LocalScope lore m) =>
a -> m b -> m b
inScopeOf LoopForm SOACS
form' (InternaliseM
   (Result,
    (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
 -> InternaliseM
      (Result,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a b. (a -> b) -> a -> b
$ do
          Result
ses <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
"loopres" ExpBase Info VName
loopbody
          [Type]
sets <- (SubExp -> InternaliseM Type) -> Result -> 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 Result
ses
          Result
shapeargs <- Closure -> [FParam] -> [Type] -> InternaliseM Result
forall (m :: * -> *).
(HasScope SOACS m, Monad m) =>
Closure -> [FParam] -> [Type] -> m Result
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]
[FParam]
mergepat' [Type]
sets
          (Result,
 (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall (m :: * -> *) a. Monad m => a -> m a
return
            ( Result
shapeargs Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ Result
ses,
              ( LoopForm SOACS
form',
                [Param DeclType]
shapepat,
                [Param DeclType]
mergepat',
                Result
mergeinit
              )
            )

    handleForm :: Result
-> LoopFormBase Info VName
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
handleForm Result
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], Result)))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a.
[TypeParamBase VName]
-> Pattern
-> ([FParam] -> [FParam] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' Pattern
mergepat (([FParam]
  -> [FParam]
  -> InternaliseM
       (Body,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> ([FParam]
    -> [FParam]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a b. (a -> b) -> a -> b
$
        \[FParam]
shapepat [FParam]
mergepat' ->
          [Pattern]
-> [Type]
-> ([LParam]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
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], Result)))
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> ([LParam]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
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]
-> Result
-> LoopForm SOACS
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forLoop [Param DeclType]
[FParam]
mergepat' [Param DeclType]
[FParam]
shapepat Result
mergeinit (LoopForm SOACS
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> LoopForm SOACS
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
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 Result
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], Result)))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a.
[TypeParamBase VName]
-> Pattern
-> ([FParam] -> [FParam] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' Pattern
mergepat (([FParam]
  -> [FParam]
  -> InternaliseM
       (Body,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> ([FParam]
    -> [FParam]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a b. (a -> b) -> a -> b
$
        \[FParam]
shapepat [FParam]
mergepat' ->
          [Param DeclType]
-> [Param DeclType]
-> Result
-> LoopForm SOACS
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forLoop [Param DeclType]
[FParam]
mergepat' [Param DeclType]
[FParam]
shapepat Result
mergeinit (LoopForm SOACS
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> LoopForm SOACS
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
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 Result
mergeinit (E.While ExpBase Info VName
cond) =
      [TypeParamBase VName]
-> Pattern
-> ([FParam]
    -> [FParam]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a.
[TypeParamBase VName]
-> Pattern
-> ([FParam] -> [FParam] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' Pattern
mergepat (([FParam]
  -> [FParam]
  -> InternaliseM
       (Body,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> ([FParam]
    -> [FParam]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a b. (a -> b) -> a -> b
$ \[FParam]
shapepat [FParam]
mergepat' -> do
        [Type]
mergeinit_ts <- (SubExp -> InternaliseM Type) -> Result -> 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 Result
mergeinit
        -- 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.
        Result
shapeinit <- Closure -> [FParam] -> [Type] -> InternaliseM Result
forall (m :: * -> *).
(HasScope SOACS m, Monad m) =>
Closure -> [FParam] -> [Type] -> m Result
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) [FParam]
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] -> Result -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
[FParam]
shapepat Result
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] -> Result -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
[FParam]
mergepat' Result
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) -> Result -> ShapeChange SubExp
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimCoercion (Result -> ShapeChange SubExp) -> Result -> ShapeChange SubExp
forall a b. (a -> b) -> a -> b
$ Type -> Result
forall u. TypeBase Shape u -> Result
arrayDims (Type -> Result) -> Type -> Result
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
  (Result,
   (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a. InternaliseM (Result, a) -> InternaliseM (Body, a)
bodyFromStms (InternaliseM
   (Result,
    (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a b. (a -> b) -> a -> b
$ do
          Result
ses <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
"loopres" ExpBase Info VName
loopbody
          [Type]
sets <- (SubExp -> InternaliseM Type) -> Result -> 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 Result
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
          Result
shapeargs <- Closure -> [FParam] -> [Type] -> InternaliseM Result
forall (m :: * -> *).
(HasScope SOACS m, Monad m) =>
Closure -> [FParam] -> [Type] -> m Result
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) [FParam]
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] -> Result -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
[FParam]
shapepat Result
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] -> Result -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
[FParam]
mergepat' Result
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) -> Result -> ShapeChange SubExp
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimCoercion (Result -> ShapeChange SubExp) -> Result -> ShapeChange SubExp
forall a b. (a -> b) -> a -> b
$ Type -> Result
forall u. TypeBase Shape u -> Result
arrayDims (Type -> Result) -> Type -> Result
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
            Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody (Result -> Body) -> InternaliseM Result -> InternaliseM Body
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
"loop_cond" ExpBase Info VName
cond
          Result
loop_end_cond <- Body (Lore InternaliseM) -> InternaliseM Result
forall (m :: * -> *). MonadBinder m => Body (Lore m) -> m Result
bodyBind Body (Lore InternaliseM)
Body
loop_end_cond_body

          (Result,
 (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall (m :: * -> *) a. Monad m => a -> m a
return
            ( Result
shapeargs Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ Result
loop_end_cond Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ Result
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]
[FParam]
mergepat',
                SubExp
loop_initial_cond SubExp -> Result -> Result
forall a. a -> [a] -> [a]
: Result
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 Result
internaliseExp String
desc (ExpBase Info VName -> InternaliseM Result)
-> ExpBase Info VName -> InternaliseM Result
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
  Result
ves <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
"lw_val" ExpBase Info VName
ve
  Closure
srcs <- String -> ExpBase Info VName -> InternaliseM Closure
internaliseExpToVars String
"src" ExpBase Info VName
src
  Result
dims <- case Closure
srcs of
    [] -> Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return [] -- Will this happen?
    VName
v : Closure
_ -> Type -> Result
forall u. TypeBase Shape u -> Result
I.arrayDims (Type -> Result) -> InternaliseM Type -> InternaliseM Result
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
-> Result
-> [DimIndexBase Info VName]
-> InternaliseM ([DimIndex SubExp], Certificates)
internaliseSlice SrcLoc
loc Result
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 -> Result -> Type
forall oldshape u.
TypeBase oldshape u -> Result -> TypeBase Shape u
`setArrayDims` [DimIndex SubExp] -> Result
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 Result -> InternaliseM Result
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
cs (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> Closure -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var (Closure -> Result) -> InternaliseM Closure -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> SubExp -> InternaliseM VName)
-> Closure -> Result -> 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 Result
ves
internaliseExp String
desc (E.RecordUpdate ExpBase Info VName
src [Name]
fields ExpBase Info VName
ve Info PatternType
_ SrcLoc
_) = do
  Result
src' <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc ExpBase Info VName
src
  Result
ve' <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc ExpBase Info VName
ve
  StructType -> [Name] -> Result -> Result -> InternaliseM Result
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 Result
ve' Result
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 Result -> InternaliseM Result
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local InternaliseEnv -> InternaliseEnv
f (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ String -> ExpBase Info VName -> InternaliseM Result
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 Result -> InternaliseM Result
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
c (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ (SubExp -> InternaliseM SubExp) -> Result -> InternaliseM Result
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 (Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> ExpBase Info VName -> InternaliseM Result
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
  Result
es' <- [Result] -> Result
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([Result] -> Result)
-> InternaliseM [Result] -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ExpBase Info VName -> InternaliseM Result)
-> [ExpBase Info VName] -> InternaliseM [Result]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> ExpBase Info VName -> InternaliseM Result
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 -> Result -> Result
forall a. a -> [a] -> [a]
:) (Result -> Result) -> InternaliseM Result -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [Type] -> [(Int, SubExp)] -> InternaliseM Result
forall (f :: * -> *) a.
(Num a, MonadBinder f, Eq a) =>
a -> [Type] -> [(a, SubExp)] -> f Result
clauses Int
0 [Type]
ts' ([Int] -> Result -> [(Int, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
js Result
es')
    Maybe (Int, [Int])
Nothing ->
      String -> InternaliseM Result
forall a. HasCallStack => String -> a
error String
"internaliseExp Constr: missing constructor"
  where
    clauses :: a -> [Type] -> [(a, SubExp)] -> f Result
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
e SubExp -> Result -> Result
forall a. a -> [a] -> [a]
:) (Result -> Result) -> f Result -> f Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> [Type] -> [(a, SubExp)] -> f Result
clauses (a
j a -> 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
blank SubExp -> Result -> Result
forall a. a -> [a] -> [a]
:) (Result -> Result) -> f Result -> f Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> [Type] -> [(a, SubExp)] -> f Result
clauses (a
j a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) [Type]
ts [(a, SubExp)]
js_to_es
    clauses a
_ [] [(a, SubExp)]
_ =
      Result -> f Result
forall (m :: * -> *) a. Monad m => a -> m a
return []
internaliseExp String
_ (E.Constr Name
_ [ExpBase Info VName]
_ (Info PatternType
t) SrcLoc
loc) =
  String -> InternaliseM Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
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
  Result
ses <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp (String
desc String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_scrutinee") ExpBase Info VName
e
  Result
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
_, Maybe (NonEmpty (CaseBase Info VName))
Nothing) -> do
        (SubExp
_, Result
pertinent) <- Pattern -> Result -> InternaliseM (SubExp, Result)
generateCond Pattern
pCase Result
ses
        Pattern
-> Result
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM Result)
-> InternaliseM Result
forall a.
Pattern
-> Result
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat' Pattern
pCase Result
pertinent ExpBase Info VName
eCase (String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc)
      (CaseBase Info VName
c, Just NonEmpty (CaseBase Info VName)
cs') -> do
        let CasePat Pattern
pLast ExpBase Info VName
eLast SrcLoc
_ = NonEmpty (CaseBase Info VName) -> CaseBase Info VName
forall a. NonEmpty a -> a
NE.last NonEmpty (CaseBase Info VName)
cs'
        Body
bFalse <- do
          (SubExp
_, Result
pertinent) <- Pattern -> Result -> InternaliseM (SubExp, Result)
generateCond Pattern
pLast Result
ses
          Body
eLast' <- Pattern
-> Result
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM Body)
-> InternaliseM Body
forall a.
Pattern
-> Result
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat' Pattern
pLast Result
pertinent ExpBase Info VName
eLast 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
$ Result -> CaseBase Info VName -> Body -> InternaliseM (ExpT SOACS)
generateCaseIf Result
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 Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (ExpT SOACS -> InternaliseM Result)
-> InternaliseM (ExpT SOACS) -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Result -> CaseBase Info VName -> Body -> InternaliseM (ExpT SOACS)
generateCaseIf Result
ses CaseBase Info VName
c Body
bFalse
  StructType -> Closure -> Result -> InternaliseM ()
bindExtSizes (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret) Closure
retext Result
res
  Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
res

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

internaliseExp String
_ (E.Literal PrimValue
v SrcLoc
_) =
  Result -> InternaliseM Result
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)) ->
      Result -> InternaliseM Result
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)) ->
      Result -> InternaliseM Result
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)) ->
      Result -> InternaliseM Result
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 Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
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)) ->
      Result -> InternaliseM Result
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 Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
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
  Result
ses <-
    String -> Exp (Lore InternaliseM) -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc
      (ExpT SOACS -> InternaliseM Result)
-> InternaliseM (ExpT SOACS) -> InternaliseM Result
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 -> Result -> InternaliseM ()
bindExtSizes (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret) Closure
retext Result
ses
  Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
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 Result
internalise <- QualName VName
-> [ExpBase Info VName]
-> SrcLoc
-> Maybe (String -> InternaliseM Result)
isOverloadedFunction QualName VName
op [ExpBase Info VName
xe, ExpBase Info VName
ye] SrcLoc
loc =
    String -> InternaliseM Result
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 Result
internaliseExp String
desc (ExpBase Info VName -> InternaliseM Result)
-> ExpBase Info VName -> InternaliseM Result
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 -> Result -> Result
forall a. Int -> [a] -> [a]
take Int
n (Result -> Result) -> (Result -> Result) -> Result -> Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Result -> Result
forall a. Int -> [a] -> [a]
drop Int
i' (Result -> Result) -> InternaliseM Result -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc ExpBase Info VName
e
internaliseExp String
_ e :: ExpBase Info VName
e@E.Lambda {} =
  String -> InternaliseM Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
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 Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
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 Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
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 Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
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 Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
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 Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
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 Result
internaliseArg String
desc (ExpBase Info VName
arg, Maybe VName
argdim) = do
  Result
arg' <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc ExpBase Info VName
arg
  case (Result
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
    (Result, Maybe VName)
_ -> () -> InternaliseM ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
  Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
arg'

generateCond :: E.Pattern -> [I.SubExp] -> InternaliseM (I.SubExp, [I.SubExp])
generateCond :: Pattern -> Result -> InternaliseM (SubExp, Result)
generateCond Pattern
orig_p Result
orig_ses = do
  (Result
cmps, Result
pertinent, Result
_) <- Pattern -> Result -> InternaliseM (Result, Result, Result)
compares Pattern
orig_p Result
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
=<< Result -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => Result -> m (Exp (Lore m))
eAll Result
cmps
  (SubExp, Result) -> InternaliseM (SubExp, Result)
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
cmp, Result
pertinent)
  where
    -- Literals are always primitive values.
    compares :: Pattern -> Result -> InternaliseM (Result, Result, Result)
compares (E.PatternLit ExpBase Info VName
e Info PatternType
_ SrcLoc
_) (SubExp
se : Result
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
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
      (Result, Result, Result) -> InternaliseM (Result, Result, Result)
forall (m :: * -> *) a. Monad m => a -> m a
return ([SubExp
cmp], [SubExp
se], Result
ses)
    compares (E.PatternConstr Name
c (Info (E.Scalar (E.Sum Map Name [PatternType]
fs))) [Pattern]
pats SrcLoc
_) (SubExp
se : Result
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 (Result
payload_ses, Result
ses') = Int -> Result -> (Result, Result)
forall a. Int -> [a] -> ([a], [a])
splitAt ([DeclExtType] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DeclExtType]
payload_ts) Result
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
          (Result
cmps, Result
pertinent, Result
_) <- [Pattern] -> Result -> InternaliseM (Result, Result, Result)
comparesMany [Pattern]
pats (Result -> InternaliseM (Result, Result, Result))
-> Result -> InternaliseM (Result, Result, Result)
forall a b. (a -> b) -> a -> b
$ (Int -> SubExp) -> [Int] -> Result
forall a b. (a -> b) -> [a] -> [b]
map (Result
payload_ses Result -> Int -> SubExp
forall a. [a] -> Int -> a
!!) [Int]
payload_is
          (Result, Result, Result) -> InternaliseM (Result, Result, Result)
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
cmp SubExp -> Result -> Result
forall a. a -> [a] -> [a]
: Result
cmps, Result
pertinent, Result
ses')
        Maybe (Int, [Int])
Nothing ->
          String -> InternaliseM (Result, Result, Result)
forall a. HasCallStack => String -> a
error String
"generateCond: missing constructor"
    compares (E.PatternConstr Name
_ (Info PatternType
t) [Pattern]
_ SrcLoc
_) Result
_ =
      String -> InternaliseM (Result, Result, Result)
forall a. HasCallStack => String -> a
error (String -> InternaliseM (Result, Result, Result))
-> String -> InternaliseM (Result, Result, Result)
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) Result
ses =
      Pattern -> Result -> InternaliseM (Result, Result, Result)
compares (Info PatternType -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
f PatternType -> SrcLoc -> PatternBase f vn
E.Wildcard Info PatternType
t SrcLoc
loc) Result
ses
    compares (E.Wildcard (Info PatternType
t) SrcLoc
_) Result
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 (Result
id_ses, Result
rest_ses) = Int -> Result -> (Result, Result)
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n Result
ses
      (Result, Result, Result) -> InternaliseM (Result, Result, Result)
forall (m :: * -> *) a. Monad m => a -> m a
return ([], Result
id_ses, Result
rest_ses)
    compares (E.PatternParens Pattern
pat SrcLoc
_) Result
ses =
      Pattern -> Result -> InternaliseM (Result, Result, Result)
compares Pattern
pat Result
ses
    compares (E.TuplePattern [Pattern]
pats SrcLoc
_) Result
ses =
      [Pattern] -> Result -> InternaliseM (Result, Result, Result)
comparesMany [Pattern]
pats Result
ses
    compares (E.RecordPattern [(Name, Pattern)]
fs SrcLoc
_) Result
ses =
      [Pattern] -> Result -> InternaliseM (Result, Result, Result)
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) Result
ses
    compares (E.PatternAscription Pattern
pat TypeDeclBase Info VName
_ SrcLoc
_) Result
ses =
      Pattern -> Result -> InternaliseM (Result, Result, Result)
compares Pattern
pat Result
ses
    compares Pattern
pat [] =
      String -> InternaliseM (Result, Result, Result)
forall a. HasCallStack => String -> a
error (String -> InternaliseM (Result, Result, Result))
-> String -> InternaliseM (Result, Result, Result)
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] -> Result -> InternaliseM (Result, Result, Result)
comparesMany [] Result
ses = (Result, Result, Result) -> InternaliseM (Result, Result, Result)
forall (m :: * -> *) a. Monad m => a -> m a
return ([], [], Result
ses)
    comparesMany (Pattern
pat : [Pattern]
pats) Result
ses = do
      (Result
cmps1, Result
pertinent1, Result
ses') <- Pattern -> Result -> InternaliseM (Result, Result, Result)
compares Pattern
pat Result
ses
      (Result
cmps2, Result
pertinent2, Result
ses'') <- [Pattern] -> Result -> InternaliseM (Result, Result, Result)
comparesMany [Pattern]
pats Result
ses'
      (Result, Result, Result) -> InternaliseM (Result, Result, Result)
forall (m :: * -> *) a. Monad m => a -> m a
return
        ( Result
cmps1 Result -> Result -> Result
forall a. Semigroup a => a -> a -> a
<> Result
cmps2,
          Result
pertinent1 Result -> Result -> Result
forall a. Semigroup a => a -> a -> a
<> Result
pertinent2,
          Result
ses''
        )

generateCaseIf :: [I.SubExp] -> Case -> I.Body -> InternaliseM I.Exp
generateCaseIf :: Result -> CaseBase Info VName -> Body -> InternaliseM (ExpT SOACS)
generateCaseIf Result
ses (CasePat Pattern
p ExpBase Info VName
eCase SrcLoc
_) Body
bFail = do
  (SubExp
cond, Result
pertinent) <- Pattern -> Result -> InternaliseM (SubExp, Result)
generateCond Pattern
p Result
ses
  Body
eCase' <- Pattern
-> Result
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM Body)
-> InternaliseM Body
forall a.
Pattern
-> Result
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat' Pattern
p Result
pertinent ExpBase Info VName
eCase 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 ->
  (E.Exp -> InternaliseM a) ->
  InternaliseM a
internalisePat :: String
-> Pattern
-> ExpBase Info VName
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat String
desc Pattern
p ExpBase Info VName
e ExpBase Info VName
body ExpBase Info VName -> InternaliseM a
m = do
  Result
ses <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc' ExpBase Info VName
e
  Pattern
-> Result
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
forall a.
Pattern
-> Result
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat' Pattern
p Result
ses ExpBase Info VName
body 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 ->
  (E.Exp -> InternaliseM a) ->
  InternaliseM a
internalisePat' :: Pattern
-> Result
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat' Pattern
p Result
ses ExpBase Info VName
body ExpBase Info VName -> InternaliseM a
m = do
  [Type]
ses_ts <- (SubExp -> InternaliseM Type) -> Result -> 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 Result
ses
  Pattern -> [Type] -> (Closure -> InternaliseM a) -> InternaliseM a
forall a.
Pattern -> [Type] -> (Closure -> InternaliseM a) -> InternaliseM a
stmPattern Pattern
p [Type]
ses_ts ((Closure -> InternaliseM a) -> InternaliseM a)
-> (Closure -> InternaliseM a) -> InternaliseM a
forall a b. (a -> b) -> a -> b
$ \Closure
pat_names -> do
    [(VName, SubExp)]
-> ((VName, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Closure -> Result -> [(VName, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip Closure
pat_names Result
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
-> Result
-> [DimIndexBase Info VName]
-> InternaliseM ([DimIndex SubExp], Certificates)
internaliseSlice SrcLoc
loc Result
dims [DimIndexBase Info VName]
idxs = do
  ([DimIndex SubExp]
idxs', Result
oks, [[ErrorMsgPart SubExp]]
parts) <- [(DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])]
-> ([DimIndex SubExp], Result, [[ErrorMsgPart SubExp]])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 ([(DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])]
 -> ([DimIndex SubExp], Result, [[ErrorMsgPart SubExp]]))
-> InternaliseM [(DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])]
-> InternaliseM
     ([DimIndex SubExp], Result, [[ErrorMsgPart SubExp]])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp
 -> DimIndexBase Info VName
 -> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp]))
-> Result
-> [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 Result
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
=<< Result -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => Result -> m (Exp (Lore m))
eAll Result
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) -> Result -> [ErrorMsgPart SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt32 (Result -> [ErrorMsgPart SubExp])
-> Result -> [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ Int -> Result -> Result
forall a. Int -> [a] -> [a]
take ([DimIndexBase Info VName] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimIndexBase Info VName]
idxs) Result
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
            (Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody [SubExp
w_minus_1])
            (Result -> Body
forall lore. Bindable lore => Result -> 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
            (Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody [SubExp
negone])
            (Result -> Body
forall lore. Bindable lore => Result -> 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
=<< Result -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => Result -> 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
=<< Result -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => Result -> 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
        (Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody [SubExp
backwards_ok])
        (Result -> Body
forall lore. Bindable lore => Result -> 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 -> Result -> Closure -> InternaliseM (SOAC SOACS))
-> (ExpBase Info VName, ExpBase Info VName, ExpBase Info VName,
    SrcLoc)
-> InternaliseM Result
internaliseScanOrReduce String
desc String
what SubExp -> Lambda -> Result -> 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
what String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_arr") ExpBase Info VName
arr
  Result
nes <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp (String
what String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_ne") ExpBase Info VName
ne
  Result
nes' <- [(SubExp, VName)]
-> ((SubExp, VName) -> InternaliseM SubExp) -> InternaliseM Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (Result -> Closure -> [(SubExp, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip Result
nes Closure
arrs) (((SubExp, VName) -> InternaliseM SubExp) -> InternaliseM Result)
-> ((SubExp, VName) -> InternaliseM SubExp) -> InternaliseM Result
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
what String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_ne_right_shape")
      SubExp
ne'
  [Type]
nests <- (SubExp -> InternaliseM Type) -> Result -> 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 Result
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 Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (ExpT SOACS -> InternaliseM Result)
-> (SOAC SOACS -> ExpT SOACS) -> SOAC SOACS -> InternaliseM Result
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 Result)
-> InternaliseM (SOAC SOACS) -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SubExp -> Lambda -> Result -> Closure -> InternaliseM (SOAC SOACS)
f SubExp
w Lambda
lam' Result
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 Result
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
  Result
ne' <- String -> ExpBase Info VName -> InternaliseM Result
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
  Result
ne_shp <- [(SubExp, VName)]
-> ((SubExp, VName) -> InternaliseM SubExp) -> InternaliseM Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (Result -> Closure -> [(SubExp, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip Result
ne' Closure
hist') (((SubExp, VName) -> InternaliseM SubExp) -> InternaliseM Result)
-> ((SubExp, VName) -> InternaliseM SubExp) -> InternaliseM Result
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) -> Result -> 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 Result
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 -> Result -> Body
forall lore. Bindable lore => Stms lore -> Result -> Body lore
mkBody Stms SOACS
forall a. Monoid a => a
mempty (Result -> Body) -> Result -> Body
forall a b. (a -> b) -> a -> b
$ (Param Type -> SubExp) -> [Param Type] -> Result
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 Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
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 -> Result -> Lambda -> HistOp SOACS
forall lore.
SubExp -> SubExp -> Closure -> Result -> Lambda lore -> HistOp lore
HistOp SubExp
w_hist SubExp
rf' Closure
hist' Result
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 Result
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 -> Result -> InternaliseM Lambda
internaliseStreamMapLambda InternaliseLambda
internaliseLambda ExpBase Info VName
lam (Result -> InternaliseM Lambda) -> Result -> InternaliseM Lambda
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> Closure -> Result
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 -> Result -> StreamForm SOACS
forall lore.
StreamOrd
-> Commutativity -> Lambda lore -> Result -> 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 -> Result -> Body
forall lore. Bindable lore => Stms lore -> Result -> Body lore
mkBody Stms SOACS
forall a. Monoid a => a
mempty []) []) []
  String -> Exp (Lore InternaliseM) -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
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 Result
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 -> Result -> 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) (Result -> BasicOp) -> Result -> BasicOp
forall a b. (a -> b) -> a -> b
$
          Type -> Result
forall u. TypeBase Shape u -> Result
I.arrayDims (Type -> Result) -> Type -> Result
forall a b. (a -> b) -> a -> b
$ Param Type -> Type
forall dec. Typed dec => Param dec -> Type
I.paramType Param Type
p
  Result
nes <- Body -> InternaliseM Result
forall (m :: * -> *). MonadBinder m => Body (Lore m) -> m Result
bodyBind (Body -> InternaliseM Result)
-> InternaliseM Body -> InternaliseM Result
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) -> Result -> 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 Result
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 (Result -> 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 (Result -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Result
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
        Result
lam_res <- Body (Lore InternaliseM) -> InternaliseM Result
forall (m :: * -> *). MonadBinder m => Body (Lore m) -> m Result
bodyBind Body (Lore InternaliseM)
Body
lam_body
        Result
lam_res' <-
          ErrorMsg SubExp
-> SrcLoc -> Closure -> [Type] -> Result -> InternaliseM Result
forall u.
Typed (TypeBase Shape u) =>
ErrorMsg SubExp
-> SrcLoc
-> Closure
-> [TypeBase Shape u]
-> Result
-> InternaliseM Result
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')
            Result
lam_res
        Result
new_lam_res <-
          Lambda (Lore InternaliseM)
-> [InternaliseM (Exp (Lore InternaliseM))] -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
Lambda (Lore m) -> [m (Exp (Lore m))] -> m Result
eLambda Lambda (Lore InternaliseM)
Lambda
lam0' ([InternaliseM (Exp (Lore InternaliseM))] -> InternaliseM Result)
-> [InternaliseM (Exp (Lore InternaliseM))] -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
            (SubExp -> InternaliseM (ExpT SOACS))
-> Result -> [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 (Result -> [InternaliseM (ExpT SOACS)])
-> Result -> [InternaliseM (ExpT SOACS)]
forall a b. (a -> b) -> a -> b
$
              (Param Type -> SubExp) -> [Param Type] -> Result
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 Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ Result
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
$ Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody Result
new_lam_res

  let form :: StreamForm SOACS
form = StreamOrd -> Commutativity -> Lambda -> Result -> StreamForm SOACS
forall lore.
StreamOrd
-> Commutativity -> Lambda lore -> Result -> StreamForm lore
I.Parallel StreamOrd
o Commutativity
comm Lambda
lam0' Result
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 Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
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
  Result
vs <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc ExpBase Info VName
e
  case Result
vs of
    [SubExp
se] -> SubExp -> InternaliseM SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return SubExp
se
    Result
_ -> 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'
    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) -> Result -> InternaliseM Closure
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM VName
asIdent (Result -> InternaliseM Closure)
-> InternaliseM Result -> InternaliseM Closure
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> ExpBase Info VName -> InternaliseM Result
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 Result
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 Result
forall (m :: * -> *).
MonadBinder m =>
String -> [Exp (Lore m)] -> m Result
letSubExps String
s ([ExpT SOACS] -> InternaliseM Result)
-> InternaliseM [ExpT SOACS] -> InternaliseM Result
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 Result
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Plus SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result -> InternaliseM Result
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
    String -> BinOp -> SubExp -> SubExp -> InternaliseM Result
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 Result -> InternaliseM Result
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
    String -> BinOp -> SubExp -> SubExp -> InternaliseM Result
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 Result
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 Result
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 Result -> InternaliseM Result
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonnegative SrcLoc
loc IntType
t SubExp
y (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
    String -> BinOp -> SubExp -> SubExp -> InternaliseM Result
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 Result
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 Result -> InternaliseM Result
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
    String -> BinOp -> SubExp -> SubExp -> InternaliseM Result
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 Result -> InternaliseM Result
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
    String -> BinOp -> SubExp -> SubExp -> InternaliseM Result
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 Result
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 Result -> InternaliseM Result
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
    String -> BinOp -> SubExp -> SubExp -> InternaliseM Result
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 Result -> InternaliseM Result
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
    String -> BinOp -> SubExp -> SubExp -> InternaliseM Result
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 Result -> InternaliseM Result
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
    String -> BinOp -> SubExp -> SubExp -> InternaliseM Result
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 Result -> InternaliseM Result
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
    String -> BinOp -> SubExp -> SubExp -> InternaliseM Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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
desc String -> 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 -> Result) -> InternaliseM SubExp -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> Result
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InternaliseM SubExp -> InternaliseM Result)
-> InternaliseM SubExp -> InternaliseM Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
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 Result
simpleBinOp String
desc BinOp
bop SubExp
x SubExp
y =
  String -> Exp (Lore InternaliseM) -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
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 Result
simpleCmpOp String
desc CmpOp
op SubExp
x SubExp
y =
  String -> Exp (Lore InternaliseM) -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
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, [Type]))
-> InternaliseM ([Param Type], Body, [Type])
forall a.
[Pattern]
-> [Type] -> ([LParam] -> InternaliseM a) -> InternaliseM a
bindingLambdaParams [Pattern]
params [Type]
rowtypes (([LParam] -> InternaliseM ([Param Type], Body, [Type]))
 -> InternaliseM ([Param Type], Body, [Type]))
-> ([LParam] -> InternaliseM ([Param Type], Body, [Type]))
-> InternaliseM ([Param Type], Body, [Type])
forall a b. (a -> b) -> a -> b
$ \[LParam]
params' -> do
    Body
body' <- ExpBase Info VName -> InternaliseM Body
internaliseBody ExpBase Info VName
body
    [Type]
rettype' <- StructType -> InternaliseM [Type]
internaliseLambdaReturnType StructType
rettype
    ([Param Type], Body, [Type])
-> InternaliseM ([Param Type], Body, [Type])
forall (m :: * -> *) a. Monad m => a -> m a
return ([Param Type]
[LParam]
params', Body
body', [Type]
rettype')
internaliseLambda ExpBase Info VName
e [Type]
_ = String -> InternaliseM ([Param Type], Body, [Type])
forall a. HasCallStack => String -> a
error (String -> InternaliseM ([Param Type], Body, [Type]))
-> String -> InternaliseM ([Param Type], Body, [Type])
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 Result)
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 Result)]
handlers =
        [ [ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM Result)
forall a.
(Eq a, IsString a) =>
[ExpBase Info VName] -> a -> Maybe (String -> InternaliseM Result)
handleSign,
          [ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM Result)
forall (f :: * -> *).
Applicative f =>
[ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM (f SubExp))
handleIntrinsicOps,
          [ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM Result)
handleOps,
          [ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM Result)
handleSOACs,
          [ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM Result)
handleRest
        ]
  [Maybe (String -> InternaliseM Result)]
-> Maybe (String -> InternaliseM Result)
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, MonadPlus m) =>
t (m a) -> m a
msum [[ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM Result)
h [ExpBase Info VName]
args (String -> Maybe (String -> InternaliseM Result))
-> String -> Maybe (String -> InternaliseM Result)
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 Result)
h <- [[ExpBase Info VName]
 -> String -> Maybe (String -> InternaliseM Result)]
handlers]
  where
    handleSign :: [ExpBase Info VName] -> a -> Maybe (String -> InternaliseM Result)
handleSign [ExpBase Info VName
x] a
"sign_i8" = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ IntType -> ExpBase Info VName -> String -> InternaliseM Result
toSigned IntType
I.Int8 ExpBase Info VName
x
    handleSign [ExpBase Info VName
x] a
"sign_i16" = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ IntType -> ExpBase Info VName -> String -> InternaliseM Result
toSigned IntType
I.Int16 ExpBase Info VName
x
    handleSign [ExpBase Info VName
x] a
"sign_i32" = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ IntType -> ExpBase Info VName -> String -> InternaliseM Result
toSigned IntType
I.Int32 ExpBase Info VName
x
    handleSign [ExpBase Info VName
x] a
"sign_i64" = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ IntType -> ExpBase Info VName -> String -> InternaliseM Result
toSigned IntType
I.Int64 ExpBase Info VName
x
    handleSign [ExpBase Info VName
x] a
"unsign_i8" = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ IntType -> ExpBase Info VName -> String -> InternaliseM Result
toUnsigned IntType
I.Int8 ExpBase Info VName
x
    handleSign [ExpBase Info VName
x] a
"unsign_i16" = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ IntType -> ExpBase Info VName -> String -> InternaliseM Result
toUnsigned IntType
I.Int16 ExpBase Info VName
x
    handleSign [ExpBase Info VName
x] a
"unsign_i32" = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ IntType -> ExpBase Info VName -> String -> InternaliseM Result
toUnsigned IntType
I.Int32 ExpBase Info VName
x
    handleSign [ExpBase Info VName
x] a
"unsign_i64" = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ IntType -> ExpBase Info VName -> String -> InternaliseM Result
toUnsigned IntType
I.Int64 ExpBase Info VName
x
    handleSign [ExpBase Info VName]
_ a
_ = Maybe (String -> InternaliseM Result)
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 Result)
handleOps [ExpBase Info VName
x, ExpBase Info VName
y] String
"&&" = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc (ExpBase Info VName -> InternaliseM Result)
-> ExpBase Info VName -> InternaliseM Result
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 Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc (ExpBase Info VName -> InternaliseM Result)
-> ExpBase Info VName -> InternaliseM Result
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 Result
cmp_f <- String -> Maybe (String -> SubExp -> InternaliseM Result)
forall a (m :: * -> *).
(IsString a, MonadBinder m, Eq a) =>
a -> Maybe (String -> SubExp -> m Result)
isEqlOp String
op = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ \String
desc -> do
        Result
xe' <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
"x" ExpBase Info VName
xe
        Result
ye' <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
"y" ExpBase Info VName
ye
        Result
rs <- (SubExp -> SubExp -> InternaliseM SubExp)
-> Result -> Result -> InternaliseM Result
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) Result
xe' Result
ye'
        String -> SubExp -> InternaliseM Result
cmp_f String
desc (SubExp -> InternaliseM Result)
-> InternaliseM SubExp -> InternaliseM Result
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
=<< Result -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => Result -> m (Exp (Lore m))
eAll Result
rs
      where
        isEqlOp :: a -> Maybe (String -> SubExp -> m Result)
isEqlOp a
"!=" = (String -> SubExp -> m Result)
-> Maybe (String -> SubExp -> m Result)
forall a. a -> Maybe a
Just ((String -> SubExp -> m Result)
 -> Maybe (String -> SubExp -> m Result))
-> (String -> SubExp -> m Result)
-> Maybe (String -> SubExp -> m Result)
forall a b. (a -> b) -> a -> b
$ \String
desc SubExp
eq ->
          String -> Exp (Lore m) -> m Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore m) -> m Result) -> Exp (Lore m) -> m Result
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 Result)
-> Maybe (String -> SubExp -> m Result)
forall a. a -> Maybe a
Just ((String -> SubExp -> m Result)
 -> Maybe (String -> SubExp -> m Result))
-> (String -> SubExp -> m Result)
-> Maybe (String -> SubExp -> m Result)
forall a b. (a -> b) -> a -> b
$ \String
_ SubExp
eq ->
          Result -> m Result
forall (m :: * -> *) a. Monad m => a -> m a
return [SubExp
eq]
        isEqlOp a
_ = Maybe (String -> SubExp -> m Result)
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 :: Result
x_dims = Type -> Result
forall u. TypeBase Shape u -> Result
I.arrayDims Type
x_t
                  y_dims :: Result
y_dims = Type -> Result
forall u. TypeBase Shape u -> Result
I.arrayDims Type
y_t
              Result
dims_match <- [(SubExp, SubExp)] -> ((SubExp, SubExp) -> m SubExp) -> m Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (Result -> Result -> [(SubExp, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip Result
x_dims Result
y_dims) (((SubExp, SubExp) -> m SubExp) -> m Result)
-> ((SubExp, SubExp) -> m SubExp) -> m Result
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
=<< Result -> m (Exp (Lore m))
forall (m :: * -> *). MonadBinder m => Result -> m (Exp (Lore m))
eAll Result
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
-> Result
-> BinderT
     (Lore m)
     (State VNameSource)
     (Exp (Lore (BinderT (Lore m) (State VNameSource))))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> SubExp -> Result -> 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)) Result
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
-> BinderT
     (Lore m)
     (State VNameSource)
     (Lambda (Lore (BinderT (Lore m) (State VNameSource))))
forall (m :: * -> *).
(MonadBinder m, Bindable (Lore m)) =>
CmpOp -> m (Lambda (Lore m))
cmpOpLambda (CmpOp
 -> BinderT
      (Lore m)
      (State VNameSource)
      (Lambda (Lore (BinderT (Lore m) (State VNameSource)))))
-> CmpOp
-> BinderT
     (Lore m)
     (State VNameSource)
     (Lambda (Lore (BinderT (Lore m) (State VNameSource))))
forall a b. (a -> b) -> a -> b
$ PrimType -> CmpOp
I.CmpEq (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) -> Result -> Reduce (Lore m)
forall lore. Commutativity -> Lambda lore -> Result -> 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
$ Result -> Body (Lore m)
forall lore. Bindable lore => Result -> 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 (Result -> Body (Lore m)
forall lore. Bindable lore => Result -> 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
name String -> 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 Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
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 Result
internaliseBinOp SrcLoc
loc String
desc BinOp
bop SubExp
x' SubExp
y' PrimType
t1 PrimType
t2
            (PatternType, PatternType)
_ -> String -> InternaliseM Result
forall a. HasCallStack => String -> a
error String
"Futhark.Internalise.internaliseExp: non-primitive type in BinOp."
    handleOps [ExpBase Info VName]
_ String
_ = Maybe (String -> InternaliseM Result)
forall a. Maybe a
Nothing

    handleSOACs :: [ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM Result)
handleSOACs [TupLit [ExpBase Info VName
lam, ExpBase Info VName
arr] SrcLoc
_] String
"map" = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
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 -> Result -> InternaliseM Lambda
internaliseMapLambda InternaliseLambda
internaliseLambda ExpBase Info VName
lam (Result -> InternaliseM Lambda) -> Result -> InternaliseM Lambda
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> Closure -> Result
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 Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
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
fromInt32 ExpBase Info VName
k
      (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
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 -> Result -> InternaliseM Lambda
internalisePartitionLambda InternaliseLambda
internaliseLambda Int
k' ExpBase Info VName
lam (Result -> InternaliseM Lambda) -> Result -> InternaliseM Lambda
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> Closure -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var Closure
arrs
        (Result -> Result -> Result) -> (Result, Result) -> Result
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Result -> Result -> Result
forall a. [a] -> [a] -> [a]
(++) ((Result, Result) -> Result)
-> InternaliseM (Result, Result) -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Lambda -> Closure -> InternaliseM (Result, Result)
partitionWithSOACS Int
k' Lambda
lam' Closure
arrs
      where
        fromInt32 :: ExpBase Info vn -> Maybe Int32
fromInt32 (Literal (SignedValue (Int32Value Int32
k')) SrcLoc
_) = Int32 -> Maybe Int32
forall a. a -> Maybe a
Just Int32
k'
        fromInt32 (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'
        fromInt32 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 Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> String
-> (SubExp
    -> Lambda -> Result -> Closure -> InternaliseM (SOAC SOACS))
-> (ExpBase Info VName, ExpBase Info VName, ExpBase Info VName,
    SrcLoc)
-> InternaliseM Result
internaliseScanOrReduce String
desc String
"reduce" SubExp -> Lambda -> Result -> Closure -> InternaliseM (SOAC SOACS)
forall (f :: * -> *) lore.
(Bindable lore, MonadFreshNames f) =>
SubExp -> Lambda lore -> Result -> Closure -> f (SOAC lore)
reduce (ExpBase Info VName
lam, ExpBase Info VName
ne, ExpBase Info VName
arr, SrcLoc
loc)
      where
        reduce :: SubExp -> Lambda lore -> Result -> Closure -> f (SOAC lore)
reduce SubExp
w Lambda lore
red_lam Result
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 -> Result -> Reduce lore
forall lore. Commutativity -> Lambda lore -> Result -> Reduce lore
Reduce Commutativity
Noncommutative Lambda lore
red_lam Result
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 Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> String
-> (SubExp
    -> Lambda -> Result -> Closure -> InternaliseM (SOAC SOACS))
-> (ExpBase Info VName, ExpBase Info VName, ExpBase Info VName,
    SrcLoc)
-> InternaliseM Result
internaliseScanOrReduce String
desc String
"reduce" SubExp -> Lambda -> Result -> Closure -> InternaliseM (SOAC SOACS)
forall (f :: * -> *) lore.
(Bindable lore, MonadFreshNames f) =>
SubExp -> Lambda lore -> Result -> Closure -> f (SOAC lore)
reduce (ExpBase Info VName
lam, ExpBase Info VName
ne, ExpBase Info VName
arr, SrcLoc
loc)
      where
        reduce :: SubExp -> Lambda lore -> Result -> Closure -> f (SOAC lore)
reduce SubExp
w Lambda lore
red_lam Result
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 -> Result -> Reduce lore
forall lore. Commutativity -> Lambda lore -> Result -> Reduce lore
Reduce Commutativity
Commutative Lambda lore
red_lam Result
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 Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> String
-> (SubExp
    -> Lambda -> Result -> Closure -> InternaliseM (SOAC SOACS))
-> (ExpBase Info VName, ExpBase Info VName, ExpBase Info VName,
    SrcLoc)
-> InternaliseM Result
internaliseScanOrReduce String
desc String
"scan" SubExp -> Lambda -> Result -> Closure -> InternaliseM (SOAC SOACS)
forall (f :: * -> *) lore.
(Bindable lore, MonadFreshNames f) =>
SubExp -> Lambda lore -> Result -> Closure -> f (SOAC lore)
reduce (ExpBase Info VName
lam, ExpBase Info VName
ne, ExpBase Info VName
arr, SrcLoc
loc)
      where
        reduce :: SubExp -> Lambda lore -> Result -> Closure -> f (SOAC lore)
reduce SubExp
w Lambda lore
scan_lam Result
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 -> Result -> Scan lore
forall lore. Lambda lore -> Result -> Scan lore
Scan Lambda lore
scan_lam Result
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 Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> StreamOrd
-> Commutativity
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> InternaliseM Result
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 Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> StreamOrd
-> Commutativity
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> InternaliseM Result
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 Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> StreamOrd
-> ExpBase Info VName
-> ExpBase Info VName
-> InternaliseM Result
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 Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> StreamOrd
-> ExpBase Info VName
-> ExpBase Info VName
-> InternaliseM Result
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 Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
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 Result
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 Result)
forall a. Maybe a
Nothing

    handleRest :: [ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM Result)
handleRest [ExpBase Info VName
x] String
"!" = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> String -> InternaliseM Result
complementF ExpBase Info VName
x
    handleRest [ExpBase Info VName
x] String
"opaque" = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      (SubExp -> InternaliseM SubExp) -> Result -> InternaliseM Result
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) (Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> ExpBase Info VName -> InternaliseM Result
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 Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> String
-> InternaliseM Result
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 Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
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 Result -> InternaliseM Result
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
dim_ok_cert (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
        Closure -> (VName -> InternaliseM SubExp) -> InternaliseM Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM Closure
arrs ((VName -> InternaliseM SubExp) -> InternaliseM Result)
-> (VName -> InternaliseM SubExp) -> InternaliseM Result
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 Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
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 Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM Closure
arrs ((VName -> InternaliseM SubExp) -> InternaliseM Result)
-> (VName -> InternaliseM SubExp) -> InternaliseM Result
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 Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
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 -> Result -> 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
          (Result -> InternaliseM SubExp)
-> InternaliseM Result -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Closure -> InternaliseM SubExp)
-> [Closure] -> InternaliseM Result
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 Result
forall (m :: * -> *).
MonadBinder m =>
String -> [Exp (Lore m)] -> m Result
letSubExps String
desc ([Exp (Lore InternaliseM)] -> InternaliseM Result)
-> [Exp (Lore InternaliseM)] -> InternaliseM Result
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 Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
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 Result
internaliseOperation String
desc ExpBase Info VName
e ((VName -> InternaliseM BasicOp) -> InternaliseM Result)
-> (VName -> InternaliseM BasicOp) -> InternaliseM Result
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 :: Result
offsets = SubExp
offset' SubExp -> Result -> Result
forall a. a -> [a] -> [a]
: Int -> SubExp -> Result
forall a. Int -> a -> [a]
replicate (Int
r Int -> 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
$ Result -> VName -> BasicOp
I.Rotate Result
offsets VName
v
    handleRest [ExpBase Info VName
e] String
"transpose" = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> ExpBase Info VName
-> (VName -> InternaliseM BasicOp)
-> InternaliseM Result
internaliseOperation String
desc ExpBase Info VName
e ((VName -> InternaliseM BasicOp) -> InternaliseM Result)
-> (VName -> InternaliseM BasicOp) -> InternaliseM Result
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
r Int -> 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 Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      Result -> Result -> Result
forall a. [a] -> [a] -> [a]
(++) (Result -> Result -> Result)
-> InternaliseM Result -> InternaliseM (Result -> Result)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> ExpBase Info VName -> InternaliseM Result
internaliseExp (String
desc String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_zip_x") ExpBase Info VName
x
        InternaliseM (Result -> Result)
-> InternaliseM Result -> InternaliseM Result
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> String -> ExpBase Info VName -> InternaliseM Result
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 Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ (String -> ExpBase Info VName -> InternaliseM Result)
-> ExpBase Info VName -> String -> InternaliseM Result
forall a b c. (a -> b -> c) -> b -> a -> c
flip String -> ExpBase Info VName -> InternaliseM Result
internaliseExp ExpBase Info VName
x
    handleRest [ExpBase Info VName
x] String
"trace" = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ (String -> ExpBase Info VName -> InternaliseM Result)
-> ExpBase Info VName -> String -> InternaliseM Result
forall a b c. (a -> b -> c) -> b -> a -> c
flip String -> ExpBase Info VName -> InternaliseM Result
internaliseExp ExpBase Info VName
x
    handleRest [ExpBase Info VName
x] String
"break" = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ (String -> ExpBase Info VName -> InternaliseM Result)
-> ExpBase Info VName -> String -> InternaliseM Result
forall a b c. (a -> b -> c) -> b -> a -> c
flip String -> ExpBase Info VName -> InternaliseM Result
internaliseExp ExpBase Info VName
x
    handleRest [ExpBase Info VName]
_ String
_ = Maybe (String -> InternaliseM Result)
forall a. Maybe a
Nothing

    toSigned :: IntType -> ExpBase Info VName -> String -> InternaliseM Result
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 Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
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'
              (Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
1])
              (Result -> Body
forall lore. Bindable lore => Result -> 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 Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
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 Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
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 Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
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 Result
forall a. HasCallStack => String -> a
error String
"Futhark.Internalise: non-numeric type in ToSigned"

    toUnsigned :: IntType -> ExpBase Info VName -> String -> InternaliseM Result
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 Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
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'
              (Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
1])
              (Result -> Body
forall lore. Bindable lore => Result -> 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 Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
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 Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
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 Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
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 Result
forall a. HasCallStack => String -> a
error String
"Futhark.Internalise.internaliseExp: non-numeric type in ToUnsigned"

    complementF :: ExpBase Info VName -> String -> InternaliseM Result
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 Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
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 Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
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 Result
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 Result
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
          Result
results <- Closure -> (VName -> InternaliseM SubExp) -> InternaliseM Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM Closure
outs ((VName -> InternaliseM SubExp) -> InternaliseM Result)
-> (VName -> InternaliseM SubExp) -> InternaliseM Result
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
$ Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody Result
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 :: Result
sa_ws = (Type -> SubExp) -> [Type] -> Result
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 Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
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
$ Result -> [Int] -> Closure -> [(SubExp, Int, VName)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 Result
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
-> Result
-> SrcLoc
-> InternaliseM (Result, [ExtType])
funcall String
desc (QualName Closure
_ VName
fname) Result
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) -> Result -> 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 Result
args

  Result
shapeargs <- Closure -> [FParam] -> [Type] -> InternaliseM Result
forall (m :: * -> *).
(HasScope SOACS m, Monad m) =>
Closure -> [FParam] -> [Type] -> m Result
argShapes Closure
shapes [Param DeclType]
[FParam]
fun_params [Type]
argts
  let 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
+ Result -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Result
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
  Result
args' <-
    ErrorMsg SubExp
-> SrcLoc -> Closure -> [Type] -> Result -> InternaliseM Result
forall u.
Typed (TypeBase Shape u) =>
ErrorMsg SubExp
-> SrcLoc
-> Closure
-> [TypeBase Shape u]
-> Result
-> InternaliseM Result
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)
      ((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]
fun_params)
      ((VName -> SubExp) -> Closure -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var Closure
closure Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ Result
shapeargs Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ Result
args)
  [Type]
argts' <- (SubExp -> InternaliseM Type) -> Result -> 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 Result
args'
  case [(SubExp, Type)] -> Maybe [DeclExtType]
rettype_fun ([(SubExp, Type)] -> Maybe [DeclExtType])
-> [(SubExp, Type)] -> Maybe [DeclExtType]
forall a b. (a -> b) -> a -> b
$ Result -> [Type] -> [(SubExp, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip Result
args' [Type]
argts' of
    Maybe [DeclExtType]
Nothing ->
      String -> InternaliseM (Result, [ExtType])
forall a. HasCallStack => String -> a
error (String -> InternaliseM (Result, [ExtType]))
-> String -> InternaliseM (Result, [ExtType])
forall a b. (a -> b) -> a -> b
$
        [String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
          [ String
"Cannot apply ",
            VName -> String
forall a. Pretty a => a -> String
pretty VName
fname,
            String
" to ",
            Int -> String
forall a. Show a => a -> String
show (Result -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Result
args'),
            String
" arguments\n ",
            Result -> String
forall a. Pretty a => a -> String
pretty Result
args',
            String
"\nof types\n ",
            [Type] -> String
forall a. Pretty a => a -> String
pretty [Type]
argts',
            String
"\nFunction has ",
            Int -> String
forall a. Show a => a -> String
show ([Param DeclType] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Param DeclType]
fun_params),
            String
" parameters\n ",
            [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
      Result
ses <-
        Attrs -> InternaliseM Result -> InternaliseM Result
forall (m :: * -> *) a. MonadBinder m => Attrs -> m a -> m a
attributing Attrs
attrs (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
          String -> Exp (Lore InternaliseM) -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
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' (Result -> [Diet] -> [(SubExp, Diet)]
forall a b. [a] -> [b] -> [(a, b)]
zip Result
args' [Diet]
diets) [DeclExtType]
[RetType SOACS]
ts (Safety
safety, SrcLoc
loc, [SrcLoc]
forall a. Monoid a => a
mempty)
      (Result, [ExtType]) -> InternaliseM (Result, [ExtType])
forall (m :: * -> *) a. Monad m => a -> m a
return (Result
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 -> Result -> InternaliseM ()
bindExtSizes StructType
ret Closure
retext Result
ses = do
  [DeclExtType]
ts <- StructType -> InternaliseM [DeclExtType]
internaliseType StructType
ret
  [Type]
ses_ts <- (SubExp -> InternaliseM Type) -> Result -> 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 Result
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
$ (ExtSize -> SubExp -> Map VName SubExp)
-> [ExtSize] -> Result -> [Map VName SubExp]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith ExtSize -> SubExp -> Map VName SubExp
combine' (DeclExtType -> [ExtSize]
forall u. TypeBase ExtShape u -> [ExtSize]
arrayExtDims DeclExtType
t1) (Type -> Result
forall u. TypeBase Shape u -> Result
arrayDims Type
t2)
      combine' :: ExtSize -> 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' ExtSize
_ 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 (Result, Result)
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
$
      (Result -> Body)
-> BinderT SOACS (State VNameSource) Result -> Binder SOACS Body
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody (BinderT SOACS (State VNameSource) Result -> Binder SOACS Body)
-> BinderT SOACS (State VNameSource) Result -> 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) Result
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) Result)
-> ((Param Type, Param Type)
    -> BinderT SOACS (State VNameSource) SubExp)
-> BinderT SOACS (State VNameSource) Result
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 :: Result
nes = Int -> SubExp -> Result
forall a. Int -> a -> [a]
replicate (Closure -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Closure
increments) (SubExp -> Result) -> SubExp -> Result
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 -> Result -> Scan SOACS
forall lore. Lambda lore -> Result -> Scan lore
I.Scan Lambda
add_lam Result
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
$
    (Result -> Body)
-> BinderT SOACS (State VNameSource) Result -> Binder SOACS Body
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody (BinderT SOACS (State VNameSource) Result -> Binder SOACS Body)
-> BinderT SOACS (State VNameSource) Result -> Binder SOACS Body
forall a b. (a -> b) -> a -> b
$
      Closure
-> (VName -> BinderT SOACS (State VNameSource) SubExp)
-> BinderT SOACS (State VNameSource) Result
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) Result)
-> (VName -> BinderT SOACS (State VNameSource) SubExp)
-> BinderT SOACS (State VNameSource) Result
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 = Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody (Result -> Body) -> Result -> Body
forall a b. (a -> b) -> a -> b
$ Int -> SubExp -> Result
forall a. Int -> a -> [a]
replicate Int
k (SubExp -> Result) -> SubExp -> Result
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 -> Result -> BasicOp
Scratch (Type -> PrimType
forall shape u. TypeBase shape u -> PrimType
elemType Type
arr_t) (SubExp
w SubExp -> Result -> Result
forall a. a -> [a] -> [a]
: Int -> Result -> Result
forall a. Int -> [a] -> [a]
drop Int
1 (Type -> Result
forall u. TypeBase Shape u -> Result
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
$
        Result -> SubExp -> Int -> [LParam] -> InternaliseM SubExp
mkOffsetLambdaBody
          ((VName -> SubExp) -> Closure -> Result
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 -> Result -> Body
forall lore. Bindable lore => Stms lore -> Result -> Body lore
mkBody Stms SOACS
offset_stms (Result -> Body) -> Result -> Body
forall a b. (a -> b) -> a -> b
$
              Int -> SubExp -> Result
forall a. Int -> a -> [a]
replicate ([Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
arr_ts) SubExp
offset
                Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ (Param Type -> SubExp) -> [Param Type] -> Result
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
$ Result -> [Int] -> Closure -> [(SubExp, Int, VName)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 (SubExp -> Result
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
$
        Result -> Type -> BasicOp
I.ArrayLit ((VName -> SubExp) -> Closure -> Result
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
  (Result, Result) -> InternaliseM (Result, Result)
forall (m :: * -> *) a. Monad m => a -> m a
return ((VName -> SubExp) -> Closure -> Result
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 :: Result -> SubExp -> Int -> [LParam] -> InternaliseM SubExp
mkOffsetLambdaBody Result
_ 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 Result
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 (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$
              IntType -> Integer -> SubExp
intConst IntType
Int32 (Integer -> SubExp) -> Integer -> SubExp
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
i
      SubExp
next_one <- Result -> SubExp -> Int -> [LParam] -> InternaliseM SubExp
mkOffsetLambdaBody Result
sizes SubExp
c (Int
i Int -> 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 -> Result -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> SubExp -> Result -> 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 -> Result -> Result
forall a. a -> [a] -> [a]
: Int -> Result -> Result
forall a. Int -> [a] -> [a]
take Int
i Result
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
          (Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody [SubExp
this_one])
          (Result -> Body
forall lore. Bindable lore => Result -> 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 Result
substs <- VName -> InternaliseM (Maybe Result)
lookupSubst (VName -> InternaliseM (Maybe Result))
-> VName -> InternaliseM (Maybe Result)
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
E.qualLeaf QualName VName
d
  SubExp
d' <- case Maybe Result
substs of
    Just [SubExp
v] -> SubExp -> InternaliseM SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return SubExp
v
    Maybe Result
_ -> 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
x String -> 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