{-# LANGUAGE Strict #-}
{-# LANGUAGE TypeFamilies #-}

-- | Conversion of a monomorphic, first-order, defunctorised source
-- program to a core Futhark program.
module Futhark.Internalise.Exps (transformProg) where

import Control.Monad.Reader
import Data.List (elemIndex, find, intercalate, intersperse, transpose)
import Data.List.NonEmpty (NonEmpty (..))
import Data.List.NonEmpty qualified as NE
import Data.Map.Strict qualified as M
import Data.Set qualified as S
import Data.Text qualified as T
import Futhark.IR.SOACS as I hiding (stmPat)
import Futhark.Internalise.AccurateSizes
import Futhark.Internalise.Bindings
import Futhark.Internalise.Entry
import Futhark.Internalise.Lambdas
import Futhark.Internalise.Monad as I
import Futhark.Internalise.TypesValues
import Futhark.Transform.Rename as I
import Futhark.Util (splitAt3)
import Futhark.Util.Pretty (align, docText, pretty)
import Language.Futhark as E hiding (TypeArg)

-- | Convert a program in source Futhark to a program in the Futhark
-- core language.
transformProg :: MonadFreshNames m => Bool -> VisibleTypes -> [E.ValBind] -> m (I.Prog SOACS)
transformProg :: forall (m :: * -> *).
MonadFreshNames m =>
Bool -> VisibleTypes -> [ValBind] -> m (Prog SOACS)
transformProg Bool
always_safe VisibleTypes
types [ValBind]
vbinds = do
  (OpaqueTypes
opaques, Stms SOACS
consts, [FunDef SOACS]
funs) <-
    forall (m :: * -> *).
MonadFreshNames m =>
Bool
-> InternaliseM () -> m (OpaqueTypes, Stms SOACS, [FunDef SOACS])
runInternaliseM Bool
always_safe (VisibleTypes -> [ValBind] -> InternaliseM ()
internaliseValBinds VisibleTypes
types [ValBind]
vbinds)
  forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Prog rep -> m (Prog rep)
I.renameProg forall a b. (a -> b) -> a -> b
$ forall rep. OpaqueTypes -> Stms rep -> [FunDef rep] -> Prog rep
I.Prog OpaqueTypes
opaques Stms SOACS
consts [FunDef SOACS]
funs

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

internaliseFunName :: VName -> Name
internaliseFunName :: VName -> Name
internaliseFunName = [Char] -> Name
nameFromString forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> [Char]
prettyString

internaliseValBind :: VisibleTypes -> E.ValBind -> InternaliseM ()
internaliseValBind :: VisibleTypes -> ValBind -> InternaliseM ()
internaliseValBind VisibleTypes
types fb :: ValBind
fb@(E.ValBind Maybe (Info EntryPoint)
entry VName
fname Maybe (TypeExp VName)
retdecl (Info StructRetType
rettype) [TypeParamBase VName]
tparams [PatBase Info VName]
params Exp
body Maybe DocComment
_ [AttrInfo VName]
attrs SrcLoc
loc) = do
  forall a.
[TypeParamBase VName]
-> [PatBase Info VName]
-> ([FParam SOACS] -> [[FParam SOACS]] -> InternaliseM a)
-> InternaliseM a
bindingFParams [TypeParamBase VName]
tparams [PatBase Info VName]
params forall a b. (a -> b) -> a -> b
$ \[FParam SOACS]
shapeparams [[FParam SOACS]]
params' -> do
    let shapenames :: [VName]
shapenames = forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
I.paramName [FParam SOACS]
shapeparams

    ErrorMsg SubExp
msg <- case Maybe (TypeExp VName)
retdecl of
      Just TypeExp VName
dt ->
        forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg
          forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ErrorMsgPart SubExp
"Function return value does not match shape of type " :)
          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 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [ErrorMsgPart SubExp
"Function return value does not match shape of declared return type."]

    (Body SOACS
body', [TypeBase ExtShape Uniqueness]
rettype') <- forall (m :: * -> *) a.
MonadBuilder m =>
m (Result, a) -> m (Body (Rep m), a)
buildBody forall a b. (a -> b) -> a -> b
$ do
      [SubExp]
body_res <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp (VName -> [Char]
baseString VName
fname forall a. Semigroup a => a -> a -> a
<> [Char]
"_res") Exp
body
      [TypeBase ExtShape Uniqueness]
rettype' <-
        forall {u}. [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall shape u.
StructRetType
-> [TypeBase shape u] -> [TypeBase ExtShape Uniqueness]
internaliseReturnType StructRetType
rettype forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
body_res
      Result
body_res' <-
        ErrorMsg SubExp
-> SrcLoc -> [ExtType] -> Result -> InternaliseM Result
ensureResultExtShape ErrorMsg SubExp
msg SrcLoc
loc (forall a b. (a -> b) -> [a] -> [b]
map forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl [TypeBase ExtShape Uniqueness]
rettype') forall a b. (a -> b) -> a -> b
$ [SubExp] -> Result
subExpsRes [SubExp]
body_res
      forall (f :: * -> *) a. Applicative f => a -> f a
pure
        ( Result
body_res',
          forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Int
length (forall u. [TypeBase ExtShape u] -> Set Int
shapeContext [TypeBase ExtShape Uniqueness]
rettype')) (forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64) forall a. [a] -> [a] -> [a]
++ [TypeBase ExtShape Uniqueness]
rettype'
        )

    let all_params :: [Param DeclType]
all_params = [FParam SOACS]
shapeparams forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[FParam SOACS]]
params'

    Attrs
attrs' <- [AttrInfo VName] -> InternaliseM Attrs
internaliseAttrs [AttrInfo VName]
attrs

    let fd :: FunDef SOACS
fd =
          forall rep.
Maybe EntryPoint
-> Attrs
-> Name
-> [RetType rep]
-> [FParam rep]
-> Body rep
-> FunDef rep
I.FunDef
            forall a. Maybe a
Nothing
            Attrs
attrs'
            (VName -> Name
internaliseFunName VName
fname)
            [TypeBase ExtShape Uniqueness]
rettype'
            [Param DeclType]
all_params
            Body SOACS
body'

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

  case Maybe (Info EntryPoint)
entry of
    Just (Info EntryPoint
entry') -> VisibleTypes -> EntryPoint -> ValBind -> InternaliseM ()
generateEntryPoint VisibleTypes
types EntryPoint
entry' ValBind
fb
    Maybe (Info EntryPoint)
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  where
    zeroExts :: [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts [TypeBase ExtShape u]
ts = forall u.
[TypeBase ExtShape u]
-> [TypeBase ExtShape u] -> [TypeBase ExtShape u]
generaliseExtTypes [TypeBase ExtShape u]
ts [TypeBase ExtShape u]
ts

generateEntryPoint :: VisibleTypes -> E.EntryPoint -> E.ValBind -> InternaliseM ()
generateEntryPoint :: VisibleTypes -> EntryPoint -> ValBind -> InternaliseM ()
generateEntryPoint VisibleTypes
types (E.EntryPoint [EntryParam]
e_params EntryType
e_rettype) ValBind
vb = do
  let (E.ValBind Maybe (Info EntryPoint)
_ VName
ofname Maybe (TypeExp VName)
_ (Info StructRetType
rettype) [TypeParamBase VName]
tparams [PatBase Info VName]
params Exp
_ Maybe DocComment
_ [AttrInfo VName]
attrs SrcLoc
loc) = ValBind
vb
  forall a.
[TypeParamBase VName]
-> [PatBase Info VName]
-> ([FParam SOACS] -> [[FParam SOACS]] -> InternaliseM a)
-> InternaliseM a
bindingFParams [TypeParamBase VName]
tparams [PatBase Info VName]
params forall a b. (a -> b) -> a -> b
$ \[FParam SOACS]
shapeparams [[FParam SOACS]]
params' -> do
    let entry_rettype :: [[TypeBase ExtShape Uniqueness]]
entry_rettype = StructRetType -> [[TypeBase ExtShape Uniqueness]]
internaliseEntryReturnType StructRetType
rettype
        (EntryPoint
entry', OpaqueTypes
opaques) =
          VisibleTypes
-> Name
-> [(EntryParam, [Param DeclType])]
-> (EntryType, [[TypeBase Rank Uniqueness]])
-> (EntryPoint, OpaqueTypes)
entryPoint
            VisibleTypes
types
            (VName -> Name
baseName VName
ofname)
            (forall a b. [a] -> [b] -> [(a, b)]
zip [EntryParam]
e_params [[FParam SOACS]]
params')
            (EntryType
e_rettype, forall a b. (a -> b) -> [a] -> [b]
map (forall a b. (a -> b) -> [a] -> [b]
map forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase Rank u
I.rankShaped) [[TypeBase ExtShape Uniqueness]]
entry_rettype)
        args :: [SubExp]
args = forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
I.Var forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dec. Param dec -> VName
I.paramName) forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[FParam SOACS]]
params'

    OpaqueTypes -> InternaliseM ()
addOpaques OpaqueTypes
opaques

    (Body SOACS
entry_body, [TypeBase ExtShape Uniqueness]
ctx_ts) <- forall (m :: * -> *) a.
MonadBuilder m =>
m (Result, a) -> m (Body (Rep m), a)
buildBody forall a b. (a -> b) -> a -> b
$ do
      -- Special case the (rare) situation where the entry point is
      -- not a function.
      Maybe [SubExp]
maybe_const <- VName -> InternaliseM (Maybe [SubExp])
lookupConst VName
ofname
      [SubExp]
vals <- case Maybe [SubExp]
maybe_const of
        Just [SubExp]
ses ->
          forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
ses
        Maybe [SubExp]
Nothing ->
          forall a b. (a, b) -> a
fst forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char]
-> QualName VName
-> [SubExp]
-> SrcLoc
-> InternaliseM ([SubExp], [ExtType])
funcall [Char]
"entry_result" (forall v. v -> QualName v
E.qualName VName
ofname) [SubExp]
args SrcLoc
loc
      [SubExp]
ctx <-
        forall u a. [TypeBase ExtShape u] -> [[a]] -> [a]
extractShapeContext (forall {u}. [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[TypeBase ExtShape Uniqueness]]
entry_rettype)
          forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall u. TypeBase Shape u -> [SubExp]
I.arrayDims forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType) [SubExp]
vals
      forall (f :: * -> *) a. Applicative f => a -> f a
pure ([SubExp] -> Result
subExpsRes forall a b. (a -> b) -> a -> b
$ [SubExp]
ctx forall a. [a] -> [a] -> [a]
++ [SubExp]
vals, forall a b. (a -> b) -> [a] -> [b]
map (forall a b. a -> b -> a
const (forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)) [SubExp]
ctx)

    Attrs
attrs' <- [AttrInfo VName] -> InternaliseM Attrs
internaliseAttrs [AttrInfo VName]
attrs
    FunDef SOACS -> InternaliseM ()
addFunDef forall a b. (a -> b) -> a -> b
$
      forall rep.
Maybe EntryPoint
-> Attrs
-> Name
-> [RetType rep]
-> [FParam rep]
-> Body rep
-> FunDef rep
I.FunDef
        (forall a. a -> Maybe a
Just EntryPoint
entry')
        Attrs
attrs'
        (Name
"entry_" forall a. Semigroup a => a -> a -> a
<> VName -> Name
baseName VName
ofname)
        ([TypeBase ExtShape Uniqueness]
ctx_ts forall a. [a] -> [a] -> [a]
++ forall {u}. [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts (forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[TypeBase ExtShape Uniqueness]]
entry_rettype))
        ([FParam SOACS]
shapeparams forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[FParam SOACS]]
params')
        Body SOACS
entry_body
  where
    zeroExts :: [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts [TypeBase ExtShape u]
ts = forall u.
[TypeBase ExtShape u]
-> [TypeBase ExtShape u] -> [TypeBase ExtShape u]
generaliseExtTypes [TypeBase ExtShape u]
ts [TypeBase ExtShape u]
ts

internaliseBody :: String -> E.Exp -> InternaliseM (Body SOACS)
internaliseBody :: [Char] -> Exp -> InternaliseM (Body SOACS)
internaliseBody [Char]
desc Exp
e =
  forall (m :: * -> *).
MonadBuilder m =>
m Result -> m (Body (Rep m))
buildBody_ forall a b. (a -> b) -> a -> b
$ [SubExp] -> Result
subExpsRes forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp ([Char]
desc forall a. Semigroup a => a -> a -> a
<> [Char]
"_res") Exp
e

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

-- | Only returns those pattern names that are not used in the pattern
-- itself (the "non-existential" part, you could say).
letValExp :: String -> I.Exp SOACS -> InternaliseM [VName]
letValExp :: [Char] -> Exp SOACS -> InternaliseM [VName]
letValExp [Char]
name Exp SOACS
e = do
  [ExtType]
e_t <- forall rep (m :: * -> *).
(HasScope rep m, TypedOp (Op rep)) =>
Exp rep -> m [ExtType]
expExtType Exp SOACS
e
  [VName]
names <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (forall (t :: * -> *) a. Foldable t => t a -> Int
length [ExtType]
e_t) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
name
  forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName]
names Exp SOACS
e
  let ctx :: Set Int
ctx = forall u. [TypeBase ExtShape u] -> Set Int
shapeContext [ExtType]
e_t
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall a. (a -> Bool) -> [a] -> [a]
filter ((forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set Int
ctx) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
names [Int
0 ..]

letValExp' :: String -> I.Exp SOACS -> InternaliseM [SubExp]
letValExp' :: [Char] -> Exp SOACS -> InternaliseM [SubExp]
letValExp' [Char]
_ (BasicOp (SubExp SubExp
se)) = forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp
se]
letValExp' [Char]
name Exp SOACS
ses = forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp SOACS -> InternaliseM [VName]
letValExp [Char]
name Exp SOACS
ses

internaliseAppExp :: String -> E.AppRes -> E.AppExp -> InternaliseM [I.SubExp]
internaliseAppExp :: [Char] -> AppRes -> AppExp -> InternaliseM [SubExp]
internaliseAppExp [Char]
desc AppRes
_ (E.Index Exp
e SliceBase Info VName
idxs SrcLoc
loc) = do
  [VName]
vs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"indexed" Exp
e
  [SubExp]
dims <- case [VName]
vs of
    [] -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [] -- Will this happen?
    VName
v : [VName]
_ -> forall u. TypeBase Shape u -> [SubExp]
I.arrayDims forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
v
  ([DimIndex SubExp]
idxs', Certs
cs) <- SrcLoc
-> [SubExp]
-> SliceBase Info VName
-> InternaliseM ([DimIndex SubExp], Certs)
internaliseSlice SrcLoc
loc [SubExp]
dims SliceBase Info VName
idxs
  let index :: VName -> InternaliseM (Exp SOACS)
index VName
v = do
        TypeBase Shape NoUniqueness
v_t <- forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
v
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp
I.Index VName
v forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoUniqueness -> [DimIndex SubExp] -> Slice SubExp
fullSlice TypeBase Shape NoUniqueness
v_t [DimIndex SubExp]
idxs'
  forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< VName -> InternaliseM (Exp SOACS)
index) [VName]
vs
internaliseAppExp [Char]
desc AppRes
_ (E.Range Exp
start Maybe Exp
maybe_second Inclusiveness Exp
end SrcLoc
loc) = do
  SubExp
start' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"range_start" Exp
start
  SubExp
end' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"range_end" forall a b. (a -> b) -> a -> b
$ case Inclusiveness Exp
end of
    DownToExclusive Exp
e -> Exp
e
    ToInclusive Exp
e -> Exp
e
    UpToExclusive Exp
e -> Exp
e
  Maybe SubExp
maybe_second' <-
    forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ([Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"range_second") Maybe Exp
maybe_second

  -- Construct an error message in case the range is invalid.
  let conv :: SubExp -> InternaliseM SubExp
conv = case Exp -> PatType
E.typeOf Exp
start of
        E.Scalar (E.Prim (E.Unsigned IntType
_)) -> forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntZ IntType
Int64
        PatType
_ -> forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64
  SubExp
start'_i64 <- SubExp -> InternaliseM SubExp
conv SubExp
start'
  SubExp
end'_i64 <- SubExp -> InternaliseM SubExp
conv SubExp
end'
  Maybe SubExp
maybe_second'_i64 <- 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 =
        forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg forall a b. (a -> b) -> a -> b
$
          [ErrorMsgPart SubExp
"Range "]
            forall a. [a] -> [a] -> [a]
++ [forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
start'_i64]
            forall a. [a] -> [a] -> [a]
++ ( case Maybe SubExp
maybe_second'_i64 of
                   Maybe SubExp
Nothing -> []
                   Just SubExp
second_i64 -> [ErrorMsgPart SubExp
"..", forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
second_i64]
               )
            forall a. [a] -> [a] -> [a]
++ ( case Inclusiveness Exp
end of
                   DownToExclusive {} -> [ErrorMsgPart SubExp
"..>"]
                   ToInclusive {} -> [ErrorMsgPart SubExp
"..."]
                   UpToExclusive {} -> [ErrorMsgPart SubExp
"..<"]
               )
            forall a. [a] -> [a] -> [a]
++ [forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
end'_i64, ErrorMsgPart SubExp
" is invalid."]

  (IntType
it, CmpOp
le_op, CmpOp
lt_op) <-
    case Exp -> PatType
E.typeOf Exp
start of
      E.Scalar (E.Prim (E.Signed IntType
it)) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (IntType
it, IntType -> CmpOp
CmpSle IntType
it, IntType -> CmpOp
CmpSlt IntType
it)
      E.Scalar (E.Prim (E.Unsigned IntType
it)) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (IntType
it, IntType -> CmpOp
CmpUle IntType
it, IntType -> CmpOp
CmpUlt IntType
it)
      PatType
start_t -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Start value in range has type " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PatType
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 Exp
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 <-
        forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"subtracted_step" forall a b. (a -> b) -> a -> b
$
          forall rep. BasicOp -> Exp rep
I.BasicOp 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 <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"step_zero" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
start' SubExp
second'
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
subtracted_step, SubExp
step_zero)
    Maybe SubExp
Nothing ->
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
default_step, forall v. IsValue v => v -> SubExp
constant Bool
False)

  SubExp
step_sign <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"s_sign" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (IntType -> UnOp
I.SSignum IntType
it) SubExp
step
  SubExp
step_sign_i64 <- forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
step_sign

  SubExp
bounds_invalid_downwards <-
    forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"bounds_invalid_downwards" forall a b. (a -> b) -> a -> b
$
      forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
        CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp CmpOp
le_op SubExp
start' SubExp
end'
  SubExp
bounds_invalid_upwards <-
    forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"bounds_invalid_upwards" forall a b. (a -> b) -> a -> b
$
      forall rep. BasicOp -> Exp rep
I.BasicOp 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 Exp
end of
    DownToExclusive {} -> do
      SubExp
step_wrong_dir <-
        forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"step_wrong_dir" forall a b. (a -> b) -> a -> b
$
          forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
one
      SubExp
distance <-
        forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"distance" forall a b. (a -> b) -> a -> b
$
          forall rep. BasicOp -> Exp rep
I.BasicOp 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_i64 <- forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
distance
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
distance_i64, SubExp
step_wrong_dir, SubExp
bounds_invalid_downwards)
    UpToExclusive {} -> do
      SubExp
step_wrong_dir <-
        forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"step_wrong_dir" forall a b. (a -> b) -> a -> b
$
          forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
negone
      SubExp
distance <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"distance" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp 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_i64 <- forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
distance
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
distance_i64, SubExp
step_wrong_dir, SubExp
bounds_invalid_upwards)
    ToInclusive {} -> do
      SubExp
downwards <-
        forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"downwards" forall a b. (a -> b) -> a -> b
$
          forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
negone
      SubExp
distance_downwards_exclusive <-
        forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"distance_downwards_exclusive" forall a b. (a -> b) -> a -> b
$
          forall rep. BasicOp -> Exp rep
I.BasicOp 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 <-
        forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"distance_upwards_exclusive" forall a b. (a -> b) -> a -> b
$
          forall rep. BasicOp -> Exp rep
I.BasicOp 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 <-
        forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"bounds_invalid"
          forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
            (forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
downwards)
            (forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
bounds_invalid_downwards])
            (forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
bounds_invalid_upwards])
      SubExp
distance_exclusive <-
        forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"distance_exclusive"
          forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
            (forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
downwards)
            (forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
distance_downwards_exclusive])
            (forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
distance_upwards_exclusive])
      SubExp
distance_exclusive_i64 <- forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
distance_exclusive
      SubExp
distance <-
        forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"distance" forall a b. (a -> b) -> a -> b
$
          forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
            BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp
              (IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
I.OverflowWrap)
              SubExp
distance_exclusive_i64
              (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1)
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
distance, forall v. IsValue v => v -> SubExp
constant Bool
False, SubExp
bounds_invalid)

  SubExp
step_invalid <-
    forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"step_invalid" forall a b. (a -> b) -> a -> b
$
      forall rep. BasicOp -> Exp rep
I.BasicOp 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 <-
    forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"range_invalid" forall a b. (a -> b) -> a -> b
$
      forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
        BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogOr SubExp
step_invalid SubExp
bounds_invalid
  SubExp
valid <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"valid" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
I.Not SubExp
invalid
  Certs
cs <- [Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert [Char]
"range_valid_c" SubExp
valid ErrorMsg SubExp
errmsg SrcLoc
loc

  SubExp
step_i64 <- forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
step
  SubExp
pos_step <-
    forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"pos_step" forall a b. (a -> b) -> a -> b
$
      forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
        BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Mul IntType
Int64 Overflow
I.OverflowWrap) SubExp
step_i64 SubExp
step_sign_i64

  SubExp
num_elems <-
    forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs forall a b. (a -> b) -> a -> b
$
      forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"num_elems" forall a b. (a -> b) -> a -> b
$
        forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
          BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Safety -> BinOp
SDivUp IntType
Int64 Safety
I.Unsafe) SubExp
distance SubExp
pos_step

  SubExp
se <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc (forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> SubExp -> SubExp -> IntType -> BasicOp
I.Iota SubExp
num_elems SubExp
start' SubExp
step IntType
it)
  forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp
se]
internaliseAppExp [Char]
desc (E.AppRes PatType
et [VName]
ext) (E.Coerce Exp
e TypeExp VName
dt SrcLoc
loc) = do
  [SubExp]
ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
  [TypeBase ExtShape Uniqueness]
ts <- forall shape u.
StructRetType
-> [TypeBase shape u] -> [TypeBase ExtShape Uniqueness]
internaliseReturnType (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
E.RetType [VName]
ext (forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatType
et)) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
ses
  [ErrorMsgPart SubExp]
dt' <- TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
dt
  forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
ses [TypeBase ExtShape Uniqueness]
ts) forall a b. (a -> b) -> a -> b
$ \(SubExp
e', TypeBase ExtShape Uniqueness
t') -> do
    [SubExp]
dims <- forall u. TypeBase Shape u -> [SubExp]
arrayDims forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType SubExp
e'
    let parts :: [ErrorMsgPart SubExp]
parts =
          [ErrorMsgPart SubExp
"Value of (core language) shape ("]
            forall a. [a] -> [a] -> [a]
++ forall a. a -> [a] -> [a]
intersperse ErrorMsgPart SubExp
", " (forall a b. (a -> b) -> [a] -> [b]
map (forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64) [SubExp]
dims)
            forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
") cannot match shape of type `"]
            forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
dt'
            forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"`."]
    ErrorMsg SubExp
-> SrcLoc -> ExtType -> [Char] -> SubExp -> InternaliseM SubExp
ensureExtShape (forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [ErrorMsgPart SubExp]
parts) SrcLoc
loc (forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl TypeBase ExtShape Uniqueness
t') [Char]
desc SubExp
e'
internaliseAppExp [Char]
desc AppRes
_ e :: AppExp
e@E.Apply {} =
  case AppExp -> (Function, [(Exp, Maybe VName)])
findFuncall AppExp
e of
    (FunctionHole PatType
t SrcLoc
loc, [(Exp, Maybe VName)]
_args) -> do
      -- The function we are supposed to call doesn't exist, but we
      -- have to synthesize some fake values of the right type.  The
      -- easy way to do this is to just ignore the arguments and
      -- create a hole whose type is the type of the entire
      -- application.
      [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc (forall (f :: * -> *) vn. f PatType -> SrcLoc -> ExpBase f vn
E.Hole (forall a. a -> Info a
Info (forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall dim as.
TypeBase dim as -> ([TypeBase dim ()], TypeBase dim ())
E.unfoldFunType PatType
t)) SrcLoc
loc)
    (FunctionName QualName VName
qfname, [(Exp, Maybe VName)]
args) -> do
      -- Argument evaluation is outermost-in so that any existential sizes
      -- created by function applications can be brought into scope.
      let fname :: Name
fname = [Char] -> Name
nameFromString forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> [Char]
prettyString forall a b. (a -> b) -> a -> b
$ VName -> Name
baseName forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname
          loc :: SrcLoc
loc = forall a. Located a => a -> SrcLoc
srclocOf AppExp
e
          arg_desc :: [Char]
arg_desc = Name -> [Char]
nameToString Name
fname forall a. [a] -> [a] -> [a]
++ [Char]
"_arg"

      -- Some functions are magical (overloaded) and we handle that here.
      case () of
        ()
          -- Short-circuiting operators are magical.
          | VName -> Int
baseTag (forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag,
            VName -> [Char]
baseString (forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) forall a. Eq a => a -> a -> Bool
== [Char]
"&&",
            [(Exp
x, Maybe VName
_), (Exp
y, Maybe VName
_)] <- [(Exp, Maybe VName)]
args ->
              [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc forall a b. (a -> b) -> a -> b
$
                forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
E.AppExp
                  (forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
E.If Exp
x Exp
y (forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
E.Literal (Bool -> PrimValue
E.BoolValue Bool
False) forall a. Monoid a => a
mempty) forall a. Monoid a => a
mempty)
                  (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ PatType -> [VName] -> AppRes
AppRes (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
E.Prim PrimType
E.Bool) [])
          | VName -> Int
baseTag (forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag,
            VName -> [Char]
baseString (forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) forall a. Eq a => a -> a -> Bool
== [Char]
"||",
            [(Exp
x, Maybe VName
_), (Exp
y, Maybe VName
_)] <- [(Exp, Maybe VName)]
args ->
              [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc forall a b. (a -> b) -> a -> b
$
                forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
E.AppExp
                  (forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
E.If Exp
x (forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
E.Literal (Bool -> PrimValue
E.BoolValue Bool
True) forall a. Monoid a => a
mempty) Exp
y forall a. Monoid a => a
mempty)
                  (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ PatType -> [VName] -> AppRes
AppRes (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
E.Prim PrimType
E.Bool) [])
          -- Overloaded and intrinsic functions never take array
          -- arguments (except equality, but those cannot be
          -- existential), so we can safely ignore the existential
          -- dimensions.
          | Just [(TypeBase Size (), [SubExp])] -> InternaliseM [SubExp]
internalise <- QualName VName
-> [Char]
-> SrcLoc
-> Maybe ([(TypeBase Size (), [SubExp])] -> InternaliseM [SubExp])
isOverloadedFunction QualName VName
qfname [Char]
desc SrcLoc
loc -> do
              let prepareArg :: (Exp, b) -> InternaliseM (TypeBase Size (), [SubExp])
prepareArg (Exp
arg, b
_) =
                    (forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct (Exp -> PatType
E.typeOf Exp
arg),) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"arg" Exp
arg
              [(TypeBase Size (), [SubExp])] -> InternaliseM [SubExp]
internalise forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall {b}. (Exp, b) -> InternaliseM (TypeBase Size (), [SubExp])
prepareArg [(Exp, Maybe VName)]
args
          | Just [Char] -> InternaliseM [SubExp]
internalise <- QualName VName
-> [Exp] -> SrcLoc -> Maybe ([Char] -> InternaliseM [SubExp])
isIntrinsicFunction QualName VName
qfname (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(Exp, Maybe VName)]
args) SrcLoc
loc ->
              [Char] -> InternaliseM [SubExp]
internalise [Char]
desc
          | VName -> Int
baseTag (forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag,
            Just (PrimType
rettype, [PrimType]
_) <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
fname Map Name (PrimType, [PrimType])
I.builtInFunctions -> do
              let tag :: [a] -> [(a, Diet)]
tag [a]
ses = [(a
se, Diet
I.Observe) | a
se <- [a]
ses]
              [[SubExp]]
args' <- forall a. [a] -> [a]
reverse forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Char] -> (Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg [Char]
arg_desc) (forall a. [a] -> [a]
reverse [(Exp, Maybe VName)]
args)
              let args'' :: [(SubExp, Diet)]
args'' = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall {a}. [a] -> [(a, Diet)]
tag [[SubExp]]
args'
              [Char] -> Exp SOACS -> InternaliseM [SubExp]
letValExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep.
Name
-> [(SubExp, Diet)]
-> [RetType rep]
-> (Safety, SrcLoc, [SrcLoc])
-> Exp rep
I.Apply Name
fname [(SubExp, Diet)]
args'' [forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
rettype] (Safety
Safe, SrcLoc
loc, [])
          | Bool
otherwise -> do
              [SubExp]
args' <- forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a] -> [a]
reverse forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Char] -> (Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg [Char]
arg_desc) (forall a. [a] -> [a]
reverse [(Exp, Maybe VName)]
args)
              forall a b. (a, b) -> a
fst forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char]
-> QualName VName
-> [SubExp]
-> SrcLoc
-> InternaliseM ([SubExp], [ExtType])
funcall [Char]
desc QualName VName
qfname [SubExp]
args' SrcLoc
loc
internaliseAppExp [Char]
desc AppRes
_ (E.LetPat [SizeBinder VName]
sizes PatBase Info VName
pat Exp
e Exp
body SrcLoc
_) =
  forall a.
[Char]
-> [SizeBinder VName]
-> PatBase Info VName
-> Exp
-> InternaliseM a
-> InternaliseM a
internalisePat [Char]
desc [SizeBinder VName]
sizes PatBase Info VName
pat Exp
e forall a b. (a -> b) -> a -> b
$ [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
body
internaliseAppExp [Char]
_ AppRes
_ (E.LetFun VName
ofname ([TypeParamBase VName], [PatBase Info VName],
 Maybe (TypeExp VName), Info StructRetType, Exp)
_ Exp
_ SrcLoc
_) =
  forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Unexpected LetFun " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString VName
ofname
internaliseAppExp [Char]
desc AppRes
_ (E.DoLoop [VName]
sparams PatBase Info VName
mergepat Exp
mergeexp LoopFormBase Info VName
form Exp
loopbody SrcLoc
loc) = do
  [SubExp]
ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"loop_init" Exp
mergeexp
  ((Body SOACS
loopbody', (LoopForm SOACS
form', [Param DeclType]
shapepat, [Param DeclType]
mergepat', [SubExp]
mergeinit')), Stms SOACS
initstms) <-
    forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms forall a b. (a -> b) -> a -> b
$ [SubExp]
-> LoopFormBase Info VName
-> InternaliseM
     (Body SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
handleForm [SubExp]
ses LoopFormBase Info VName
form

  forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
addStms Stms SOACS
initstms
  [TypeBase Shape NoUniqueness]
mergeinit_ts' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
mergeinit'

  [SubExp]
ctxinit <- [VName]
-> [FParam SOACS]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM [SubExp]
argShapes (forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
I.paramName [Param DeclType]
shapepat) [Param DeclType]
mergepat' [TypeBase Shape NoUniqueness]
mergeinit_ts'

  -- Ensure that the initial loop values match the shapes of the loop
  -- 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 args :: [SubExp]
args = [SubExp]
ctxinit forall a. [a] -> [a] -> [a]
++ [SubExp]
mergeinit'
  [SubExp]
args' <-
    forall u.
Typed (TypeBase Shape u) =>
ErrorMsg SubExp
-> SrcLoc
-> [VName]
-> [TypeBase Shape u]
-> [SubExp]
-> InternaliseM [SubExp]
ensureArgShapes
      ErrorMsg SubExp
"initial loop values have right shape"
      SrcLoc
loc
      (forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
I.paramName [Param DeclType]
shapepat)
      (forall a b. (a -> b) -> [a] -> [b]
map forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
paramType forall a b. (a -> b) -> a -> b
$ [Param DeclType]
shapepat forall a. [a] -> [a] -> [a]
++ [Param DeclType]
mergepat')
      [SubExp]
args

  let dropCond :: [VName] -> [VName]
dropCond = case LoopFormBase Info VName
form of
        E.While {} -> forall a. Int -> [a] -> [a]
drop Int
1
        LoopFormBase Info VName
_ -> forall a. a -> a
id

  -- As above, ensure that the result has the right shape.
  let merge :: [(Param DeclType, SubExp)]
merge = forall a b. [a] -> [b] -> [(a, b)]
zip ([Param DeclType]
shapepat forall a. [a] -> [a] -> [a]
++ [Param DeclType]
mergepat') [SubExp]
args'
      merge_ts :: [TypeBase Shape NoUniqueness]
merge_ts = forall a b. (a -> b) -> [a] -> [b]
map (forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
I.paramType forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(Param DeclType, SubExp)]
merge
  Body SOACS
loopbody'' <-
    forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(Param DeclType, SubExp)]
merge) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf LoopForm SOACS
form' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *).
MonadBuilder m =>
m Result -> m (Body (Rep m))
buildBody_ forall a b. (a -> b) -> a -> b
$
      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [SubExp] -> Result
subExpsRes
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall u.
Typed (TypeBase Shape u) =>
ErrorMsg SubExp
-> SrcLoc
-> [VName]
-> [TypeBase Shape u]
-> [SubExp]
-> InternaliseM [SubExp]
ensureArgShapes
          ErrorMsg SubExp
"shape of loop result does not match shapes in loop parameter"
          SrcLoc
loc
          (forall a b. (a -> b) -> [a] -> [b]
map (forall dec. Param dec -> VName
I.paramName forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(Param DeclType, SubExp)]
merge)
          [TypeBase Shape NoUniqueness]
merge_ts
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp
        forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind Body SOACS
loopbody'

  Attrs
attrs <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks InternaliseEnv -> Attrs
envAttrs
  forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var forall b c a. (b -> c) -> (a -> b) -> a -> c
. [VName] -> [VName]
dropCond
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a. MonadBuilder m => Attrs -> m a -> m a
attributing
      Attrs
attrs
      ([Char] -> Exp SOACS -> InternaliseM [VName]
letValExp [Char]
desc (forall rep.
[(FParam rep, SubExp)] -> LoopForm rep -> Body rep -> Exp rep
I.DoLoop [(Param DeclType, SubExp)]
merge LoopForm SOACS
form' Body SOACS
loopbody''))
  where
    sparams' :: [TypeParamBase VName]
sparams' = forall a b. (a -> b) -> [a] -> [b]
map (forall vn. vn -> SrcLoc -> TypeParamBase vn
`TypeParamDim` forall a. Monoid a => a
mempty) [VName]
sparams

    forLoop :: [Param DeclType]
-> [Param DeclType]
-> [SubExp]
-> LoopForm SOACS
-> InternaliseM
     (Body SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forLoop [Param DeclType]
mergepat' [Param DeclType]
shapepat [SubExp]
mergeinit LoopForm SOACS
form' =
      forall a. InternaliseM (Result, a) -> InternaliseM (Body SOACS, a)
bodyFromStms forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf LoopForm SOACS
form' forall a b. (a -> b) -> a -> b
$ do
        [SubExp]
ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"loopres" Exp
loopbody
        [TypeBase Shape NoUniqueness]
sets <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
ses
        [SubExp]
shapeargs <- [VName]
-> [FParam SOACS]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM [SubExp]
argShapes (forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
I.paramName [Param DeclType]
shapepat) [Param DeclType]
mergepat' [TypeBase Shape NoUniqueness]
sets
        forall (f :: * -> *) a. Applicative f => a -> f a
pure
          ( [SubExp] -> Result
subExpsRes forall a b. (a -> b) -> a -> b
$ [SubExp]
shapeargs forall a. [a] -> [a] -> [a]
++ [SubExp]
ses,
            ( LoopForm SOACS
form',
              [Param DeclType]
shapepat,
              [Param DeclType]
mergepat',
              [SubExp]
mergeinit
            )
          )

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

      VName
i <- forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"i"

      [TypeBase Shape NoUniqueness]
ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
mergeinit
      forall a.
[TypeParamBase VName]
-> PatBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([FParam SOACS] -> [FParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' PatBase Info VName
mergepat [TypeBase Shape NoUniqueness]
ts forall a b. (a -> b) -> a -> b
$ \[FParam SOACS]
shapepat [FParam SOACS]
mergepat' ->
        forall a.
[PatBase Info VName]
-> [TypeBase Shape NoUniqueness]
-> ([LParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLambdaParams [PatBase Info VName
x] (forall a b. (a -> b) -> [a] -> [b]
map forall u. TypeBase Shape u -> TypeBase Shape u
rowType [TypeBase Shape NoUniqueness]
arr_ts) forall a b. (a -> b) -> a -> b
$ \[LParam SOACS]
x_params -> do
          let loopvars :: [(Param (TypeBase Shape NoUniqueness), VName)]
loopvars = forall a b. [a] -> [b] -> [(a, b)]
zip [LParam SOACS]
x_params [VName]
arr'
          [Param DeclType]
-> [Param DeclType]
-> [SubExp]
-> LoopForm SOACS
-> InternaliseM
     (Body SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forLoop [FParam SOACS]
mergepat' [FParam SOACS]
shapepat [SubExp]
mergeinit forall a b. (a -> b) -> a -> b
$
            forall rep.
VName -> IntType -> SubExp -> [(LParam rep, VName)] -> LoopForm rep
I.ForLoop VName
i IntType
Int64 SubExp
w [(Param (TypeBase Shape NoUniqueness), VName)]
loopvars
    handleForm [SubExp]
mergeinit (E.For IdentBase Info VName
i Exp
num_iterations) = do
      SubExp
num_iterations' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"upper_bound" Exp
num_iterations
      TypeBase Shape NoUniqueness
num_iterations_t <- forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
I.subExpType SubExp
num_iterations'
      IntType
it <- case TypeBase Shape NoUniqueness
num_iterations_t of
        I.Prim (IntType IntType
it) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure IntType
it
        TypeBase Shape NoUniqueness
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"internaliseExp DoLoop: invalid type"

      [TypeBase Shape NoUniqueness]
ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
mergeinit
      forall a.
[TypeParamBase VName]
-> PatBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([FParam SOACS] -> [FParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' PatBase Info VName
mergepat [TypeBase Shape NoUniqueness]
ts forall a b. (a -> b) -> a -> b
$
        \[FParam SOACS]
shapepat [FParam SOACS]
mergepat' ->
          [Param DeclType]
-> [Param DeclType]
-> [SubExp]
-> LoopForm SOACS
-> InternaliseM
     (Body SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forLoop [FParam SOACS]
mergepat' [FParam SOACS]
shapepat [SubExp]
mergeinit forall a b. (a -> b) -> a -> b
$
            forall rep.
VName -> IntType -> SubExp -> [(LParam rep, VName)] -> LoopForm rep
I.ForLoop (forall (f :: * -> *) vn. IdentBase f vn -> vn
E.identName IdentBase Info VName
i) IntType
it SubExp
num_iterations' []
    handleForm [SubExp]
mergeinit (E.While Exp
cond) = do
      [TypeBase Shape NoUniqueness]
ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
mergeinit
      forall a.
[TypeParamBase VName]
-> PatBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([FParam SOACS] -> [FParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' PatBase Info VName
mergepat [TypeBase Shape NoUniqueness]
ts forall a b. (a -> b) -> a -> b
$ \[FParam SOACS]
shapepat [FParam SOACS]
mergepat' -> do
        [TypeBase Shape NoUniqueness]
mergeinit_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
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.
        [SubExp]
shapeinit <- [VName]
-> [FParam SOACS]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM [SubExp]
argShapes (forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
I.paramName [FParam SOACS]
shapepat) [FParam SOACS]
mergepat' [TypeBase Shape NoUniqueness]
mergeinit_ts

        (SubExp
loop_initial_cond, Stms SOACS
init_loop_cond_stms) <- forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms forall a b. (a -> b) -> a -> b
$ do
          forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (forall a b. [a] -> [b] -> [(a, b)]
zip [FParam SOACS]
shapepat [SubExp]
shapeinit) forall a b. (a -> b) -> a -> b
$ \(Param DeclType
p, SubExp
se) ->
            forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [forall dec. Param dec -> VName
I.paramName Param DeclType
p] forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
          forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (forall a b. [a] -> [b] -> [(a, b)]
zip [FParam SOACS]
mergepat' [SubExp]
mergeinit) forall a b. (a -> b) -> a -> b
$ \(Param DeclType
p, SubExp
se) ->
            forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (SubExp
se forall a. Eq a => a -> a -> Bool
== VName -> SubExp
I.Var (forall dec. Param dec -> VName
I.paramName Param DeclType
p)) forall a b. (a -> b) -> a -> b
$
              forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [forall dec. Param dec -> VName
I.paramName Param DeclType
p] forall a b. (a -> b) -> a -> b
$
                forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$
                  case SubExp
se of
                    I.Var VName
v
                      | Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall shape u. TypeBase shape u -> Bool
primType forall a b. (a -> b) -> a -> b
$ forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
paramType Param DeclType
p ->
                          ReshapeKind -> Shape -> VName -> BasicOp
Reshape ReshapeKind
I.ReshapeCoerce (forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape forall a b. (a -> b) -> a -> b
$ forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
paramType Param DeclType
p) VName
v
                    SubExp
_ -> SubExp -> BasicOp
SubExp SubExp
se
          [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"loop_cond" Exp
cond

        forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
addStms Stms SOACS
init_loop_cond_stms

        forall a. InternaliseM (Result, a) -> InternaliseM (Body SOACS, a)
bodyFromStms forall a b. (a -> b) -> a -> b
$ do
          [SubExp]
ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"loopres" Exp
loopbody
          [TypeBase Shape NoUniqueness]
sets <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
ses
          Param DeclType
loop_while <- forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"loop_while" forall a b. (a -> b) -> a -> b
$ forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.Bool
          [SubExp]
shapeargs <- [VName]
-> [FParam SOACS]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM [SubExp]
argShapes (forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
I.paramName [FParam SOACS]
shapepat) [FParam SOACS]
mergepat' [TypeBase Shape NoUniqueness]
sets

          -- Careful not to clobber anything.
          Body SOACS
loop_end_cond_body <- forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Body rep -> m (Body rep)
renameBody forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< forall (m :: * -> *).
MonadBuilder m =>
m Result -> m (Body (Rep m))
buildBody_ forall a b. (a -> b) -> a -> b
$ do
            forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (forall a b. [a] -> [b] -> [(a, b)]
zip [FParam SOACS]
shapepat [SubExp]
shapeargs) forall a b. (a -> b) -> a -> b
$ \(Param DeclType
p, SubExp
se) ->
              forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (SubExp
se forall a. Eq a => a -> a -> Bool
== VName -> SubExp
I.Var (forall dec. Param dec -> VName
I.paramName Param DeclType
p)) forall a b. (a -> b) -> a -> b
$
                forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [forall dec. Param dec -> VName
I.paramName Param DeclType
p] forall a b. (a -> b) -> a -> b
$
                  forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$
                    SubExp -> BasicOp
SubExp SubExp
se
            forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (forall a b. [a] -> [b] -> [(a, b)]
zip [FParam SOACS]
mergepat' [SubExp]
ses) forall a b. (a -> b) -> a -> b
$ \(Param DeclType
p, SubExp
se) ->
              forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (SubExp
se forall a. Eq a => a -> a -> Bool
== VName -> SubExp
I.Var (forall dec. Param dec -> VName
I.paramName Param DeclType
p)) forall a b. (a -> b) -> a -> b
$
                forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [forall dec. Param dec -> VName
I.paramName Param DeclType
p] forall a b. (a -> b) -> a -> b
$
                  forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$
                    case SubExp
se of
                      I.Var VName
v
                        | Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall shape u. TypeBase shape u -> Bool
primType forall a b. (a -> b) -> a -> b
$ forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
paramType Param DeclType
p ->
                            ReshapeKind -> Shape -> VName -> BasicOp
Reshape ReshapeKind
I.ReshapeCoerce (forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape forall a b. (a -> b) -> a -> b
$ forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
paramType Param DeclType
p) VName
v
                      SubExp
_ -> SubExp -> BasicOp
SubExp SubExp
se
            [SubExp] -> Result
subExpsRes forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"loop_cond" Exp
cond
          Result
loop_end_cond <- forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind Body SOACS
loop_end_cond_body

          forall (f :: * -> *) a. Applicative f => a -> f a
pure
            ( [SubExp] -> Result
subExpsRes [SubExp]
shapeargs forall a. [a] -> [a] -> [a]
++ Result
loop_end_cond forall a. [a] -> [a] -> [a]
++ [SubExp] -> Result
subExpsRes [SubExp]
ses,
              ( forall rep. VName -> LoopForm rep
I.WhileLoop forall a b. (a -> b) -> a -> b
$ forall dec. Param dec -> VName
I.paramName Param DeclType
loop_while,
                [FParam SOACS]
shapepat,
                Param DeclType
loop_while forall a. a -> [a] -> [a]
: [FParam SOACS]
mergepat',
                SubExp
loop_initial_cond forall a. a -> [a] -> [a]
: [SubExp]
mergeinit
              )
            )
internaliseAppExp [Char]
desc AppRes
_ (E.LetWith IdentBase Info VName
name IdentBase Info VName
src SliceBase Info VName
idxs Exp
ve Exp
body SrcLoc
loc) = do
  let pat :: PatBase Info VName
pat = forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
E.Id (forall (f :: * -> *) vn. IdentBase f vn -> vn
E.identName IdentBase Info VName
name) (forall (f :: * -> *) vn. IdentBase f vn -> f PatType
E.identType IdentBase Info VName
name) SrcLoc
loc
      src_t :: Info PatType
src_t = forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
E.fromStruct forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) vn. IdentBase f vn -> f PatType
E.identType IdentBase Info VName
src
      e :: Exp
e = forall (f :: * -> *) vn.
ExpBase f vn
-> SliceBase f vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
E.Update (forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
E.Var (forall v. v -> QualName v
E.qualName forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. IdentBase f vn -> vn
E.identName IdentBase Info VName
src) Info PatType
src_t SrcLoc
loc) SliceBase Info VName
idxs Exp
ve SrcLoc
loc
  [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc forall a b. (a -> b) -> a -> b
$
    forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
E.AppExp
      (forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
E.LetPat [] PatBase Info VName
pat Exp
e Exp
body SrcLoc
loc)
      (forall a. a -> Info a
Info (PatType -> [VName] -> AppRes
AppRes (Exp -> PatType
E.typeOf Exp
body) forall a. Monoid a => a
mempty))
internaliseAppExp [Char]
desc AppRes
_ (E.Match Exp
e NonEmpty (CaseBase Info VName)
orig_cs SrcLoc
_) = do
  [SubExp]
ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp ([Char]
desc forall a. [a] -> [a] -> [a]
++ [Char]
"_scrutinee") Exp
e
  NonEmpty (Case (InternaliseM (Body SOACS)))
cs <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([SubExp]
-> CaseBase Info VName
-> InternaliseM (Case (InternaliseM (Body SOACS)))
onCase [SubExp]
ses) NonEmpty (CaseBase Info VName)
orig_cs
  case forall a. NonEmpty a -> (a, Maybe (NonEmpty a))
NE.uncons NonEmpty (Case (InternaliseM (Body SOACS)))
cs of
    (I.Case [Maybe PrimValue]
_ InternaliseM (Body SOACS)
body, Maybe (NonEmpty (Case (InternaliseM (Body SOACS))))
Nothing) ->
      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Body SOACS)
body
    (Case (InternaliseM (Body SOACS)),
 Maybe (NonEmpty (Case (InternaliseM (Body SOACS)))))
_ -> do
      [Char] -> Exp SOACS -> InternaliseM [SubExp]
letValExp' [Char]
desc forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
[SubExp]
-> [Case (m (Body (Rep m)))] -> m (Body (Rep m)) -> m (Exp (Rep m))
eMatch [SubExp]
ses (forall a. NonEmpty a -> [a]
NE.init NonEmpty (Case (InternaliseM (Body SOACS)))
cs) (forall body. Case body -> body
I.caseBody forall a b. (a -> b) -> a -> b
$ forall a. NonEmpty a -> a
NE.last NonEmpty (Case (InternaliseM (Body SOACS)))
cs)
  where
    onCase :: [SubExp]
-> CaseBase Info VName
-> InternaliseM (Case (InternaliseM (Body SOACS)))
onCase [SubExp]
ses (E.CasePat PatBase Info VName
p Exp
case_e SrcLoc
_) = do
      ([Maybe PrimValue]
cmps, [SubExp]
pertinent) <- PatBase Info VName
-> [SubExp] -> InternaliseM ([Maybe PrimValue], [SubExp])
generateCond PatBase Info VName
p [SubExp]
ses
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall body. [Maybe PrimValue] -> body -> Case body
I.Case [Maybe PrimValue]
cmps forall a b. (a -> b) -> a -> b
$
        forall a.
[SizeBinder VName]
-> PatBase Info VName
-> [SubExp]
-> InternaliseM a
-> InternaliseM a
internalisePat' [] PatBase Info VName
p [SubExp]
pertinent forall a b. (a -> b) -> a -> b
$
          [Char] -> Exp -> InternaliseM (Body SOACS)
internaliseBody [Char]
"case" Exp
case_e
internaliseAppExp [Char]
desc AppRes
_ (E.If Exp
ce Exp
te Exp
fe SrcLoc
_) =
  [Char] -> Exp SOACS -> InternaliseM [SubExp]
letValExp' [Char]
desc
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
      (forall rep. BasicOp -> Exp rep
BasicOp forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
SubExp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"cond" Exp
ce)
      ([Char] -> Exp -> InternaliseM (Body SOACS)
internaliseBody ([Char]
desc forall a. Semigroup a => a -> a -> a
<> [Char]
"_t") Exp
te)
      ([Char] -> Exp -> InternaliseM (Body SOACS)
internaliseBody ([Char]
desc forall a. Semigroup a => a -> a -> a
<> [Char]
"_f") Exp
fe)
internaliseAppExp [Char]
_ AppRes
_ e :: AppExp
e@E.BinOp {} =
  forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseAppExp: Unexpected BinOp " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString AppExp
e

internaliseExp :: String -> E.Exp -> InternaliseM [I.SubExp]
internaliseExp :: [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc (E.Parens Exp
e SrcLoc
_) =
  [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
internaliseExp [Char]
desc (E.Hole (Info PatType
t) SrcLoc
loc) = do
  let msg :: Text
msg = forall a. Doc a -> Text
docText forall a b. (a -> b) -> a -> b
$ Doc Any
"Reached hole of type: " forall a. Semigroup a => a -> a -> a
<> forall ann. Doc ann -> Doc ann
align (forall a ann. Pretty a => a -> Doc ann
pretty PatType
t)
      ts :: [TypeBase ExtShape Uniqueness]
ts = TypeBase Size () -> [TypeBase ExtShape Uniqueness]
internaliseType (forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatType
t)
  Certs
c <- [Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert [Char]
"hole_c" (forall v. IsValue v => v -> SubExp
constant Bool
False) (forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [forall a. Text -> ErrorMsgPart a
ErrorString Text
msg]) SrcLoc
loc
  case forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall u. TypeBase ExtShape u -> Maybe (TypeBase Shape u)
hasStaticShape [TypeBase ExtShape Uniqueness]
ts of
    Maybe [DeclType]
Nothing ->
      forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Hole at " forall a. Semigroup a => a -> a -> a
<> forall a. Located a => a -> [Char]
locStr SrcLoc
loc forall a. Semigroup a => a -> a -> a
<> [Char]
" has existential type:\n" forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show [TypeBase ExtShape Uniqueness]
ts
    Just [DeclType]
ts' ->
      -- Make sure we always generate a binding, even for primitives.
      forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap VName -> SubExp
I.Var forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
desc forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< forall (m :: * -> *).
MonadBuilder m =>
TypeBase Shape NoUniqueness -> m (Exp (Rep m))
eBlank forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl) [DeclType]
ts'
internaliseExp [Char]
desc (E.QualParens (QualName VName, SrcLoc)
_ Exp
e SrcLoc
_) =
  [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
internaliseExp [Char]
desc (E.StringLit [Word8]
vs SrcLoc
_) =
  forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$
    forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
      [SubExp] -> TypeBase Shape NoUniqueness -> BasicOp
I.ArrayLit (forall a b. (a -> b) -> [a] -> [b]
map forall v. IsValue v => v -> SubExp
constant [Word8]
vs) forall a b. (a -> b) -> a -> b
$
        forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int8
internaliseExp [Char]
_ (E.Var (E.QualName [VName]
_ VName
name) Info PatType
_ SrcLoc
_) = do
  Maybe [SubExp]
subst <- VName -> InternaliseM (Maybe [SubExp])
lookupSubst VName
name
  case Maybe [SubExp]
subst of
    Just [SubExp]
substs -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
substs
    Maybe [SubExp]
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [VName -> SubExp
I.Var VName
name]
internaliseExp [Char]
desc (E.AppExp AppExp
e (Info AppRes
appres)) = do
  [SubExp]
ses <- [Char] -> AppRes -> AppExp -> InternaliseM [SubExp]
internaliseAppExp [Char]
desc AppRes
appres AppExp
e
  AppRes -> [SubExp] -> InternaliseM ()
bindExtSizes AppRes
appres [SubExp]
ses
  forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
ses
internaliseExp [Char]
_ (E.TupLit [] SrcLoc
_) =
  forall (f :: * -> *) a. Applicative f => a -> f a
pure [forall v. IsValue v => v -> SubExp
constant PrimValue
UnitValue]
internaliseExp [Char]
_ (E.RecordLit [] SrcLoc
_) =
  forall (f :: * -> *) a. Applicative f => a -> f a
pure [forall v. IsValue v => v -> SubExp
constant PrimValue
UnitValue]
internaliseExp [Char]
desc (E.TupLit [Exp]
es SrcLoc
_) = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc) [Exp]
es
internaliseExp [Char]
desc (E.RecordLit [FieldBase Info VName]
orig_fields SrcLoc
_) =
  forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Map Name a -> [(Name, a)]
sortFields forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
M.unions forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM FieldBase Info VName -> InternaliseM (Map Name [SubExp])
internaliseField [FieldBase Info VName]
orig_fields
  where
    internaliseField :: FieldBase Info VName -> InternaliseM (Map Name [SubExp])
internaliseField (E.RecordFieldExplicit Name
name Exp
e SrcLoc
_) =
      forall k a. k -> a -> Map k a
M.singleton Name
name forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
    internaliseField (E.RecordFieldImplicit VName
name Info PatType
t SrcLoc
loc) =
      FieldBase Info VName -> InternaliseM (Map Name [SubExp])
internaliseField forall a b. (a -> b) -> a -> b
$
        forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
E.RecordFieldExplicit
          (VName -> Name
baseName VName
name)
          (forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
E.Var (forall v. v -> QualName v
E.qualName VName
name) Info PatType
t SrcLoc
loc)
          SrcLoc
loc
internaliseExp [Char]
desc (E.ArrayLit [Exp]
es (Info PatType
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, [Exp]
e') : [([Int], [Exp])]
es') <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> Maybe ([Int], [Exp])
isArrayLiteral [Exp]
es,
    Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
eshape,
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (([Int]
eshape ==) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [([Int], [Exp])]
es',
    Just PatType
basetype <- forall dim as. Int -> TypeBase dim as -> Maybe (TypeBase dim as)
E.peelArray (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
eshape) PatType
arr_t = do
      let flat_lit :: Exp
flat_lit = forall (f :: * -> *) vn.
[ExpBase f vn] -> f PatType -> SrcLoc -> ExpBase f vn
E.ArrayLit ([Exp]
e' forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a b. (a, b) -> b
snd [([Int], [Exp])]
es') (forall a. a -> Info a
Info PatType
basetype) SrcLoc
loc
          new_shape :: [Int]
new_shape = forall (t :: * -> *) a. Foldable t => t a -> Int
length [Exp]
es forall a. a -> [a] -> [a]
: [Int]
eshape
      [VName]
flat_arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"flat_literal" Exp
flat_lit
      forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
flat_arrs forall a b. (a -> b) -> a -> b
$ \VName
flat_arr -> do
        TypeBase Shape NoUniqueness
flat_arr_t <- forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
flat_arr
        let new_shape' :: Shape
new_shape' =
              Shape -> Int -> Shape -> Shape
reshapeOuter
                (forall d. [d] -> ShapeBase d
I.Shape forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (IntType -> Integer -> SubExp
intConst IntType
Int64 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Integral a => a -> Integer
toInteger) [Int]
new_shape)
                Int
1
                forall a b. (a -> b) -> a -> b
$ forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape TypeBase Shape NoUniqueness
flat_arr_t
        forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ ReshapeKind -> Shape -> VName -> BasicOp
I.Reshape ReshapeKind
I.ReshapeArbitrary Shape
new_shape' VName
flat_arr
  | Bool
otherwise = do
      [[SubExp]]
es' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"arr_elem") [Exp]
es
      let arr_t_ext :: [TypeBase ExtShape Uniqueness]
arr_t_ext = TypeBase Size () -> [TypeBase ExtShape Uniqueness]
internaliseType forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatType
arr_t

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

      let arraylit :: [SubExp] -> TypeBase Shape NoUniqueness -> InternaliseM (Exp SOACS)
arraylit [SubExp]
ks TypeBase Shape NoUniqueness
rt = do
            [SubExp]
ks' <-
              forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM
                ( ErrorMsg SubExp
-> SrcLoc
-> TypeBase Shape NoUniqueness
-> [Char]
-> SubExp
-> InternaliseM SubExp
ensureShape
                    ErrorMsg SubExp
"shape of element differs from shape of first element"
                    SrcLoc
loc
                    TypeBase Shape NoUniqueness
rt
                    [Char]
"elem_reshaped"
                )
                [SubExp]
ks
            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ [SubExp] -> TypeBase Shape NoUniqueness -> BasicOp
I.ArrayLit [SubExp]
ks' TypeBase Shape NoUniqueness
rt

      forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc)
        forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[SubExp]]
es'
          then forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([SubExp] -> TypeBase Shape NoUniqueness -> InternaliseM (Exp SOACS)
arraylit []) [TypeBase Shape NoUniqueness]
rowtypes
          else forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM [SubExp] -> TypeBase Shape NoUniqueness -> InternaliseM (Exp SOACS)
arraylit (forall a. [[a]] -> [[a]]
transpose [[SubExp]]
es') [TypeBase Shape NoUniqueness]
rowtypes
  where
    isArrayLiteral :: E.Exp -> Maybe ([Int], [E.Exp])
    isArrayLiteral :: Exp -> Maybe ([Int], [Exp])
isArrayLiteral (E.ArrayLit [Exp]
inner_es Info PatType
_ SrcLoc
_) = do
      ([Int]
eshape, [Exp]
e) : [([Int], [Exp])]
inner_es' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> Maybe ([Int], [Exp])
isArrayLiteral [Exp]
inner_es
      forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (([Int]
eshape ==) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [([Int], [Exp])]
inner_es'
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Exp]
inner_es forall a. a -> [a] -> [a]
: [Int]
eshape, [Exp]
e forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a b. (a, b) -> b
snd [([Int], [Exp])]
inner_es')
    isArrayLiteral Exp
e =
      forall a. a -> Maybe a
Just ([], [Exp
e])
internaliseExp [Char]
desc (E.Ascript Exp
e TypeExp VName
_ SrcLoc
_) =
  [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
internaliseExp [Char]
desc (E.Negate Exp
e SrcLoc
_) = do
  SubExp
e' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"negate_arg" Exp
e
  TypeBase Shape NoUniqueness
et <- forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType SubExp
e'
  case TypeBase Shape NoUniqueness
et of
    I.Prim (I.IntType IntType
t) ->
      forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp 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) ->
      forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp 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'
    TypeBase Shape NoUniqueness
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise.internaliseExp: non-numeric type in Negate"
internaliseExp [Char]
desc (E.Not Exp
e SrcLoc
_) = do
  SubExp
e' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"not_arg" Exp
e
  TypeBase Shape NoUniqueness
et <- forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType SubExp
e'
  case TypeBase Shape NoUniqueness
et of
    I.Prim (I.IntType IntType
t) ->
      forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp 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 ->
      forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
I.Not SubExp
e'
    TypeBase Shape NoUniqueness
_ ->
      forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise.internaliseExp: non-int/bool type in Not"
internaliseExp [Char]
desc (E.Update Exp
src SliceBase Info VName
slice Exp
ve SrcLoc
loc) = do
  [SubExp]
ves <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"lw_val" Exp
ve
  [VName]
srcs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"src" Exp
src
  [SubExp]
dims <- case [VName]
srcs of
    [] -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [] -- Will this happen?
    VName
v : [VName]
_ -> forall u. TypeBase Shape u -> [SubExp]
I.arrayDims forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
v
  ([DimIndex SubExp]
idxs', Certs
cs) <- SrcLoc
-> [SubExp]
-> SliceBase Info VName
-> InternaliseM ([DimIndex SubExp], Certs)
internaliseSlice SrcLoc
loc [SubExp]
dims SliceBase Info VName
slice

  let comb :: VName -> SubExp -> InternaliseM VName
comb VName
sname SubExp
ve' = do
        TypeBase Shape NoUniqueness
sname_t <- forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
sname
        let full_slice :: Slice SubExp
full_slice = TypeBase Shape NoUniqueness -> [DimIndex SubExp] -> Slice SubExp
fullSlice TypeBase Shape NoUniqueness
sname_t [DimIndex SubExp]
idxs'
            rowtype :: TypeBase Shape NoUniqueness
rowtype = TypeBase Shape NoUniqueness
sname_t forall oldshape u.
TypeBase oldshape u -> [SubExp] -> TypeBase Shape u
`setArrayDims` forall d. Slice d -> [d]
sliceDims Slice SubExp
full_slice
        SubExp
ve'' <-
          ErrorMsg SubExp
-> SrcLoc
-> TypeBase Shape NoUniqueness
-> [Char]
-> SubExp
-> InternaliseM SubExp
ensureShape
            ErrorMsg SubExp
"shape of value does not match shape of source array"
            SrcLoc
loc
            TypeBase Shape NoUniqueness
rowtype
            [Char]
"lw_val_correct_shape"
            SubExp
ve'
        forall (m :: * -> *).
MonadBuilder m =>
[Char] -> VName -> Slice SubExp -> Exp (Rep m) -> m VName
letInPlace [Char]
desc VName
sname Slice SubExp
full_slice forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
ve''
  forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM VName -> SubExp -> InternaliseM VName
comb [VName]
srcs [SubExp]
ves
internaliseExp [Char]
desc (E.RecordUpdate Exp
src [Name]
fields Exp
ve Info PatType
_ SrcLoc
_) = do
  [SubExp]
src' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
src
  [SubExp]
ve' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
ve
  forall {m :: * -> *} {als} {a}.
Monad m =>
TypeBase Size als -> [Name] -> [a] -> [a] -> m [a]
replace (Exp -> PatType
E.typeOf Exp
src forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` ()) [Name]
fields [SubExp]
ve' [SubExp]
src'
  where
    replace :: TypeBase Size als -> [Name] -> [a] -> [a] -> m [a]
replace (E.Scalar (E.Record Map Name (TypeBase Size als)
m)) (Name
f : [Name]
fs) [a]
ve' [a]
src'
      | Just TypeBase Size als
t <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f Map Name (TypeBase Size als)
m = do
          let i :: Int
i =
                forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall als. TypeBase Size als -> Int
internalisedTypeSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) forall a b. (a -> b) -> a -> b
$
                  forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((forall a. Eq a => a -> a -> Bool
/= Name
f) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Map Name a -> [(Name, a)]
sortFields forall a b. (a -> b) -> a -> b
$
                    Map Name (TypeBase Size als)
m
              k :: Int
k = forall als. TypeBase Size als -> Int
internalisedTypeSize TypeBase Size als
t
              ([a]
bef, [a]
to_update, [a]
aft) = forall a. Int -> Int -> [a] -> ([a], [a], [a])
splitAt3 Int
i Int
k [a]
src'
          [a]
src'' <- TypeBase Size als -> [Name] -> [a] -> [a] -> m [a]
replace TypeBase Size als
t [Name]
fs [a]
ve' [a]
to_update
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [a]
bef forall a. [a] -> [a] -> [a]
++ [a]
src'' forall a. [a] -> [a] -> [a]
++ [a]
aft
    replace TypeBase Size als
_ [Name]
_ [a]
ve' [a]
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure [a]
ve'
internaliseExp [Char]
desc (E.Attr AttrInfo VName
attr Exp
e SrcLoc
loc) = do
  Attr
attr' <- AttrInfo VName -> InternaliseM Attr
internaliseAttr AttrInfo VName
attr
  [SubExp]
e' <- forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (Attr -> InternaliseEnv -> InternaliseEnv
f Attr
attr') forall a b. (a -> b) -> a -> b
$ [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
  case Attr
attr' of
    Attr
"trace" ->
      Text -> [SubExp] -> InternaliseM [SubExp]
traceRes ([Char] -> Text
T.pack forall a b. (a -> b) -> a -> b
$ forall a. Located a => a -> [Char]
locStr SrcLoc
loc) [SubExp]
e'
    I.AttrComp Name
"trace" [I.AttrName Name
tag] ->
      Text -> [SubExp] -> InternaliseM [SubExp]
traceRes (Name -> Text
nameToText Name
tag) [SubExp]
e'
    Attr
"opaque" ->
      forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. BasicOp -> Exp rep
BasicOp forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpaqueOp -> SubExp -> BasicOp
Opaque OpaqueOp
OpaqueNil) [SubExp]
e'
    Attr
_ ->
      forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
e'
  where
    traceRes :: Text -> [SubExp] -> InternaliseM [SubExp]
traceRes Text
tag' =
      forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. BasicOp -> Exp rep
BasicOp forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpaqueOp -> SubExp -> BasicOp
Opaque (Text -> OpaqueOp
OpaqueTrace Text
tag'))
    f :: Attr -> InternaliseEnv -> InternaliseEnv
f Attr
attr' InternaliseEnv
env
      | Attr
attr' forall a. Eq a => a -> a -> Bool
== Attr
"unsafe",
        Bool -> Bool
not 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 forall a. Semigroup a => a -> a -> a
<> Attr -> Attrs
oneAttr Attr
attr'}
internaliseExp [Char]
desc (E.Assert Exp
e1 Exp
e2 (Info Text
check) SrcLoc
loc) = do
  SubExp
e1' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"assert_cond" Exp
e1
  Certs
c <- [Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert [Char]
"assert_c" SubExp
e1' (forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [forall a. Text -> ErrorMsgPart a
ErrorString forall a b. (a -> b) -> a -> b
$ Text
"Assertion is false: " forall a. Semigroup a => a -> a -> a
<> Text
check]) SrcLoc
loc
  -- Make sure there are some bindings to certify.
  forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall {m :: * -> *}. MonadBuilder m => SubExp -> m SubExp
rebind forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e2
  where
    rebind :: SubExp -> m SubExp
rebind SubExp
v = do
      VName
v' <- forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"assert_res"
      forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
v'] forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
v
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
v'
internaliseExp [Char]
_ (E.Constr Name
c [Exp]
es (Info (E.Scalar (E.Sum Map Name [PatType]
fs))) SrcLoc
_) = do
  ([TypeBase ExtShape Uniqueness]
ts, Map Name (Int, [Int])
constr_map) <- Map Name [TypeBase Size ()]
-> InternaliseM
     ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
internaliseSumType forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> b) -> Map k a -> Map k b
M.map (forall a b. (a -> b) -> [a] -> [b]
map forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct) Map Name [PatType]
fs
  [SubExp]
es' <- forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"payload") [Exp]
es

  let noExt :: p -> f SubExp
noExt p
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0
  [TypeBase Shape NoUniqueness]
ts' <- forall (m :: * -> *) u.
Monad m =>
(Int -> m SubExp) -> [TypeBase ExtShape u] -> m [TypeBase Shape u]
instantiateShapes forall {f :: * -> *} {p}. Applicative f => p -> f SubExp
noExt forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
fromDecl [TypeBase ExtShape Uniqueness]
ts

  case 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 (forall a. Integral a => a -> Integer
toInteger Int
i) :) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {f :: * -> *} {t}.
(Num t, MonadBuilder f, Eq t) =>
t -> [TypeBase Shape NoUniqueness] -> [(t, SubExp)] -> f [SubExp]
clauses Int
0 [TypeBase Shape NoUniqueness]
ts' (forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
js [SubExp]
es')
    Maybe (Int, [Int])
Nothing ->
      forall a. HasCallStack => [Char] -> a
error [Char]
"internaliseExp Constr: missing constructor"
  where
    clauses :: t -> [TypeBase Shape NoUniqueness] -> [(t, SubExp)] -> f [SubExp]
clauses t
j (TypeBase Shape NoUniqueness
t : [TypeBase Shape NoUniqueness]
ts) [(t, SubExp)]
js_to_es
      | Just SubExp
e <- t
j forall a b. Eq a => a -> [(a, b)] -> Maybe b
`lookup` [(t, SubExp)]
js_to_es =
          (SubExp
e :) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> t -> [TypeBase Shape NoUniqueness] -> [(t, SubExp)] -> f [SubExp]
clauses (t
j forall a. Num a => a -> a -> a
+ t
1) [TypeBase Shape NoUniqueness]
ts [(t, SubExp)]
js_to_es
      | Bool
otherwise = do
          SubExp
blank <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"zero" forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
MonadBuilder m =>
TypeBase Shape NoUniqueness -> m (Exp (Rep m))
eBlank TypeBase Shape NoUniqueness
t
          (SubExp
blank :) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> t -> [TypeBase Shape NoUniqueness] -> [(t, SubExp)] -> f [SubExp]
clauses (t
j forall a. Num a => a -> a -> a
+ t
1) [TypeBase Shape NoUniqueness]
ts [(t, SubExp)]
js_to_es
    clauses t
_ [] [(t, SubExp)]
_ =
      forall (f :: * -> *) a. Applicative f => a -> f a
pure []
internaliseExp [Char]
_ (E.Constr Name
_ [Exp]
_ (Info PatType
t) SrcLoc
loc) =
  forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: constructor with type " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PatType
t forall a. [a] -> [a] -> [a]
++ [Char]
" at " forall a. [a] -> [a] -> [a]
++ forall a. Located a => a -> [Char]
locStr SrcLoc
loc
-- The "interesting" cases are over, now it's mostly boilerplate.

internaliseExp [Char]
_ (E.Literal PrimValue
v SrcLoc
_) =
  forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
I.Constant forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimValue
internalisePrimValue PrimValue
v]
internaliseExp [Char]
_ (E.IntLit Integer
v (Info PatType
t) SrcLoc
_) =
  case PatType
t of
    E.Scalar (E.Prim (E.Signed IntType
it)) ->
      forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
I.Constant forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
I.IntValue forall a b. (a -> b) -> a -> b
$ forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
v]
    E.Scalar (E.Prim (E.Unsigned IntType
it)) ->
      forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
I.Constant forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
I.IntValue forall a b. (a -> b) -> a -> b
$ forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
v]
    E.Scalar (E.Prim (E.FloatType FloatType
ft)) ->
      forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
I.Constant forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
I.FloatValue forall a b. (a -> b) -> a -> b
$ forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
ft Integer
v]
    PatType
_ -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: nonsensical type for integer literal: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PatType
t
internaliseExp [Char]
_ (E.FloatLit Double
v (Info PatType
t) SrcLoc
_) =
  case PatType
t of
    E.Scalar (E.Prim (E.FloatType FloatType
ft)) ->
      forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
I.Constant forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
I.FloatValue forall a b. (a -> b) -> a -> b
$ forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
ft Double
v]
    PatType
_ -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: nonsensical type for float literal: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PatType
t
-- Builtin operators are handled specially because they are
-- overloaded.
internaliseExp [Char]
desc (E.Project Name
k Exp
e (Info PatType
rt) SrcLoc
_) = do
  let i' :: Int
i' = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall als. TypeBase Size als -> Int
internalisedTypeSize forall a b. (a -> b) -> a -> b
$
        case Exp -> PatType
E.typeOf Exp
e forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` () of
          E.Scalar (Record Map Name (TypeBase Size ())
fs) ->
            forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((forall a. Eq a => a -> a -> Bool
/= Name
k) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) forall a b. (a -> b) -> a -> b
$ forall a. Map Name a -> [(Name, a)]
sortFields Map Name (TypeBase Size ())
fs
          TypeBase Size ()
t -> [TypeBase Size ()
t]
  forall a. Int -> [a] -> [a]
take (forall als. TypeBase Size als -> Int
internalisedTypeSize forall a b. (a -> b) -> a -> b
$ PatType
rt forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` ()) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Int -> [a] -> [a]
drop Int
i'
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
internaliseExp [Char]
_ e :: Exp
e@E.Lambda {} =
  forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected lambda at " forall a. [a] -> [a] -> [a]
++ forall a. Located a => a -> [Char]
locStr (forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp [Char]
_ e :: Exp
e@E.OpSection {} =
  forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected operator section at " forall a. [a] -> [a] -> [a]
++ forall a. Located a => a -> [Char]
locStr (forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp [Char]
_ e :: Exp
e@E.OpSectionLeft {} =
  forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected left operator section at " forall a. [a] -> [a] -> [a]
++ forall a. Located a => a -> [Char]
locStr (forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp [Char]
_ e :: Exp
e@E.OpSectionRight {} =
  forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected right operator section at " forall a. [a] -> [a] -> [a]
++ forall a. Located a => a -> [Char]
locStr (forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp [Char]
_ e :: Exp
e@E.ProjectSection {} =
  forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected projection section at " forall a. [a] -> [a] -> [a]
++ forall a. Located a => a -> [Char]
locStr (forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp [Char]
_ e :: Exp
e@E.IndexSection {} =
  forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected index section at " forall a. [a] -> [a] -> [a]
++ forall a. Located a => a -> [Char]
locStr (forall a. Located a => a -> SrcLoc
srclocOf Exp
e)

internaliseArg :: String -> (E.Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg :: [Char] -> (Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg [Char]
desc (Exp
arg, Maybe VName
argdim) = do
  Scope SOACS
exists <- forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  case Maybe VName
argdim of
    Just VName
d | VName
d forall k a. Ord k => k -> Map k a -> Bool
`M.member` Scope SOACS
exists -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [VName -> SubExp
I.Var VName
d]
    Maybe VName
_ -> do
      [SubExp]
arg' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
arg
      case ([SubExp]
arg', Maybe VName
argdim) of
        ([SubExp
se], Just VName
d) -> do
          forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
d] forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
        ([SubExp], Maybe VName)
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
arg'

internalisePatLit :: E.PatLit -> E.PatType -> I.PrimValue
internalisePatLit :: PatLit -> PatType -> PrimValue
internalisePatLit (E.PatLitPrim PrimValue
v) PatType
_ =
  PrimValue -> PrimValue
internalisePrimValue PrimValue
v
internalisePatLit (E.PatLitInt Integer
x) (E.Scalar (E.Prim (E.Signed IntType
it))) =
  IntValue -> PrimValue
I.IntValue forall a b. (a -> b) -> a -> b
$ forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
x
internalisePatLit (E.PatLitInt Integer
x) (E.Scalar (E.Prim (E.Unsigned IntType
it))) =
  IntValue -> PrimValue
I.IntValue forall a b. (a -> b) -> a -> b
$ forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
x
internalisePatLit (E.PatLitFloat Double
x) (E.Scalar (E.Prim (E.FloatType FloatType
ft))) =
  FloatValue -> PrimValue
I.FloatValue forall a b. (a -> b) -> a -> b
$ forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
ft Double
x
internalisePatLit PatLit
l PatType
t =
  forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Nonsensical pattern and type: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show (PatLit
l, PatType
t)

generateCond ::
  E.Pat ->
  [I.SubExp] ->
  InternaliseM ([Maybe I.PrimValue], [I.SubExp])
generateCond :: PatBase Info VName
-> [SubExp] -> InternaliseM ([Maybe PrimValue], [SubExp])
generateCond PatBase Info VName
orig_p [SubExp]
orig_ses = do
  ([Maybe PrimValue]
cmps, [SubExp]
pertinent, [SubExp]
_) <- forall {vn} {a}.
(Eq vn, IsName vn) =>
PatBase Info vn
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares PatBase Info VName
orig_p [SubExp]
orig_ses
  forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Maybe PrimValue]
cmps, [SubExp]
pertinent)
  where
    compares :: PatBase Info vn
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares (E.PatLit PatLit
l (Info PatType
t) SrcLoc
_) (a
se : [a]
ses) =
      forall (f :: * -> *) a. Applicative f => a -> f a
pure ([forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ PatLit -> PatType -> PrimValue
internalisePatLit PatLit
l PatType
t], [a
se], [a]
ses)
    compares (E.PatConstr Name
c (Info (E.Scalar (E.Sum Map Name [PatType]
fs))) [PatBase Info vn]
pats SrcLoc
_) (a
_ : [a]
ses) = do
      ([TypeBase ExtShape Uniqueness]
payload_ts, Map Name (Int, [Int])
m) <- Map Name [TypeBase Size ()]
-> InternaliseM
     ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
internaliseSumType forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> b) -> Map k a -> Map k b
M.map (forall a b. (a -> b) -> [a] -> [b]
map forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct) Map Name [PatType]
fs
      case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
c Map Name (Int, [Int])
m of
        Just (Int
tag, [Int]
payload_is) -> do
          let ([a]
payload_ses, [a]
ses') = forall a. Int -> [a] -> ([a], [a])
splitAt (forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase ExtShape Uniqueness]
payload_ts) [a]
ses
          ([Maybe PrimValue]
cmps, [a]
pertinent, [a]
_) <-
            [PatBase Info vn]
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
comparesMany [PatBase Info vn]
pats forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map ([a]
payload_ses !!) [Int]
payload_is
          let missingCmps :: Int -> a -> Maybe PrimValue
missingCmps Int
i a
_ =
                case Int
i forall a. Eq a => a -> [a] -> Maybe Int
`elemIndex` [Int]
payload_is of
                  Just Int
j -> [Maybe PrimValue]
cmps forall a. [a] -> Int -> a
!! Int
j
                  Maybe Int
Nothing -> forall a. Maybe a
Nothing
          forall (f :: * -> *) a. Applicative f => a -> f a
pure
            ( forall a. a -> Maybe a
Just (IntValue -> PrimValue
I.IntValue forall a b. (a -> b) -> a -> b
$ forall int. Integral int => IntType -> int -> IntValue
intValue IntType
Int8 forall a b. (a -> b) -> a -> b
$ forall a. Integral a => a -> Integer
toInteger Int
tag)
                forall a. a -> [a] -> [a]
: forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> a -> Maybe PrimValue
missingCmps [Int
0 ..] [a]
payload_ses,
              [a]
pertinent,
              [a]
ses'
            )
        Maybe (Int, [Int])
Nothing ->
          forall a. HasCallStack => [Char] -> a
error [Char]
"generateCond: missing constructor"
    compares (E.PatConstr Name
_ (Info PatType
t) [PatBase Info vn]
_ SrcLoc
_) [a]
_ =
      forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"generateCond: PatConstr has nonsensical type: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PatType
t
    compares (E.Id vn
_ Info PatType
t SrcLoc
loc) [a]
ses =
      PatBase Info vn
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares (forall (f :: * -> *) vn. f PatType -> SrcLoc -> PatBase f vn
E.Wildcard Info PatType
t SrcLoc
loc) [a]
ses
    compares (E.Wildcard (Info PatType
t) SrcLoc
_) [a]
ses = do
      let ([a]
id_ses, [a]
rest_ses) = forall a. Int -> [a] -> ([a], [a])
splitAt (forall als. TypeBase Size als -> Int
internalisedTypeSize forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatType
t) [a]
ses
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a b. (a -> b) -> [a] -> [b]
map (forall a b. a -> b -> a
const forall a. Maybe a
Nothing) [a]
id_ses, [a]
id_ses, [a]
rest_ses)
    compares (E.PatParens PatBase Info vn
pat SrcLoc
_) [a]
ses =
      PatBase Info vn
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares PatBase Info vn
pat [a]
ses
    compares (E.PatAttr AttrInfo vn
_ PatBase Info vn
pat SrcLoc
_) [a]
ses =
      PatBase Info vn
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares PatBase Info vn
pat [a]
ses
    compares (E.TuplePat [] SrcLoc
loc) [a]
ses =
      PatBase Info vn
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares (forall (f :: * -> *) vn. f PatType -> SrcLoc -> PatBase f vn
E.Wildcard (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
E.Record forall a. Monoid a => a
mempty) SrcLoc
loc) [a]
ses
    compares (E.RecordPat [] SrcLoc
loc) [a]
ses =
      PatBase Info vn
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares (forall (f :: * -> *) vn. f PatType -> SrcLoc -> PatBase f vn
E.Wildcard (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
E.Record forall a. Monoid a => a
mempty) SrcLoc
loc) [a]
ses
    compares (E.TuplePat [PatBase Info vn]
pats SrcLoc
_) [a]
ses =
      [PatBase Info vn]
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
comparesMany [PatBase Info vn]
pats [a]
ses
    compares (E.RecordPat [(Name, PatBase Info vn)]
fs SrcLoc
_) [a]
ses =
      [PatBase Info vn]
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
comparesMany (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall a. Map Name a -> [(Name, a)]
E.sortFields forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, PatBase Info vn)]
fs) [a]
ses
    compares (E.PatAscription PatBase Info vn
pat TypeExp vn
_ SrcLoc
_) [a]
ses =
      PatBase Info vn
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares PatBase Info vn
pat [a]
ses
    compares PatBase Info vn
pat [] =
      forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"generateCond: No values left for pattern " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PatBase Info vn
pat

    comparesMany :: [PatBase Info vn]
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
comparesMany [] [a]
ses = forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], [], [a]
ses)
    comparesMany (PatBase Info vn
pat : [PatBase Info vn]
pats) [a]
ses = do
      ([Maybe PrimValue]
cmps1, [a]
pertinent1, [a]
ses') <- PatBase Info vn
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares PatBase Info vn
pat [a]
ses
      ([Maybe PrimValue]
cmps2, [a]
pertinent2, [a]
ses'') <- [PatBase Info vn]
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
comparesMany [PatBase Info vn]
pats [a]
ses'
      forall (f :: * -> *) a. Applicative f => a -> f a
pure
        ( [Maybe PrimValue]
cmps1 forall a. Semigroup a => a -> a -> a
<> [Maybe PrimValue]
cmps2,
          [a]
pertinent1 forall a. Semigroup a => a -> a -> a
<> [a]
pertinent2,
          [a]
ses''
        )

internalisePat ::
  String ->
  [E.SizeBinder VName] ->
  E.Pat ->
  E.Exp ->
  InternaliseM a ->
  InternaliseM a
internalisePat :: forall a.
[Char]
-> [SizeBinder VName]
-> PatBase Info VName
-> Exp
-> InternaliseM a
-> InternaliseM a
internalisePat [Char]
desc [SizeBinder VName]
sizes PatBase Info VName
p Exp
e InternaliseM a
m = do
  [SubExp]
ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc' Exp
e
  forall a.
[SizeBinder VName]
-> PatBase Info VName
-> [SubExp]
-> InternaliseM a
-> InternaliseM a
internalisePat' [SizeBinder VName]
sizes PatBase Info VName
p [SubExp]
ses InternaliseM a
m
  where
    desc' :: [Char]
desc' = case forall a. Set a -> [a]
S.toList forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set (IdentBase f vn)
E.patIdents PatBase Info VName
p of
      [IdentBase Info VName
v] -> VName -> [Char]
baseString forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. IdentBase f vn -> vn
E.identName IdentBase Info VName
v
      [IdentBase Info VName]
_ -> [Char]
desc

internalisePat' ::
  [E.SizeBinder VName] ->
  E.Pat ->
  [I.SubExp] ->
  InternaliseM a ->
  InternaliseM a
internalisePat' :: forall a.
[SizeBinder VName]
-> PatBase Info VName
-> [SubExp]
-> InternaliseM a
-> InternaliseM a
internalisePat' [SizeBinder VName]
sizes PatBase Info VName
p [SubExp]
ses InternaliseM a
m = do
  [TypeBase Shape NoUniqueness]
ses_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
ses
  forall a.
PatBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([VName] -> InternaliseM a)
-> InternaliseM a
stmPat PatBase Info VName
p [TypeBase Shape NoUniqueness]
ses_ts forall a b. (a -> b) -> a -> b
$ \[VName]
pat_names -> do
    AppRes -> [SubExp] -> InternaliseM ()
bindExtSizes (PatType -> [VName] -> AppRes
AppRes (PatBase Info VName -> PatType
E.patternType PatBase Info VName
p) (forall a b. (a -> b) -> [a] -> [b]
map forall vn. SizeBinder vn -> vn
E.sizeName [SizeBinder VName]
sizes)) [SubExp]
ses
    forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
pat_names [SubExp]
ses) forall a b. (a -> b) -> a -> b
$ \(VName
v, SubExp
se) ->
      forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
v] forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
se
    InternaliseM a
m

internaliseSlice ::
  SrcLoc ->
  [SubExp] ->
  [E.DimIndex] ->
  InternaliseM ([I.DimIndex SubExp], Certs)
internaliseSlice :: SrcLoc
-> [SubExp]
-> SliceBase Info VName
-> InternaliseM ([DimIndex SubExp], Certs)
internaliseSlice SrcLoc
loc [SubExp]
dims SliceBase Info VName
idxs = do
  ([DimIndex SubExp]
idxs', [SubExp]
oks, [[ErrorMsgPart SubExp]]
parts) <- forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM SubExp
-> DimIndex
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
internaliseDimIndex [SubExp]
dims SliceBase Info VName
idxs
  SubExp
ok <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"index_ok" forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
eAll [SubExp]
oks
  let msg :: ErrorMsg SubExp
msg =
        forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg forall a b. (a -> b) -> a -> b
$
          [ErrorMsgPart SubExp
"Index ["]
            forall a. [a] -> [a] -> [a]
++ forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
", "] [[ErrorMsgPart SubExp]]
parts
            forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"] out of bounds for array of shape ["]
            forall a. [a] -> [a] -> [a]
++ forall a. a -> [a] -> [a]
intersperse ErrorMsgPart SubExp
"][" (forall a b. (a -> b) -> [a] -> [b]
map (forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64) forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> [a]
take (forall (t :: * -> *) a. Foldable t => t a -> Int
length SliceBase Info VName
idxs) [SubExp]
dims)
            forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"]."]
  Certs
c <- [Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert [Char]
"index_certs" SubExp
ok ErrorMsg SubExp
msg SrcLoc
loc
  forall (f :: * -> *) a. Applicative f => a -> f a
pure ([DimIndex SubExp]
idxs', Certs
c)

internaliseDimIndex ::
  SubExp ->
  E.DimIndex ->
  InternaliseM (I.DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
internaliseDimIndex :: SubExp
-> DimIndex
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
internaliseDimIndex SubExp
w (E.DimFix Exp
i) = do
  (SubExp
i', IntType
_) <- [Char] -> Exp -> InternaliseM (SubExp, IntType)
internaliseSizeExp [Char]
"i" Exp
i
  let lowerBound :: Exp SOACS
lowerBound =
        forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
I.Int64) (forall v. IsValue v => v -> SubExp
I.constant (Int64
0 :: I.Int64)) SubExp
i'
      upperBound :: Exp SOACS
upperBound =
        forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSlt IntType
I.Int64) SubExp
i' SubExp
w
  SubExp
ok <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"bounds_check" forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
MonadBuilder m =>
BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eBinOp BinOp
I.LogAnd (forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp SOACS
lowerBound) (forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp SOACS
upperBound)
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall d. d -> DimIndex d
I.DimFix SubExp
i', SubExp
ok, [forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
i'])

-- Special-case an important common case that otherwise leads to horrible code.
internaliseDimIndex
  SubExp
w
  ( E.DimSlice
      Maybe Exp
Nothing
      Maybe Exp
Nothing
      (Just (E.Negate (E.IntLit Integer
1 Info PatType
_ SrcLoc
_) SrcLoc
_))
    ) = do
    SubExp
w_minus_1 <-
      forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"w_minus_1" forall a b. (a -> b) -> a -> b
$
        forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$
          BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) SubExp
w SubExp
one
    forall (f :: * -> *) a. Applicative f => a -> f a
pure
      ( forall d. d -> d -> d -> DimIndex d
I.DimSlice SubExp
w_minus_1 SubExp
w forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 (-Integer
1),
        forall v. IsValue v => v -> SubExp
constant Bool
True,
        forall a. Monoid a => a
mempty
      )
    where
      one :: SubExp
one = forall v. IsValue v => v -> SubExp
constant (Int64
1 :: Int64)
internaliseDimIndex SubExp
w (E.DimSlice Maybe Exp
i Maybe Exp
j Maybe Exp
s) = do
  SubExp
s' <- forall b a. b -> (a -> b) -> Maybe a -> b
maybe (forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
one) (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp -> InternaliseM (SubExp, IntType)
internaliseSizeExp [Char]
"s") Maybe Exp
s
  SubExp
s_sign <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"s_sign" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (IntType -> UnOp
I.SSignum IntType
Int64) SubExp
s'
  SubExp
backwards <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"backwards" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
int64) SubExp
s_sign SubExp
negone
  SubExp
w_minus_1 <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"w_minus_1" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) SubExp
w SubExp
one
  let i_def :: InternaliseM SubExp
i_def =
        forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"i_def"
          forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
            (forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
backwards)
            (forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
w_minus_1])
            (forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
zero])
      j_def :: InternaliseM SubExp
j_def =
        forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"j_def"
          forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
            (forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
backwards)
            (forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
negone])
            (forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
w])
  SubExp
i' <- forall b a. b -> (a -> b) -> Maybe a -> b
maybe InternaliseM SubExp
i_def (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp -> InternaliseM (SubExp, IntType)
internaliseSizeExp [Char]
"i") Maybe Exp
i
  SubExp
j' <- forall b a. b -> (a -> b) -> Maybe a -> b
maybe InternaliseM SubExp
j_def (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp -> InternaliseM (SubExp, IntType)
internaliseSizeExp [Char]
"j") Maybe Exp
j
  SubExp
j_m_i <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"j_m_i" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) SubExp
j' SubExp
i'
  -- Something like a division-rounding-up, but accomodating negative
  -- operands.
  let divRounding :: InternaliseM (Exp SOACS)
-> InternaliseM (Exp SOACS)
-> InternaliseM (Exp (Rep InternaliseM))
divRounding InternaliseM (Exp SOACS)
x InternaliseM (Exp SOACS)
y =
        forall (m :: * -> *).
MonadBuilder m =>
BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eBinOp
          (IntType -> Safety -> BinOp
SQuot IntType
Int64 Safety
Safe)
          ( forall (m :: * -> *).
MonadBuilder m =>
BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eBinOp
              (IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
I.OverflowWrap)
              InternaliseM (Exp SOACS)
x
              (forall (m :: * -> *).
MonadBuilder m =>
BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eBinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) InternaliseM (Exp SOACS)
y (forall (m :: * -> *).
MonadBuilder m =>
m (Exp (Rep m)) -> m (Exp (Rep m))
eSignum forall a b. (a -> b) -> a -> b
$ forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
toExp SubExp
s'))
          )
          InternaliseM (Exp SOACS)
y
  SubExp
n <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"n" forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp SOACS)
-> InternaliseM (Exp SOACS)
-> InternaliseM (Exp (Rep InternaliseM))
divRounding (forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
toExp SubExp
j_m_i) (forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
toExp SubExp
s')

  SubExp
zero_stride <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"zero_stride" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
CmpEq PrimType
int64) SubExp
s_sign SubExp
zero
  SubExp
nonzero_stride <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"nonzero_stride" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
I.Not SubExp
zero_stride

  -- 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 <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"empty_slice" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
CmpEq PrimType
int64) SubExp
n SubExp
zero

  SubExp
m <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"m" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) SubExp
n SubExp
one
  SubExp
m_t_s <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"m_t_s" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Mul IntType
Int64 Overflow
I.OverflowWrap) SubExp
m SubExp
s'
  SubExp
i_p_m_t_s <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"i_p_m_t_s" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
I.OverflowWrap) SubExp
i' SubExp
m_t_s
  SubExp
zero_leq_i_p_m_t_s <-
    forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"zero_leq_i_p_m_t_s" forall a b. (a -> b) -> a -> b
$
      forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
        CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int64) SubExp
zero SubExp
i_p_m_t_s
  SubExp
i_p_m_t_s_leq_w <-
    forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"i_p_m_t_s_leq_w" forall a b. (a -> b) -> a -> b
$
      forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
        CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int64) SubExp
i_p_m_t_s SubExp
w
  SubExp
i_p_m_t_s_lth_w <-
    forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"i_p_m_t_s_leq_w" forall a b. (a -> b) -> a -> b
$
      forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
        CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSlt IntType
Int64) SubExp
i_p_m_t_s SubExp
w

  SubExp
zero_lte_i <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"zero_lte_i" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int64) SubExp
zero SubExp
i'
  SubExp
i_lte_j <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"i_lte_j" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int64) SubExp
i' SubExp
j'
  SubExp
forwards_ok <-
    forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"forwards_ok"
      forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep 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 <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"negone_lte_j" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int64) SubExp
negone SubExp
j'
  SubExp
j_lte_i <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"j_lte_i" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int64) SubExp
j' SubExp
i'
  SubExp
backwards_ok <-
    forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"backwards_ok"
      forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep 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 <-
    forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"slice_ok"
      forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
        (forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
backwards)
        (forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
backwards_ok])
        (forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
forwards_ok])

  SubExp
ok_or_empty <-
    forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"ok_or_empty" forall a b. (a -> b) -> a -> b
$
      forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
        BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogOr SubExp
empty_slice SubExp
slice_ok

  SubExp
acceptable <-
    forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"slice_acceptable" forall a b. (a -> b) -> a -> b
$
      forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
        BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogAnd SubExp
nonzero_stride SubExp
ok_or_empty

  let parts :: [ErrorMsgPart SubExp]
parts = case (Maybe Exp
i, Maybe Exp
j, Maybe Exp
s) of
        (Maybe Exp
_, Maybe Exp
_, Just {}) ->
          [ forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
i') Maybe Exp
i,
            ErrorMsgPart SubExp
":",
            forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
j') Maybe Exp
j,
            ErrorMsgPart SubExp
":",
            forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
s'
          ]
        (Maybe Exp
_, Just {}, Maybe Exp
_) ->
          [ forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
i') Maybe Exp
i,
            ErrorMsgPart SubExp
":",
            forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
j'
          ]
            forall a. [a] -> [a] -> [a]
++ forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty (forall a b. a -> b -> a
const [ErrorMsgPart SubExp
":", forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
s']) Maybe Exp
s
        (Maybe Exp
_, Maybe Exp
Nothing, Maybe Exp
Nothing) ->
          [forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
i', ErrorMsgPart SubExp
":"]
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall d. d -> d -> d -> DimIndex d
I.DimSlice SubExp
i' SubExp
n SubExp
s', SubExp
acceptable, [ErrorMsgPart SubExp]
parts)
  where
    zero :: SubExp
zero = forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)
    negone :: SubExp
negone = forall v. IsValue v => v -> SubExp
constant (-Int64
1 :: Int64)
    one :: SubExp
one = forall v. IsValue v => v -> SubExp
constant (Int64
1 :: Int64)

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

internaliseHist ::
  Int ->
  String ->
  E.Exp ->
  E.Exp ->
  E.Exp ->
  E.Exp ->
  E.Exp ->
  E.Exp ->
  SrcLoc ->
  InternaliseM [SubExp]
internaliseHist :: Int
-> [Char]
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> SrcLoc
-> InternaliseM [SubExp]
internaliseHist Int
dim [Char]
desc Exp
rf Exp
hist Exp
op Exp
ne Exp
buckets Exp
img SrcLoc
loc = do
  SubExp
rf' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"hist_rf" Exp
rf
  [SubExp]
ne' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"hist_ne" Exp
ne
  [VName]
hist' <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"hist_hist" Exp
hist
  [VName]
buckets' <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"hist_buckets" Exp
buckets
  [VName]
img' <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"hist_img" Exp
img

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

  -- reshape return type of bucket function to have same size as neutral element
  -- (modulo the index)
  [Param (TypeBase Shape NoUniqueness)]
bucket_params <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
dim (forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"bucket_p" forall a b. (a -> b) -> a -> b
$ forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
  [Param (TypeBase Shape NoUniqueness)]
img_params <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"img_p" forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall u. TypeBase Shape u -> TypeBase Shape u
rowType) forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
img'
  let params :: [Param (TypeBase Shape NoUniqueness)]
params = [Param (TypeBase Shape NoUniqueness)]
bucket_params forall a. [a] -> [a] -> [a]
++ [Param (TypeBase Shape NoUniqueness)]
img_params
      rettype :: [TypeBase Shape NoUniqueness]
rettype = forall a. Int -> a -> [a]
replicate Int
dim (forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64) forall a. [a] -> [a] -> [a]
++ [TypeBase Shape NoUniqueness]
ne_ts
      body :: Body SOACS
body = forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$ [VName] -> Result
varsRes forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
I.paramName [Param (TypeBase Shape NoUniqueness)]
params
  Lambda SOACS
lam' <-
    forall (m :: * -> *).
MonadBuilder m =>
[LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
mkLambda [Param (TypeBase Shape NoUniqueness)]
params forall a b. (a -> b) -> a -> b
$
      ErrorMsg SubExp
-> SrcLoc
-> [TypeBase Shape NoUniqueness]
-> Result
-> InternaliseM Result
ensureResultShape
        ErrorMsg SubExp
"Row shape of value array does not match row shape of hist target"
        (forall a. Located a => a -> SrcLoc
srclocOf Exp
img)
        [TypeBase Shape NoUniqueness]
rettype
        forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind Body SOACS
body

  -- get sizes of histogram and image arrays
  Shape
shape_hist <- forall d. [d] -> ShapeBase d
I.Shape forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Int -> [a] -> [a]
take Int
dim forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall u. TypeBase Shape u -> [SubExp]
I.arrayDims forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType (forall a. [a] -> a
head [VName]
hist')
  SubExp
w_img <- forall u. Int -> TypeBase Shape u -> SubExp
I.arraySize Int
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType (forall a. [a] -> a
head [VName]
img')

  [Char] -> Exp SOACS -> InternaliseM [SubExp]
letValExp' [Char]
desc forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. Op rep -> Exp rep
I.Op forall a b. (a -> b) -> a -> b
$
    forall rep.
SubExp -> [VName] -> [HistOp rep] -> Lambda rep -> SOAC rep
I.Hist SubExp
w_img ([VName]
buckets' forall a. [a] -> [a] -> [a]
++ [VName]
img') [forall rep.
Shape -> SubExp -> [VName] -> [SubExp] -> Lambda rep -> HistOp rep
HistOp Shape
shape_hist SubExp
rf' [VName]
hist' [SubExp]
ne_shp Lambda SOACS
op'] Lambda SOACS
lam'

internaliseStreamAcc ::
  String ->
  E.Exp ->
  Maybe (E.Exp, E.Exp) ->
  E.Exp ->
  E.Exp ->
  InternaliseM [SubExp]
internaliseStreamAcc :: [Char]
-> Exp -> Maybe (Exp, Exp) -> Exp -> Exp -> InternaliseM [SubExp]
internaliseStreamAcc [Char]
desc Exp
dest Maybe (Exp, Exp)
op Exp
lam Exp
bs = do
  [VName]
dest' <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"scatter_dest" Exp
dest
  [VName]
bs' <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"scatter_input" Exp
bs

  VName
acc_cert_v <- forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"acc_cert"
  [TypeBase Shape NoUniqueness]
dest_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
dest'
  let dest_w :: SubExp
dest_w = forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [TypeBase Shape NoUniqueness]
dest_ts
      acc_t :: TypeBase Shape NoUniqueness
acc_t = forall shape u.
VName
-> Shape -> [TypeBase Shape NoUniqueness] -> u -> TypeBase shape u
Acc VName
acc_cert_v (forall d. [d] -> ShapeBase d
I.Shape [SubExp
dest_w]) (forall a b. (a -> b) -> [a] -> [b]
map forall u. TypeBase Shape u -> TypeBase Shape u
rowType [TypeBase Shape NoUniqueness]
dest_ts) NoUniqueness
NoUniqueness
  Param (TypeBase Shape NoUniqueness)
acc_p <- forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"acc_p" TypeBase Shape NoUniqueness
acc_t
  Lambda SOACS
withacc_lam <- forall (m :: * -> *).
MonadBuilder m =>
[LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
mkLambda [forall dec. Attrs -> VName -> dec -> Param dec
Param forall a. Monoid a => a
mempty VName
acc_cert_v (forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.Unit), Param (TypeBase Shape NoUniqueness)
acc_p] forall a b. (a -> b) -> a -> b
$ do
    [TypeBase Shape NoUniqueness]
bs_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
bs'
    Lambda SOACS
lam' <- Exp -> [TypeBase Shape NoUniqueness] -> InternaliseM (Lambda SOACS)
internaliseLambdaCoerce Exp
lam forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall u. TypeBase Shape u -> TypeBase Shape u
rowType forall a b. (a -> b) -> a -> b
$ forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
paramType Param (TypeBase Shape NoUniqueness)
acc_p forall a. a -> [a] -> [a]
: [TypeBase Shape NoUniqueness]
bs_ts
    let w :: SubExp
w = forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [TypeBase Shape NoUniqueness]
bs_ts
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [SubExp] -> Result
subExpsRes forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp SOACS -> InternaliseM [SubExp]
letValExp' [Char]
"acc_res" forall a b. (a -> b) -> a -> b
$
      forall rep. Op rep -> Exp rep
I.Op forall a b. (a -> b) -> a -> b
$
        forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w (forall dec. Param dec -> VName
I.paramName Param (TypeBase Shape NoUniqueness)
acc_p forall a. a -> [a] -> [a]
: [VName]
bs') (forall rep. Lambda rep -> ScremaForm rep
I.mapSOAC Lambda SOACS
lam')

  Maybe (Lambda SOACS, [SubExp])
op' <-
    case Maybe (Exp, Exp)
op of
      Just (Exp
op_lam, Exp
ne) -> do
        [SubExp]
ne' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"hist_ne" Exp
ne
        [TypeBase Shape NoUniqueness]
ne_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
I.subExpType [SubExp]
ne'
        ([Param (TypeBase Shape NoUniqueness)]
lam_params, Body SOACS
lam_body, [TypeBase Shape NoUniqueness]
lam_rettype) <-
          InternaliseLambda
internaliseLambda Exp
op_lam forall a b. (a -> b) -> a -> b
$ [TypeBase Shape NoUniqueness]
ne_ts forall a. [a] -> [a] -> [a]
++ [TypeBase Shape NoUniqueness]
ne_ts
        Param (TypeBase Shape NoUniqueness)
idxp <- forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"idx" forall a b. (a -> b) -> a -> b
$ forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64
        let op_lam' :: Lambda SOACS
op_lam' = forall rep.
[LParam rep]
-> Body rep -> [TypeBase Shape NoUniqueness] -> Lambda rep
I.Lambda (Param (TypeBase Shape NoUniqueness)
idxp forall a. a -> [a] -> [a]
: [Param (TypeBase Shape NoUniqueness)]
lam_params) Body SOACS
lam_body [TypeBase Shape NoUniqueness]
lam_rettype
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (Lambda SOACS
op_lam', [SubExp]
ne')
      Maybe (Exp, Exp)
Nothing ->
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing

  SubExp
destw <- forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
dest'
  forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var) forall a b. (a -> b) -> a -> b
$
    forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [VName]
letTupExp [Char]
desc forall a b. (a -> b) -> a -> b
$
      forall rep. [WithAccInput rep] -> Lambda rep -> Exp rep
WithAcc [(forall d. [d] -> ShapeBase d
I.Shape [SubExp
destw], [VName]
dest', Maybe (Lambda SOACS, [SubExp])
op')] Lambda SOACS
withacc_lam

internaliseExp1 :: String -> E.Exp -> InternaliseM I.SubExp
internaliseExp1 :: [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
desc Exp
e = do
  [SubExp]
vs <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
  case [SubExp]
vs of
    [SubExp
se] -> forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
se
    [SubExp]
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"Internalise.internaliseExp1: was passed not just a single subexpression"

-- | Promote to dimension type as appropriate for the original type.
-- Also return original type.
internaliseSizeExp :: String -> E.Exp -> InternaliseM (I.SubExp, IntType)
internaliseSizeExp :: [Char] -> Exp -> InternaliseM (SubExp, IntType)
internaliseSizeExp [Char]
s Exp
e = do
  SubExp
e' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
s Exp
e
  case Exp -> PatType
E.typeOf Exp
e of
    E.Scalar (E.Prim (E.Signed IntType
it)) -> (,IntType
it) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
e'
    PatType
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"internaliseSizeExp: bad type"

internaliseExpToVars :: String -> E.Exp -> InternaliseM [I.VName]
internaliseExpToVars :: [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
desc Exp
e =
  forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM VName
asIdent forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
  where
    asIdent :: SubExp -> InternaliseM VName
asIdent (I.Var VName
v) = forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v
    asIdent SubExp
se = forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp 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 :: [Char]
-> Exp -> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
internaliseOperation [Char]
s Exp
e VName -> InternaliseM BasicOp
op = do
  [VName]
vs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
s Exp
e
  forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
s forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. BasicOp -> Exp rep
I.BasicOp forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< VName -> InternaliseM BasicOp
op) [VName]
vs

certifyingNonzero ::
  SrcLoc ->
  IntType ->
  SubExp ->
  InternaliseM a ->
  InternaliseM a
certifyingNonzero :: forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
x InternaliseM a
m = do
  SubExp
zero <-
    forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"zero" forall a b. (a -> b) -> a -> b
$
      forall rep. BasicOp -> Exp rep
I.BasicOp 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 <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"nonzero" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
UnOp UnOp
I.Not SubExp
zero
  Certs
c <- [Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert [Char]
"nonzero_cert" SubExp
nonzero ErrorMsg SubExp
"division by zero" SrcLoc
loc
  forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c InternaliseM a
m

certifyingNonnegative ::
  SrcLoc ->
  IntType ->
  SubExp ->
  InternaliseM a ->
  InternaliseM a
certifyingNonnegative :: forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonnegative SrcLoc
loc IntType
t SubExp
x InternaliseM a
m = do
  SubExp
nonnegative <-
    forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"nonnegative" forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. BasicOp -> Exp rep
I.BasicOp 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
  Certs
c <- [Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert [Char]
"nonzero_cert" SubExp
nonnegative ErrorMsg SubExp
"negative exponent" SrcLoc
loc
  forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c InternaliseM a
m

internaliseBinOp ::
  SrcLoc ->
  String ->
  E.BinOp ->
  I.SubExp ->
  I.SubExp ->
  E.PrimType ->
  E.PrimType ->
  InternaliseM [I.SubExp]
internaliseBinOp :: SrcLoc
-> [Char]
-> BinOp
-> SubExp
-> SubExp
-> PrimType
-> PrimType
-> InternaliseM [SubExp]
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.LogAnd SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc BinOp
I.LogAnd SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.LogOr SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc BinOp
I.LogOr SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Plus SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Overflow -> BinOp
I.Add IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Plus SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Overflow -> BinOp
I.Add IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Plus SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (FloatType -> BinOp
I.FAdd FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Minus SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Overflow -> BinOp
I.Sub IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Minus SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Overflow -> BinOp
I.Sub IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Minus SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (FloatType -> BinOp
I.FSub FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Times SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Overflow -> BinOp
I.Mul IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Times SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Overflow -> BinOp
I.Mul IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Times SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (FloatType -> BinOp
I.FMul FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc [Char]
desc BinOp
E.Divide SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y forall a b. (a -> b) -> a -> b
$
    [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Safety -> BinOp
I.SDiv IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc [Char]
desc BinOp
E.Divide SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y forall a b. (a -> b) -> a -> b
$
    [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Safety -> BinOp
I.UDiv IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Divide SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (FloatType -> BinOp
I.FDiv FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Pow SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (FloatType -> BinOp
I.FPow FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc [Char]
desc BinOp
E.Pow SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonnegative SrcLoc
loc IntType
t SubExp
y forall a b. (a -> b) -> a -> b
$
    [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.Pow IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Pow SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.Pow IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc [Char]
desc BinOp
E.Mod SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y forall a b. (a -> b) -> a -> b
$
    [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Safety -> BinOp
I.SMod IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc [Char]
desc BinOp
E.Mod SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y forall a b. (a -> b) -> a -> b
$
    [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Safety -> BinOp
I.UMod IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Mod SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (FloatType -> BinOp
I.FMod FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc [Char]
desc BinOp
E.Quot SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y forall a b. (a -> b) -> a -> b
$
    [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Safety -> BinOp
I.SQuot IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc [Char]
desc BinOp
E.Quot SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y forall a b. (a -> b) -> a -> b
$
    [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Safety -> BinOp
I.UDiv IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc [Char]
desc BinOp
E.Rem SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y forall a b. (a -> b) -> a -> b
$
    [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Safety -> BinOp
I.SRem IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc [Char]
desc BinOp
E.Rem SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y forall a b. (a -> b) -> a -> b
$
    [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Safety -> BinOp
I.UMod IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.ShiftR SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.AShr IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.ShiftR SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.LShr IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.ShiftL SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.Shl IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.ShiftL SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.Shl IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Band SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.And IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Band SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.And IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Xor SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.Xor IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Xor SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.Xor IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Bor SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.Or IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Bor SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.Or IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Equal SubExp
x SubExp
y PrimType
t PrimType
_ =
  [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (PrimType -> CmpOp
I.CmpEq forall a b. (a -> b) -> a -> b
$ PrimType -> PrimType
internalisePrimType PrimType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.NotEqual SubExp
x SubExp
y PrimType
t PrimType
_ = do
  SubExp
eq <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp ([Char]
desc forall a. [a] -> [a] -> [a]
++ [Char]
"true") forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq forall a b. (a -> b) -> a -> b
$ PrimType -> PrimType
internalisePrimType PrimType
t) SubExp
x SubExp
y
  forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
I.Not SubExp
eq
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Less SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (IntType -> CmpOp
I.CmpSlt IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Less SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (IntType -> CmpOp
I.CmpUlt IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Leq SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (IntType -> CmpOp
I.CmpSle IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Leq SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (IntType -> CmpOp
I.CmpUle IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Greater SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (IntType -> CmpOp
I.CmpSlt IntType
t) SubExp
y SubExp
x -- Note the swapped x and y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Greater SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (IntType -> CmpOp
I.CmpUlt IntType
t) SubExp
y SubExp
x -- Note the swapped x and y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Geq SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (IntType -> CmpOp
I.CmpSle IntType
t) SubExp
y SubExp
x -- Note the swapped x and y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Geq SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (IntType -> CmpOp
I.CmpUle IntType
t) SubExp
y SubExp
x -- Note the swapped x and y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Less SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (FloatType -> CmpOp
I.FCmpLt FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Leq SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (FloatType -> CmpOp
I.FCmpLe FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Greater SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (FloatType -> CmpOp
I.FCmpLt FloatType
t) SubExp
y SubExp
x -- Note the swapped x and y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Geq SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (FloatType -> CmpOp
I.FCmpLe FloatType
t) SubExp
y SubExp
x -- Note the swapped x and y

-- Relational operators for booleans.
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Less SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
  [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc CmpOp
I.CmpLlt SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Leq SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
  [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc CmpOp
I.CmpLle SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Greater SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
  [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc CmpOp
I.CmpLlt SubExp
y SubExp
x -- Note the swapped x and y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Geq SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
  [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc CmpOp
I.CmpLle SubExp
y SubExp
x -- Note the swapped x and y
internaliseBinOp SrcLoc
_ [Char]
_ BinOp
op SubExp
_ SubExp
_ PrimType
t1 PrimType
t2 =
  forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
    [Char]
"Invalid binary operator "
      forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString BinOp
op
      forall a. [a] -> [a] -> [a]
++ [Char]
" with operand types "
      forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PrimType
t1
      forall a. [a] -> [a] -> [a]
++ [Char]
", "
      forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PrimType
t2

simpleBinOp ::
  String ->
  I.BinOp ->
  I.SubExp ->
  I.SubExp ->
  InternaliseM [I.SubExp]
simpleBinOp :: [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc BinOp
bop SubExp
x SubExp
y =
  forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp 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 :: [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc CmpOp
op SubExp
x SubExp
y =
  forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp CmpOp
op SubExp
x SubExp
y

data Function
  = FunctionName (E.QualName VName)
  | FunctionHole E.PatType SrcLoc
  deriving (Int -> Function -> ShowS
[Function] -> ShowS
Function -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
showList :: [Function] -> ShowS
$cshowList :: [Function] -> ShowS
show :: Function -> [Char]
$cshow :: Function -> [Char]
showsPrec :: Int -> Function -> ShowS
$cshowsPrec :: Int -> Function -> ShowS
Show)

findFuncall :: E.AppExp -> (Function, [(E.Exp, Maybe VName)])
findFuncall :: AppExp -> (Function, [(Exp, Maybe VName)])
findFuncall (E.Apply Exp
f Exp
arg (Info (Diet
_, Maybe VName
argext)) SrcLoc
_)
  | E.AppExp AppExp
f_e Info AppRes
_ <- Exp
f =
      let (Function
f_e', [(Exp, Maybe VName)]
args) = AppExp -> (Function, [(Exp, Maybe VName)])
findFuncall AppExp
f_e
       in (Function
f_e', [(Exp, Maybe VName)]
args forall a. [a] -> [a] -> [a]
++ [(Exp
arg, Maybe VName
argext)])
  | E.Var QualName VName
fname Info PatType
_ SrcLoc
_ <- Exp
f =
      (QualName VName -> Function
FunctionName QualName VName
fname, [(Exp
arg, Maybe VName
argext)])
  | E.Hole (Info PatType
t) SrcLoc
loc <- Exp
f =
      (PatType -> SrcLoc -> Function
FunctionHole PatType
t SrcLoc
loc, [(Exp
arg, Maybe VName
argext)])
findFuncall AppExp
e =
  forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Invalid function expression in application:\n" forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString AppExp
e

-- The type of a body.  Watch out: this only works for the degenerate
-- case where the body does not already return its context.
bodyExtType :: Body SOACS -> InternaliseM [ExtType]
bodyExtType :: Body SOACS -> InternaliseM [ExtType]
bodyExtType (Body BodyDec SOACS
_ Stms SOACS
stms Result
res) =
  [VName] -> [ExtType] -> [ExtType]
existentialiseExtTypes (forall k a. Map k a -> [k]
M.keys Scope SOACS
stmsscope) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall u. [TypeBase Shape u] -> [TypeBase ExtShape u]
staticShapes
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *) a.
ExtendedScope rep m a -> Scope rep -> m a
extendedScope (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall t (m :: * -> *).
HasScope t m =>
SubExpRes -> m (TypeBase Shape NoUniqueness)
subExpResType Result
res) Scope SOACS
stmsscope
  where
    stmsscope :: Scope SOACS
stmsscope = forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stms SOACS
stms

internaliseLambda :: InternaliseLambda
internaliseLambda :: InternaliseLambda
internaliseLambda (E.Parens Exp
e SrcLoc
_) [TypeBase Shape NoUniqueness]
rowtypes =
  InternaliseLambda
internaliseLambda Exp
e [TypeBase Shape NoUniqueness]
rowtypes
internaliseLambda (E.Lambda [PatBase Info VName]
params Exp
body Maybe (TypeExp VName)
_ (Info (Set Alias
_, RetType [VName]
_ TypeBase Size ()
rettype)) SrcLoc
_) [TypeBase Shape NoUniqueness]
rowtypes =
  forall a.
[PatBase Info VName]
-> [TypeBase Shape NoUniqueness]
-> ([LParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLambdaParams [PatBase Info VName]
params [TypeBase Shape NoUniqueness]
rowtypes forall a b. (a -> b) -> a -> b
$ \[LParam SOACS]
params' -> do
    Body SOACS
body' <- [Char] -> Exp -> InternaliseM (Body SOACS)
internaliseBody [Char]
"lam" Exp
body
    [TypeBase Shape NoUniqueness]
rettype' <- forall shape u.
TypeBase Size ()
-> [TypeBase shape u] -> InternaliseM [TypeBase Shape NoUniqueness]
internaliseLambdaReturnType TypeBase Size ()
rettype forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Body SOACS -> InternaliseM [ExtType]
bodyExtType Body SOACS
body'
    forall (f :: * -> *) a. Applicative f => a -> f a
pure ([LParam SOACS]
params', Body SOACS
body', [TypeBase Shape NoUniqueness]
rettype')
internaliseLambda Exp
e [TypeBase Shape NoUniqueness]
_ = forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseLambda: unexpected expression:\n" forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString Exp
e

internaliseLambdaCoerce :: E.Exp -> [Type] -> InternaliseM (I.Lambda SOACS)
internaliseLambdaCoerce :: Exp -> [TypeBase Shape NoUniqueness] -> InternaliseM (Lambda SOACS)
internaliseLambdaCoerce Exp
lam [TypeBase Shape NoUniqueness]
argtypes = do
  ([Param (TypeBase Shape NoUniqueness)]
params, Body SOACS
body, [TypeBase Shape NoUniqueness]
rettype) <- InternaliseLambda
internaliseLambda Exp
lam [TypeBase Shape NoUniqueness]
argtypes
  forall (m :: * -> *).
MonadBuilder m =>
[LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
mkLambda [Param (TypeBase Shape NoUniqueness)]
params forall a b. (a -> b) -> a -> b
$
    ErrorMsg SubExp
-> SrcLoc
-> [TypeBase Shape NoUniqueness]
-> Result
-> InternaliseM Result
ensureResultShape
      (forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg [forall a. Text -> ErrorMsgPart a
ErrorString Text
"unexpected lambda result size"])
      (forall a. Located a => a -> SrcLoc
srclocOf Exp
lam)
      [TypeBase Shape NoUniqueness]
rettype
      forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind Body SOACS
body

-- | Overloaded operators are treated here.
isOverloadedFunction ::
  E.QualName VName ->
  String ->
  SrcLoc ->
  Maybe ([(E.StructType, [SubExp])] -> InternaliseM [SubExp])
isOverloadedFunction :: QualName VName
-> [Char]
-> SrcLoc
-> Maybe ([(TypeBase Size (), [SubExp])] -> InternaliseM [SubExp])
isOverloadedFunction QualName VName
qname [Char]
desc SrcLoc
loc = do
  forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ VName -> Int
baseTag (forall vn. QualName vn -> vn
qualLeaf QualName VName
qname) forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag
  [Char]
-> Maybe ([(TypeBase Size (), [SubExp])] -> InternaliseM [SubExp])
handle forall a b. (a -> b) -> a -> b
$ VName -> [Char]
baseString forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> vn
qualLeaf QualName VName
qname
  where
    -- Handle equality and inequality specially, to treat the case of
    -- arrays.
    handle :: [Char]
-> Maybe ([(TypeBase Size (), [SubExp])] -> InternaliseM [SubExp])
handle [Char]
op
      | Just SubExp -> InternaliseM [SubExp]
cmp_f <- [Char] -> Maybe (SubExp -> InternaliseM [SubExp])
isEqlOp [Char]
op = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[(TypeBase Size ()
_, [SubExp]
xe'), (TypeBase Size ()
_, [SubExp]
ye')] -> do
          [SubExp]
rs <- forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM SubExp -> SubExp -> InternaliseM SubExp
doComparison [SubExp]
xe' [SubExp]
ye'
          SubExp -> InternaliseM [SubExp]
cmp_f forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"eq" forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
eAll [SubExp]
rs
      where
        isEqlOp :: [Char] -> Maybe (SubExp -> InternaliseM [SubExp])
isEqlOp [Char]
"!=" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \SubExp
eq ->
          forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
I.Not SubExp
eq
        isEqlOp [Char]
"==" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \SubExp
eq ->
          forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp
eq]
        isEqlOp [Char]
_ = forall a. Maybe a
Nothing

        doComparison :: SubExp -> SubExp -> InternaliseM SubExp
doComparison SubExp
x SubExp
y = do
          TypeBase Shape NoUniqueness
x_t <- forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
I.subExpType SubExp
x
          TypeBase Shape NoUniqueness
y_t <- forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
I.subExpType SubExp
y
          case TypeBase Shape NoUniqueness
x_t of
            I.Prim PrimType
t -> forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
t) SubExp
x SubExp
y
            TypeBase Shape NoUniqueness
_ -> do
              let x_dims :: [SubExp]
x_dims = forall u. TypeBase Shape u -> [SubExp]
I.arrayDims TypeBase Shape NoUniqueness
x_t
                  y_dims :: [SubExp]
y_dims = forall u. TypeBase Shape u -> [SubExp]
I.arrayDims TypeBase Shape NoUniqueness
y_t
              [SubExp]
dims_match <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
x_dims [SubExp]
y_dims) forall a b. (a -> b) -> a -> b
$ \(SubExp
x_dim, SubExp
y_dim) ->
                forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"dim_eq" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
int64) SubExp
x_dim SubExp
y_dim
              SubExp
shapes_match <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"shapes_match" forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
eAll [SubExp]
dims_match
              let compare_elems_body :: InternaliseM (Body SOACS)
compare_elems_body = forall rep (m :: * -> *) somerep.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
 SameScope somerep rep) =>
Builder rep (Body rep) -> m (Body rep)
runBodyBuilder forall a b. (a -> b) -> a -> b
$ do
                    -- Flatten both x and y.
                    SubExp
x_num_elems <-
                      forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"x_num_elems"
                        forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
MonadBuilder m =>
BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
foldBinOp (IntType -> Overflow -> BinOp
I.Mul IntType
Int64 Overflow
I.OverflowUndef) (forall v. IsValue v => v -> SubExp
constant (Int64
1 :: Int64)) [SubExp]
x_dims
                    VName
x' <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
"x" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
x
                    VName
y' <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
"x" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
y
                    VName
x_flat <-
                      forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
"x_flat" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ ReshapeKind -> Shape -> VName -> BasicOp
I.Reshape ReshapeKind
I.ReshapeArbitrary (forall d. [d] -> ShapeBase d
I.Shape [SubExp
x_num_elems]) VName
x'
                    VName
y_flat <-
                      forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
"y_flat" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ ReshapeKind -> Shape -> VName -> BasicOp
I.Reshape ReshapeKind
I.ReshapeArbitrary (forall d. [d] -> ShapeBase d
I.Shape [SubExp
x_num_elems]) VName
y'

                    -- Compare the elements.
                    Lambda SOACS
cmp_lam <- forall (m :: * -> *).
(MonadBuilder m, Buildable (Rep m)) =>
CmpOp -> m (Lambda (Rep m))
cmpOpLambda forall a b. (a -> b) -> a -> b
$ PrimType -> CmpOp
I.CmpEq (forall shape u. TypeBase shape u -> PrimType
elemType TypeBase Shape NoUniqueness
x_t)
                    VName
cmps <-
                      forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
"cmps" forall a b. (a -> b) -> a -> b
$
                        forall rep. Op rep -> Exp rep
I.Op forall a b. (a -> b) -> a -> b
$
                          forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
x_num_elems [VName
x_flat, VName
y_flat] (forall rep. Lambda rep -> ScremaForm rep
I.mapSOAC Lambda SOACS
cmp_lam)

                    -- Check that all were equal.
                    Lambda SOACS
and_lam <- forall (m :: * -> *).
(MonadBuilder m, Buildable (Rep m)) =>
BinOp -> PrimType -> m (Lambda (Rep m))
binOpLambda BinOp
I.LogAnd PrimType
I.Bool
                    ScremaForm SOACS
reduce <- forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Reduce rep] -> m (ScremaForm rep)
I.reduceSOAC [forall rep. Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
Reduce Commutativity
Commutative Lambda SOACS
and_lam [forall v. IsValue v => v -> SubExp
constant Bool
True]]
                    SubExp
all_equal <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"all_equal" forall a b. (a -> b) -> a -> b
$ forall rep. Op rep -> Exp rep
I.Op forall a b. (a -> b) -> a -> b
$ forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
x_num_elems [VName
cmps] ScremaForm SOACS
reduce
                    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall rep. Buildable rep => [SubExp] -> Body rep
resultBody [SubExp
all_equal]

              forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"arrays_equal"
                forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf (forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
shapes_match) InternaliseM (Body SOACS)
compare_elems_body (forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [forall v. IsValue v => v -> SubExp
constant Bool
False])
    handle [Char]
name
      | Just BinOp
bop <- forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (([Char]
name ==) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> [Char]
prettyString) [forall a. Bounded a => a
minBound .. forall a. Bounded a => a
maxBound :: E.BinOp] =
          forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[(TypeBase Size ()
x_t, [SubExp
x']), (TypeBase Size ()
y_t, [SubExp
y'])] ->
            case (TypeBase Size ()
x_t, TypeBase Size ()
y_t) of
              (E.Scalar (E.Prim PrimType
t1), E.Scalar (E.Prim PrimType
t2)) ->
                SrcLoc
-> [Char]
-> BinOp
-> SubExp
-> SubExp
-> PrimType
-> PrimType
-> InternaliseM [SubExp]
internaliseBinOp SrcLoc
loc [Char]
desc BinOp
bop SubExp
x' SubExp
y' PrimType
t1 PrimType
t2
              (TypeBase Size (), TypeBase Size ())
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise.internaliseExp: non-primitive type in BinOp."
    handle [Char]
_ = forall a. Maybe a
Nothing

-- | Handle intrinsic functions.  These are only allowed to be called
-- in the prelude, and their internalisation may involve inspecting
-- the AST.
isIntrinsicFunction ::
  E.QualName VName ->
  [E.Exp] ->
  SrcLoc ->
  Maybe (String -> InternaliseM [SubExp])
isIntrinsicFunction :: QualName VName
-> [Exp] -> SrcLoc -> Maybe ([Char] -> InternaliseM [SubExp])
isIntrinsicFunction QualName VName
qname [Exp]
args SrcLoc
loc = do
  forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ VName -> Int
baseTag (forall vn. QualName vn -> vn
qualLeaf QualName VName
qname) forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag
  let handlers :: [[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])]
handlers =
        [ forall {a}.
(Eq a, IsString a) =>
[Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleSign,
          forall {f :: * -> *}.
Applicative f =>
[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM (f SubExp))
handleOps,
          [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
handleSOACs,
          forall {a}.
(Eq a, IsString a) =>
[Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleAccs,
          forall {a}.
(IsString a, Eq a) =>
[Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleAD,
          [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
handleRest
        ]
  forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, MonadPlus m) =>
t (m a) -> m a
msum [[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
h [Exp]
args forall a b. (a -> b) -> a -> b
$ VName -> [Char]
baseString forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> vn
qualLeaf QualName VName
qname | [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
h <- [[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])]
handlers]
  where
    handleSign :: [Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleSign [Exp
x] a
"sign_i8" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
I.Int8 Exp
x
    handleSign [Exp
x] a
"sign_i16" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
I.Int16 Exp
x
    handleSign [Exp
x] a
"sign_i32" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
I.Int32 Exp
x
    handleSign [Exp
x] a
"sign_i64" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
I.Int64 Exp
x
    handleSign [Exp
x] a
"unsign_i8" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
I.Int8 Exp
x
    handleSign [Exp
x] a
"unsign_i16" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
I.Int16 Exp
x
    handleSign [Exp
x] a
"unsign_i32" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
I.Int32 Exp
x
    handleSign [Exp
x] a
"unsign_i64" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
I.Int64 Exp
x
    handleSign [Exp]
_ a
_ = forall a. Maybe a
Nothing

    handleOps :: [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM (f SubExp))
handleOps [Exp
x] [Char]
s
      | Just UnOp
unop <- forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((forall a. Eq a => a -> a -> Bool
== [Char]
s) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> [Char]
prettyString) [UnOp]
allUnOps = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
          SubExp
x' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"x" Exp
x
          forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
unop SubExp
x'
    handleOps [TupLit [Exp
x, Exp
y] SrcLoc
_] [Char]
s
      | Just BinOp
bop <- forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((forall a. Eq a => a -> a -> Bool
== [Char]
s) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> [Char]
prettyString) [BinOp]
allBinOps = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
          SubExp
x' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"x" Exp
x
          SubExp
y' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"y" Exp
y
          forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
bop SubExp
x' SubExp
y'
      | Just CmpOp
cmp <- forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((forall a. Eq a => a -> a -> Bool
== [Char]
s) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> [Char]
prettyString) [CmpOp]
allCmpOps = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
          SubExp
x' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"x" Exp
x
          SubExp
y' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"y" Exp
y
          forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp CmpOp
cmp SubExp
x' SubExp
y'
    handleOps [Exp
x] [Char]
s
      | Just ConvOp
conv <- forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((forall a. Eq a => a -> a -> Bool
== [Char]
s) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> [Char]
prettyString) [ConvOp]
allConvOps = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
          SubExp
x' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"x" Exp
x
          forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp ConvOp
conv SubExp
x'
    handleOps [Exp]
_ [Char]
_ = forall a. Maybe a
Nothing

    handleSOACs :: [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
handleSOACs [TupLit [Exp
lam, Exp
arr] SrcLoc
_] [Char]
"map" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      [VName]
arr' <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"map_arr" Exp
arr
      [TypeBase Shape NoUniqueness]
arr_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arr'
      Lambda SOACS
lam' <- Exp -> [TypeBase Shape NoUniqueness] -> InternaliseM (Lambda SOACS)
internaliseLambdaCoerce Exp
lam forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall u. TypeBase Shape u -> TypeBase Shape u
rowType [TypeBase Shape NoUniqueness]
arr_ts
      let w :: SubExp
w = forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [TypeBase Shape NoUniqueness]
arr_ts
      forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. Op rep -> Exp rep
I.Op forall a b. (a -> b) -> a -> b
$ forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w [VName]
arr' (forall rep. Lambda rep -> ScremaForm rep
I.mapSOAC Lambda SOACS
lam')
    handleSOACs [TupLit [Exp
k, Exp
lam, Exp
arr] SrcLoc
_] [Char]
"partition" = do
      Int
k' <- forall a b. (Integral a, Num b) => a -> b
fromIntegral forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {vn}. ExpBase Info vn -> Maybe Int32
fromInt32 Exp
k
      forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
_desc -> do
        [VName]
arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"partition_input" Exp
arr
        Lambda SOACS
lam' <- InternaliseLambda
-> Int -> Exp -> [SubExp] -> InternaliseM (Lambda SOACS)
internalisePartitionLambda InternaliseLambda
internaliseLambda Int
k' Exp
lam forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var [VName]
arrs
        forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall a. [a] -> [a] -> [a]
(++) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Lambda SOACS -> [VName] -> InternaliseM ([SubExp], [SubExp])
partitionWithSOACS (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k') Lambda SOACS
lam' [VName]
arrs
      where
        fromInt32 :: ExpBase Info vn -> Maybe Int32
fromInt32 (Literal (SignedValue (Int32Value Int32
k')) SrcLoc
_) = forall a. a -> Maybe a
Just Int32
k'
        fromInt32 (IntLit Integer
k' (Info (E.Scalar (E.Prim (E.Signed IntType
Int32)))) SrcLoc
_) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a. Num a => Integer -> a
fromInteger Integer
k'
        fromInt32 ExpBase Info vn
_ = forall a. Maybe a
Nothing
    handleSOACs [TupLit [Exp
lam, Exp
ne, Exp
arr] SrcLoc
_] [Char]
"reduce" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char]
-> [Char]
-> (SubExp
    -> Lambda SOACS
    -> [SubExp]
    -> [VName]
    -> InternaliseM (SOAC SOACS))
-> (Exp, Exp, Exp, SrcLoc)
-> InternaliseM [SubExp]
internaliseScanOrReduce [Char]
desc [Char]
"reduce" forall {f :: * -> *} {rep}.
(Buildable rep, MonadFreshNames f) =>
SubExp -> Lambda rep -> [SubExp] -> [VName] -> f (SOAC rep)
reduce (Exp
lam, Exp
ne, Exp
arr, SrcLoc
loc)
      where
        reduce :: SubExp -> Lambda rep -> [SubExp] -> [VName] -> f (SOAC rep)
reduce SubExp
w Lambda rep
red_lam [SubExp]
nes [VName]
arrs =
          forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w [VName]
arrs
            forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Reduce rep] -> m (ScremaForm rep)
I.reduceSOAC [forall rep. Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
Reduce Commutativity
Noncommutative Lambda rep
red_lam [SubExp]
nes]
    handleSOACs [TupLit [Exp
lam, Exp
ne, Exp
arr] SrcLoc
_] [Char]
"reduce_comm" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char]
-> [Char]
-> (SubExp
    -> Lambda SOACS
    -> [SubExp]
    -> [VName]
    -> InternaliseM (SOAC SOACS))
-> (Exp, Exp, Exp, SrcLoc)
-> InternaliseM [SubExp]
internaliseScanOrReduce [Char]
desc [Char]
"reduce" forall {f :: * -> *} {rep}.
(Buildable rep, MonadFreshNames f) =>
SubExp -> Lambda rep -> [SubExp] -> [VName] -> f (SOAC rep)
reduce (Exp
lam, Exp
ne, Exp
arr, SrcLoc
loc)
      where
        reduce :: SubExp -> Lambda rep -> [SubExp] -> [VName] -> f (SOAC rep)
reduce SubExp
w Lambda rep
red_lam [SubExp]
nes [VName]
arrs =
          forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w [VName]
arrs
            forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Reduce rep] -> m (ScremaForm rep)
I.reduceSOAC [forall rep. Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
Reduce Commutativity
Commutative Lambda rep
red_lam [SubExp]
nes]
    handleSOACs [TupLit [Exp
lam, Exp
ne, Exp
arr] SrcLoc
_] [Char]
"scan" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char]
-> [Char]
-> (SubExp
    -> Lambda SOACS
    -> [SubExp]
    -> [VName]
    -> InternaliseM (SOAC SOACS))
-> (Exp, Exp, Exp, SrcLoc)
-> InternaliseM [SubExp]
internaliseScanOrReduce [Char]
desc [Char]
"scan" forall {f :: * -> *} {rep}.
(Buildable rep, MonadFreshNames f) =>
SubExp -> Lambda rep -> [SubExp] -> [VName] -> f (SOAC rep)
reduce (Exp
lam, Exp
ne, Exp
arr, SrcLoc
loc)
      where
        reduce :: SubExp -> Lambda rep -> [SubExp] -> [VName] -> f (SOAC rep)
reduce SubExp
w Lambda rep
scan_lam [SubExp]
nes [VName]
arrs =
          forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w [VName]
arrs forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Scan rep] -> m (ScremaForm rep)
I.scanSOAC [forall rep. Lambda rep -> [SubExp] -> Scan rep
Scan Lambda rep
scan_lam [SubExp]
nes]
    handleSOACs [TupLit [Exp
rf, Exp
dest, Exp
op, Exp
ne, Exp
buckets, Exp
img] SrcLoc
_] [Char]
"hist_1d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      Int
-> [Char]
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> SrcLoc
-> InternaliseM [SubExp]
internaliseHist Int
1 [Char]
desc Exp
rf Exp
dest Exp
op Exp
ne Exp
buckets Exp
img SrcLoc
loc
    handleSOACs [TupLit [Exp
rf, Exp
dest, Exp
op, Exp
ne, Exp
buckets, Exp
img] SrcLoc
_] [Char]
"hist_2d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      Int
-> [Char]
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> SrcLoc
-> InternaliseM [SubExp]
internaliseHist Int
2 [Char]
desc Exp
rf Exp
dest Exp
op Exp
ne Exp
buckets Exp
img SrcLoc
loc
    handleSOACs [TupLit [Exp
rf, Exp
dest, Exp
op, Exp
ne, Exp
buckets, Exp
img] SrcLoc
_] [Char]
"hist_3d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      Int
-> [Char]
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> SrcLoc
-> InternaliseM [SubExp]
internaliseHist Int
3 [Char]
desc Exp
rf Exp
dest Exp
op Exp
ne Exp
buckets Exp
img SrcLoc
loc
    handleSOACs [Exp]
_ [Char]
_ = forall a. Maybe a
Nothing

    handleAccs :: [Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleAccs [TupLit [Exp
dest, Exp
f, Exp
bs] SrcLoc
_] a
"scatter_stream" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char]
-> Exp -> Maybe (Exp, Exp) -> Exp -> Exp -> InternaliseM [SubExp]
internaliseStreamAcc [Char]
desc Exp
dest forall a. Maybe a
Nothing Exp
f Exp
bs
    handleAccs [TupLit [Exp
dest, Exp
op, Exp
ne, Exp
f, Exp
bs] SrcLoc
_] a
"hist_stream" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char]
-> Exp -> Maybe (Exp, Exp) -> Exp -> Exp -> InternaliseM [SubExp]
internaliseStreamAcc [Char]
desc Exp
dest (forall a. a -> Maybe a
Just (Exp
op, Exp
ne)) Exp
f Exp
bs
    handleAccs [TupLit [Exp
acc, Exp
i, Exp
v] SrcLoc
_] a
"acc_write" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      VName
acc' <- forall a. [a] -> a
head forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"acc" Exp
acc
      SubExp
i' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"acc_i" Exp
i
      [SubExp]
vs <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"acc_v" Exp
v
      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ VName -> [SubExp] -> [SubExp] -> BasicOp
UpdateAcc VName
acc' [SubExp
i'] [SubExp]
vs
    handleAccs [Exp]
_ a
_ = forall a. Maybe a
Nothing

    handleAD :: [Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleAD [TupLit [Exp
f, Exp
x, Exp
v] SrcLoc
_] a
fname
      | a
fname forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [a
"jvp2", a
"vjp2"] = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
          [SubExp]
x' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"ad_x" Exp
x
          [SubExp]
v' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"ad_v" Exp
v
          Lambda SOACS
lam <- Exp -> [TypeBase Shape NoUniqueness] -> InternaliseM (Lambda SOACS)
internaliseLambdaCoerce Exp
f forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
x'
          forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [VName]
letTupExp [Char]
desc forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. Op rep -> Exp rep
Op forall a b. (a -> b) -> a -> b
$
            case a
fname of
              a
"jvp2" -> forall rep. Lambda rep -> [SubExp] -> [SubExp] -> SOAC rep
JVP Lambda SOACS
lam [SubExp]
x' [SubExp]
v'
              a
_ -> forall rep. Lambda rep -> [SubExp] -> [SubExp] -> SOAC rep
VJP Lambda SOACS
lam [SubExp]
x' [SubExp]
v'
    handleAD [Exp]
_ a
_ = forall a. Maybe a
Nothing

    handleRest :: [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
handleRest [E.TupLit [Exp
a, Exp
si, Exp
v] SrcLoc
_] [Char]
"scatter" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Int -> Exp -> Exp -> Exp -> [Char] -> InternaliseM [SubExp]
scatterF Int
1 Exp
a Exp
si Exp
v
    handleRest [E.TupLit [Exp
a, Exp
si, Exp
v] SrcLoc
_] [Char]
"scatter_2d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Int -> Exp -> Exp -> Exp -> [Char] -> InternaliseM [SubExp]
scatterF Int
2 Exp
a Exp
si Exp
v
    handleRest [E.TupLit [Exp
a, Exp
si, Exp
v] SrcLoc
_] [Char]
"scatter_3d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Int -> Exp -> Exp -> Exp -> [Char] -> InternaliseM [SubExp]
scatterF Int
3 Exp
a Exp
si Exp
v
    handleRest [E.TupLit [Exp
n, Exp
m, Exp
arr] SrcLoc
_] [Char]
"unflatten" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      [VName]
arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"unflatten_arr" Exp
arr
      SubExp
n' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"n" Exp
n
      SubExp
m' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"m" Exp
m
      -- The unflattened dimension needs to have the same number of elements
      -- as the original dimension.
      SubExp
old_dim <- forall u. Int -> [TypeBase Shape u] -> SubExp
I.arraysSize Int
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
      SubExp
dim_ok <-
        forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"dim_ok"
          forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
MonadBuilder m =>
CmpOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eCmpOp
            (PrimType -> CmpOp
I.CmpEq PrimType
I.int64)
            (forall (m :: * -> *).
MonadBuilder m =>
BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eBinOp (IntType -> Overflow -> BinOp
I.Mul IntType
Int64 Overflow
I.OverflowUndef) (forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
n') (forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
m'))
            (forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
old_dim)
      Certs
dim_ok_cert <-
        [Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert
          [Char]
"dim_ok_cert"
          SubExp
dim_ok
          ( forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg
              [ ErrorMsgPart SubExp
"Cannot unflatten array of shape [",
                forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
old_dim,
                ErrorMsgPart SubExp
"] to array of shape [",
                forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
n',
                ErrorMsgPart SubExp
"][",
                forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
m',
                ErrorMsgPart SubExp
"]"
              ]
          )
          SrcLoc
loc
      forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
dim_ok_cert forall a b. (a -> b) -> a -> b
$
        forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
arrs forall a b. (a -> b) -> a -> b
$ \VName
arr' -> do
          TypeBase Shape NoUniqueness
arr_t <- forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
arr'
          forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
            ReshapeKind -> Shape -> VName -> BasicOp
I.Reshape
              ReshapeKind
I.ReshapeArbitrary
              (Shape -> Int -> Shape -> Shape
reshapeOuter (forall d. [d] -> ShapeBase d
I.Shape [SubExp
n', SubExp
m']) Int
1 forall a b. (a -> b) -> a -> b
$ forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape TypeBase Shape NoUniqueness
arr_t)
              VName
arr'
    handleRest [Exp
arr] [Char]
"flatten" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      [VName]
arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"flatten_arr" Exp
arr
      forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
arrs forall a b. (a -> b) -> a -> b
$ \VName
arr' -> do
        TypeBase Shape NoUniqueness
arr_t <- forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
arr'
        let n :: SubExp
n = forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
0 TypeBase Shape NoUniqueness
arr_t
            m :: SubExp
m = forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
1 TypeBase Shape NoUniqueness
arr_t
        SubExp
k <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"flat_dim" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Mul IntType
Int64 Overflow
I.OverflowUndef) SubExp
n SubExp
m
        forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
          ReshapeKind -> Shape -> VName -> BasicOp
I.Reshape
            ReshapeKind
I.ReshapeArbitrary
            (Shape -> Int -> Shape -> Shape
reshapeOuter (forall d. [d] -> ShapeBase d
I.Shape [SubExp
k]) Int
2 forall a b. (a -> b) -> a -> b
$ forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape TypeBase Shape NoUniqueness
arr_t)
            VName
arr'
    handleRest [TupLit [Exp
x, Exp
y] SrcLoc
_] [Char]
"concat" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      [VName]
xs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"concat_x" Exp
x
      [VName]
ys <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"concat_y" Exp
y
      SubExp
outer_size <- forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
xs
      let sumdims :: SubExp -> SubExp -> m SubExp
sumdims SubExp
xsize SubExp
ysize =
            forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"conc_tmp" forall a b. (a -> b) -> a -> b
$
              forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
                BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Add IntType
I.Int64 Overflow
I.OverflowUndef) SubExp
xsize SubExp
ysize
      SubExp
ressize <-
        forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM forall {m :: * -> *}.
MonadBuilder m =>
SubExp -> SubExp -> m SubExp
sumdims SubExp
outer_size
          forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType) [[VName]
ys]

      let conc :: VName -> VName -> Exp SOACS
conc VName
xarr VName
yarr =
            forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ Int -> NonEmpty VName -> SubExp -> BasicOp
I.Concat Int
0 (VName
xarr forall a. a -> [a] -> NonEmpty a
:| [VName
yarr]) SubExp
ressize
      forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc) forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith VName -> VName -> Exp SOACS
conc [VName]
xs [VName]
ys
    handleRest [TupLit [Exp
offset, Exp
e] SrcLoc
_] [Char]
"rotate" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      SubExp
offset' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"rotation_offset" Exp
offset
      [Char]
-> Exp -> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
internaliseOperation [Char]
desc Exp
e forall a b. (a -> b) -> a -> b
$ \VName
v -> do
        Int
r <- forall shape u. ArrayShape shape => TypeBase shape u -> Int
I.arrayRank forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
v
        let zero :: SubExp
zero = IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0
            offsets :: [SubExp]
offsets = SubExp
offset' forall a. a -> [a] -> [a]
: forall a. Int -> a -> [a]
replicate (Int
r forall a. Num a => a -> a -> a
- Int
1) SubExp
zero
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [SubExp] -> VName -> BasicOp
I.Rotate [SubExp]
offsets VName
v
    handleRest [Exp
e] [Char]
"transpose" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char]
-> Exp -> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
internaliseOperation [Char]
desc Exp
e forall a b. (a -> b) -> a -> b
$ \VName
v -> do
        Int
r <- forall shape u. ArrayShape shape => TypeBase shape u -> Int
I.arrayRank forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
v
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [Int] -> VName -> BasicOp
I.Rearrange ([Int
1, Int
0] forall a. [a] -> [a] -> [a]
++ [Int
2 .. Int
r forall a. Num a => a -> a -> a
- Int
1]) VName
v
    handleRest [TupLit [Exp
x, Exp
y] SrcLoc
_] [Char]
"zip" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"zip_copy" forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. BasicOp -> Exp rep
BasicOp forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> BasicOp
Copy)
        forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ( forall a. [a] -> [a] -> [a]
(++)
                forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars ([Char]
desc forall a. [a] -> [a] -> [a]
++ [Char]
"_zip_x") Exp
x
                forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars ([Char]
desc forall a. [a] -> [a] -> [a]
++ [Char]
"_zip_y") Exp
y
            )
    handleRest [Exp
x] [Char]
"unzip" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> b -> a -> c
flip [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp Exp
x
    handleRest [TupLit [Exp
arr, Exp
offset, Exp
n1, Exp
s1, Exp
n2, Exp
s2] SrcLoc
_] [Char]
"flat_index_2d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      [Char]
-> SrcLoc -> Exp -> Exp -> [(Exp, Exp)] -> InternaliseM [SubExp]
flatIndexHelper [Char]
desc SrcLoc
loc Exp
arr Exp
offset [(Exp
n1, Exp
s1), (Exp
n2, Exp
s2)]
    handleRest [TupLit [Exp
arr1, Exp
offset, Exp
s1, Exp
s2, Exp
arr2] SrcLoc
_] [Char]
"flat_update_2d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      [Char]
-> SrcLoc -> Exp -> Exp -> [Exp] -> Exp -> InternaliseM [SubExp]
flatUpdateHelper [Char]
desc SrcLoc
loc Exp
arr1 Exp
offset [Exp
s1, Exp
s2] Exp
arr2
    handleRest [TupLit [Exp
arr, Exp
offset, Exp
n1, Exp
s1, Exp
n2, Exp
s2, Exp
n3, Exp
s3] SrcLoc
_] [Char]
"flat_index_3d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      [Char]
-> SrcLoc -> Exp -> Exp -> [(Exp, Exp)] -> InternaliseM [SubExp]
flatIndexHelper [Char]
desc SrcLoc
loc Exp
arr Exp
offset [(Exp
n1, Exp
s1), (Exp
n2, Exp
s2), (Exp
n3, Exp
s3)]
    handleRest [TupLit [Exp
arr1, Exp
offset, Exp
s1, Exp
s2, Exp
s3, Exp
arr2] SrcLoc
_] [Char]
"flat_update_3d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      [Char]
-> SrcLoc -> Exp -> Exp -> [Exp] -> Exp -> InternaliseM [SubExp]
flatUpdateHelper [Char]
desc SrcLoc
loc Exp
arr1 Exp
offset [Exp
s1, Exp
s2, Exp
s3] Exp
arr2
    handleRest [TupLit [Exp
arr, Exp
offset, Exp
n1, Exp
s1, Exp
n2, Exp
s2, Exp
n3, Exp
s3, Exp
n4, Exp
s4] SrcLoc
_] [Char]
"flat_index_4d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      [Char]
-> SrcLoc -> Exp -> Exp -> [(Exp, Exp)] -> InternaliseM [SubExp]
flatIndexHelper [Char]
desc SrcLoc
loc Exp
arr Exp
offset [(Exp
n1, Exp
s1), (Exp
n2, Exp
s2), (Exp
n3, Exp
s3), (Exp
n4, Exp
s4)]
    handleRest [TupLit [Exp
arr1, Exp
offset, Exp
s1, Exp
s2, Exp
s3, Exp
s4, Exp
arr2] SrcLoc
_] [Char]
"flat_update_4d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      [Char]
-> SrcLoc -> Exp -> Exp -> [Exp] -> Exp -> InternaliseM [SubExp]
flatUpdateHelper [Char]
desc SrcLoc
loc Exp
arr1 Exp
offset [Exp
s1, Exp
s2, Exp
s3, Exp
s4] Exp
arr2
    handleRest [Exp]
_ [Char]
_ = forall a. Maybe a
Nothing

    toSigned :: IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
int_to Exp
e [Char]
desc = do
      SubExp
e' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"trunc_arg" Exp
e
      case Exp -> PatType
E.typeOf Exp
e of
        E.Scalar (E.Prim PrimType
E.Bool) ->
          forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc
            forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
              (forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
e')
              (forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
1])
              (forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
0])
        E.Scalar (E.Prim (E.Signed IntType
int_from)) ->
          forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp 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)) ->
          forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp 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)) ->
          forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp 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'
        PatType
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise: non-numeric type in ToSigned"

    toUnsigned :: IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
int_to Exp
e [Char]
desc = do
      SubExp
e' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"trunc_arg" Exp
e
      case Exp -> PatType
E.typeOf Exp
e of
        E.Scalar (E.Prim PrimType
E.Bool) ->
          forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc
            forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
              (forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
e')
              (forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
1])
              (forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
0])
        E.Scalar (E.Prim (E.Signed IntType
int_from)) ->
          forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp 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)) ->
          forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp 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)) ->
          forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp 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'
        PatType
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise.internaliseExp: non-numeric type in ToUnsigned"

    scatterF :: Int -> Exp -> Exp -> Exp -> [Char] -> InternaliseM [SubExp]
scatterF Int
dim Exp
a Exp
si Exp
v [Char]
desc = do
      [VName]
si' <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"write_arg_i" Exp
si
      [VName]
svs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"write_arg_v" Exp
v
      [VName]
sas <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"write_arg_a" Exp
a

      SubExp
si_w <- forall u. Int -> [TypeBase Shape u] -> SubExp
I.arraysSize Int
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
si'
      [TypeBase Shape NoUniqueness]
sv_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
svs

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

        -- Generate an assertion and reshapes to ensure that sv and si' are the same
        -- size.
        SubExp
cmp <-
          forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"write_cmp" forall a b. (a -> b) -> a -> b
$
            forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
              CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
I.int64) SubExp
si_w SubExp
sv_w
        Certs
c <-
          [Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert
            [Char]
"write_cert"
            SubExp
cmp
            ErrorMsg SubExp
"length of index and value array does not match"
            SrcLoc
loc
        forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c forall a b. (a -> b) -> a -> b
$
          forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp (VName -> [Char]
baseString VName
sv forall a. [a] -> [a] -> [a]
++ [Char]
"_write_sv") forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
            ReshapeKind -> Shape -> VName -> BasicOp
I.Reshape ReshapeKind
I.ReshapeCoerce (Shape -> Int -> Shape -> Shape
reshapeOuter (forall d. [d] -> ShapeBase d
I.Shape [SubExp
si_w]) Int
1 Shape
sv_shape) VName
sv

      [TypeBase Shape NoUniqueness]
indexType <- forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall u. TypeBase Shape u -> TypeBase Shape u
rowType forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
si'
      [VName]
indexName <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\TypeBase Shape NoUniqueness
_ -> forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"write_index") [TypeBase Shape NoUniqueness]
indexType
      [VName]
valueNames <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase Shape NoUniqueness]
sv_ts) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"write_value"

      [TypeBase Shape NoUniqueness]
sa_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
sas
      let bodyTypes :: [TypeBase Shape NoUniqueness]
bodyTypes = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase Shape NoUniqueness]
sv_ts) [TypeBase Shape NoUniqueness]
indexType) forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map (forall u. Int -> TypeBase Shape u -> TypeBase Shape u
I.stripArray Int
dim) [TypeBase Shape NoUniqueness]
sa_ts
          paramTypes :: [TypeBase Shape NoUniqueness]
paramTypes = [TypeBase Shape NoUniqueness]
indexType forall a. Semigroup a => a -> a -> a
<> forall a b. (a -> b) -> [a] -> [b]
map forall u. TypeBase Shape u -> TypeBase Shape u
rowType [TypeBase Shape NoUniqueness]
sv_ts
          bodyNames :: [VName]
bodyNames = [VName]
indexName forall a. Semigroup a => a -> a -> a
<> [VName]
valueNames
          bodyParams :: [Param (TypeBase Shape NoUniqueness)]
bodyParams = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (forall dec. Attrs -> VName -> dec -> Param dec
I.Param forall a. Monoid a => a
mempty) [VName]
bodyNames [TypeBase Shape NoUniqueness]
paramTypes

      -- This body is prettyString boring right now, as every input is exactly the output.
      -- But it can get funky later on if fused with something else.
      Body SOACS
body <- forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams [Param (TypeBase Shape NoUniqueness)]
bodyParams) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *).
MonadBuilder m =>
m Result -> m (Body (Rep m))
buildBody_ forall a b. (a -> b) -> a -> b
$ do
        let outs :: [VName]
outs = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
valueNames) [VName]
indexName) forall a. [a] -> [a] -> [a]
++ [VName]
valueNames
        [SubExp]
results <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
outs forall a b. (a -> b) -> a -> b
$ \VName
name ->
          forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"write_res" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
name
        ErrorMsg SubExp
-> SrcLoc
-> [TypeBase Shape NoUniqueness]
-> Result
-> InternaliseM Result
ensureResultShape
          ErrorMsg SubExp
"scatter value has wrong size"
          SrcLoc
loc
          [TypeBase Shape NoUniqueness]
bodyTypes
          ([SubExp] -> Result
subExpsRes [SubExp]
results)

      let lam :: Lambda SOACS
lam =
            I.Lambda
              { lambdaParams :: [LParam SOACS]
I.lambdaParams = [Param (TypeBase Shape NoUniqueness)]
bodyParams,
                lambdaReturnType :: [TypeBase Shape NoUniqueness]
I.lambdaReturnType = [TypeBase Shape NoUniqueness]
bodyTypes,
                lambdaBody :: Body SOACS
I.lambdaBody = Body SOACS
body
              }
          sivs :: [VName]
sivs = [VName]
si' forall a. Semigroup a => a -> a -> a
<> [VName]
svs'

      let sa_ws :: [Shape]
sa_ws = forall a b. (a -> b) -> [a] -> [b]
map (forall d. [d] -> ShapeBase d
I.Shape forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Int -> [a] -> [a]
take Int
dim forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall u. TypeBase Shape u -> [SubExp]
arrayDims) [TypeBase Shape NoUniqueness]
sa_ts
      forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. Op rep -> Exp rep
I.Op forall a b. (a -> b) -> a -> b
$ forall rep.
SubExp
-> [VName] -> Lambda rep -> [(Shape, Int, VName)] -> SOAC rep
I.Scatter SubExp
si_w [VName]
sivs Lambda SOACS
lam forall a b. (a -> b) -> a -> b
$ forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Shape]
sa_ws (forall a. a -> [a]
repeat Int
1) [VName]
sas

flatIndexHelper :: String -> SrcLoc -> E.Exp -> E.Exp -> [(E.Exp, E.Exp)] -> InternaliseM [SubExp]
flatIndexHelper :: [Char]
-> SrcLoc -> Exp -> Exp -> [(Exp, Exp)] -> InternaliseM [SubExp]
flatIndexHelper [Char]
desc SrcLoc
loc Exp
arr Exp
offset [(Exp, Exp)]
slices = do
  [VName]
arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"arr" Exp
arr
  SubExp
offset' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"offset" Exp
offset
  SubExp
old_dim <- forall u. Int -> [TypeBase Shape u] -> SubExp
I.arraysSize Int
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
  SubExp
offset_inbounds_down <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"offset_inbounds_down" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpUle IntType
Int64) (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0) SubExp
offset'
  SubExp
offset_inbounds_up <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"offset_inbounds_up" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpUlt IntType
Int64) SubExp
offset' SubExp
old_dim
  [(SubExp, SubExp)]
slices' <-
    forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM
      ( \(Exp
n, Exp
s) -> do
          SubExp
n' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"n" Exp
n
          SubExp
s' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"s" Exp
s
          forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
n', SubExp
s')
      )
      [(Exp, Exp)]
slices
  (SubExp
min_bound, SubExp
max_bound) <-
    forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
      ( \(SubExp
lower, SubExp
upper) (SubExp
n, SubExp
s) -> do
          SubExp
n_m1 <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
Int64 Overflow
I.OverflowUndef) SubExp
n (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1)
          SubExp
spn <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Mul IntType
Int64 Overflow
I.OverflowUndef) SubExp
n_m1 SubExp
s

          SubExp
span_and_lower <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span_and_lower" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Add IntType
Int64 Overflow
I.OverflowUndef) SubExp
spn SubExp
lower
          SubExp
span_and_upper <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span_and_upper" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Add IntType
Int64 Overflow
I.OverflowUndef) SubExp
spn SubExp
upper

          SubExp
lower' <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"minimum" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> BinOp
I.UMin IntType
Int64) SubExp
span_and_lower SubExp
lower
          SubExp
upper' <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"maximum" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> BinOp
I.UMax IntType
Int64) SubExp
span_and_upper SubExp
upper

          forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
lower', SubExp
upper')
      )
      (SubExp
offset', SubExp
offset')
      [(SubExp, SubExp)]
slices'
  SubExp
min_in_bounds <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"min_in_bounds" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpUle IntType
Int64) (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0) SubExp
min_bound
  SubExp
max_in_bounds <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"max_in_bounds" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpUlt IntType
Int64) SubExp
max_bound SubExp
old_dim

  SubExp
all_bounds <-
    forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
      (\SubExp
x SubExp
y -> forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"inBounds" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogAnd SubExp
x SubExp
y)
      SubExp
offset_inbounds_down
      [SubExp
offset_inbounds_up, SubExp
min_in_bounds, SubExp
max_in_bounds]

  Certs
c <- [Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert [Char]
"bounds_cert" SubExp
all_bounds (forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg [forall a. Text -> ErrorMsgPart a
ErrorString forall a b. (a -> b) -> a -> b
$ Text
"Flat slice out of bounds: " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText SubExp
old_dim forall a. Semigroup a => a -> a -> a
<> Text
" and " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText [(SubExp, SubExp)]
slices']) SrcLoc
loc
  let slice :: FlatSlice SubExp
slice = forall d. d -> [FlatDimIndex d] -> FlatSlice d
I.FlatSlice SubExp
offset' forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall d. d -> d -> FlatDimIndex d
FlatDimIndex) [(SubExp, SubExp)]
slices'
  forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c forall a b. (a -> b) -> a -> b
$
    forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
arrs forall a b. (a -> b) -> a -> b
$ \VName
arr' ->
      forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ VName -> FlatSlice SubExp -> BasicOp
I.FlatIndex VName
arr' FlatSlice SubExp
slice

flatUpdateHelper :: String -> SrcLoc -> E.Exp -> E.Exp -> [E.Exp] -> E.Exp -> InternaliseM [SubExp]
flatUpdateHelper :: [Char]
-> SrcLoc -> Exp -> Exp -> [Exp] -> Exp -> InternaliseM [SubExp]
flatUpdateHelper [Char]
desc SrcLoc
loc Exp
arr1 Exp
offset [Exp]
slices Exp
arr2 = do
  [VName]
arrs1 <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"arr" Exp
arr1
  SubExp
offset' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"offset" Exp
offset
  SubExp
old_dim <- forall u. Int -> [TypeBase Shape u] -> SubExp
I.arraysSize Int
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs1
  SubExp
offset_inbounds_down <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"offset_inbounds_down" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpUle IntType
Int64) (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0) SubExp
offset'
  SubExp
offset_inbounds_up <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"offset_inbounds_up" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpUlt IntType
Int64) SubExp
offset' SubExp
old_dim
  [VName]
arrs2 <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"arr" Exp
arr2
  [TypeBase Shape NoUniqueness]
ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs2
  [(SubExp, SubExp)]
slices' <-
    forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM
      ( \(Exp
s, Int
i) -> do
          SubExp
s' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"s" Exp
s
          let n :: SubExp
n = forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
i [TypeBase Shape NoUniqueness]
ts
          forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
n, SubExp
s')
      )
      forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [Exp]
slices [Int
0 ..]
  (SubExp
min_bound, SubExp
max_bound) <-
    forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
      ( \(SubExp
lower, SubExp
upper) (SubExp
n, SubExp
s) -> do
          SubExp
n_m1 <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
Int64 Overflow
I.OverflowUndef) SubExp
n (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1)
          SubExp
spn <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Mul IntType
Int64 Overflow
I.OverflowUndef) SubExp
n_m1 SubExp
s

          SubExp
span_and_lower <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span_and_lower" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Add IntType
Int64 Overflow
I.OverflowUndef) SubExp
spn SubExp
lower
          SubExp
span_and_upper <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span_and_upper" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Add IntType
Int64 Overflow
I.OverflowUndef) SubExp
spn SubExp
upper

          SubExp
lower' <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"minimum" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> BinOp
I.UMin IntType
Int64) SubExp
span_and_lower SubExp
lower
          SubExp
upper' <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"maximum" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> BinOp
I.UMax IntType
Int64) SubExp
span_and_upper SubExp
upper

          forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
lower', SubExp
upper')
      )
      (SubExp
offset', SubExp
offset')
      [(SubExp, SubExp)]
slices'
  SubExp
min_in_bounds <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"min_in_bounds" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpUle IntType
Int64) (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0) SubExp
min_bound
  SubExp
max_in_bounds <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"max_in_bounds" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpUlt IntType
Int64) SubExp
max_bound SubExp
old_dim

  SubExp
all_bounds <-
    forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
      (\SubExp
x SubExp
y -> forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"inBounds" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogAnd SubExp
x SubExp
y)
      SubExp
offset_inbounds_down
      [SubExp
offset_inbounds_up, SubExp
min_in_bounds, SubExp
max_in_bounds]

  Certs
c <- [Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert [Char]
"bounds_cert" SubExp
all_bounds (forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg [forall a. Text -> ErrorMsgPart a
ErrorString forall a b. (a -> b) -> a -> b
$ Text
"Flat slice out of bounds: " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText SubExp
old_dim forall a. Semigroup a => a -> a -> a
<> Text
" and " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText [(SubExp, SubExp)]
slices']) SrcLoc
loc
  let slice :: FlatSlice SubExp
slice = forall d. d -> [FlatDimIndex d] -> FlatSlice d
I.FlatSlice SubExp
offset' forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall d. d -> d -> FlatDimIndex d
FlatDimIndex) [(SubExp, SubExp)]
slices'
  forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c forall a b. (a -> b) -> a -> b
$
    forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
arrs1 [VName]
arrs2) forall a b. (a -> b) -> a -> b
$ \(VName
arr1', VName
arr2') ->
      forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ VName -> FlatSlice SubExp -> VName -> BasicOp
I.FlatUpdate VName
arr1' FlatSlice SubExp
slice VName
arr2'

funcall ::
  String ->
  QualName VName ->
  [SubExp] ->
  SrcLoc ->
  InternaliseM ([SubExp], [I.ExtType])
funcall :: [Char]
-> QualName VName
-> [SubExp]
-> SrcLoc
-> InternaliseM ([SubExp], [ExtType])
funcall [Char]
desc (QualName [VName]
_ VName
fname) [SubExp]
args SrcLoc
loc = do
  ([VName]
shapes, [DeclType]
value_paramts, [Param DeclType]
fun_params, [(SubExp, TypeBase Shape NoUniqueness)]
-> Maybe [TypeBase ExtShape Uniqueness]
rettype_fun) <-
    VName -> InternaliseM FunInfo
lookupFunction VName
fname
  [TypeBase Shape NoUniqueness]
argts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
args

  [SubExp]
shapeargs <- [VName]
-> [FParam SOACS]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM [SubExp]
argShapes [VName]
shapes [Param DeclType]
fun_params [TypeBase Shape NoUniqueness]
argts
  let diets :: [Diet]
diets =
        forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
shapeargs) Diet
I.ObservePrim
          forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map forall shape. TypeBase shape Uniqueness -> Diet
I.diet [DeclType]
value_paramts
  [SubExp]
args' <-
    forall u.
Typed (TypeBase Shape u) =>
ErrorMsg SubExp
-> SrcLoc
-> [VName]
-> [TypeBase Shape u]
-> [SubExp]
-> InternaliseM [SubExp]
ensureArgShapes
      ErrorMsg SubExp
"function arguments of wrong shape"
      SrcLoc
loc
      (forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
I.paramName [Param DeclType]
fun_params)
      (forall a b. (a -> b) -> [a] -> [b]
map forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
I.paramType [Param DeclType]
fun_params)
      ([SubExp]
shapeargs forall a. [a] -> [a] -> [a]
++ [SubExp]
args)
  [TypeBase Shape NoUniqueness]
argts' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
args'
  case [(SubExp, TypeBase Shape NoUniqueness)]
-> Maybe [TypeBase ExtShape Uniqueness]
rettype_fun forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
args' [TypeBase Shape NoUniqueness]
argts' of
    Maybe [TypeBase ExtShape Uniqueness]
Nothing ->
      forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
        forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
          [ [Char]
"Cannot apply ",
            forall a. Pretty a => a -> [Char]
prettyString VName
fname,
            [Char]
" to ",
            forall a. Show a => a -> [Char]
show (forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
args'),
            [Char]
" arguments\n ",
            forall a. Pretty a => a -> [Char]
prettyString [SubExp]
args',
            [Char]
"\nof types\n ",
            forall a. Pretty a => a -> [Char]
prettyString [TypeBase Shape NoUniqueness]
argts',
            [Char]
"\nFunction has ",
            forall a. Show a => a -> [Char]
show (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Param DeclType]
fun_params),
            [Char]
" parameters\n ",
            forall a. Pretty a => a -> [Char]
prettyString [Param DeclType]
fun_params
          ]
    Just [TypeBase ExtShape Uniqueness]
ts -> do
      Safety
safety <- InternaliseM Safety
askSafety
      Attrs
attrs <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks InternaliseEnv -> Attrs
envAttrs
      [SubExp]
ses <-
        forall (m :: * -> *) a. MonadBuilder m => Attrs -> m a -> m a
attributing Attrs
attrs forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp SOACS -> InternaliseM [SubExp]
letValExp' [Char]
desc forall a b. (a -> b) -> a -> b
$
          forall rep.
Name
-> [(SubExp, Diet)]
-> [RetType rep]
-> (Safety, SrcLoc, [SrcLoc])
-> Exp rep
I.Apply (VName -> Name
internaliseFunName VName
fname) (forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
args' [Diet]
diets) [TypeBase ExtShape Uniqueness]
ts (Safety
safety, SrcLoc
loc, forall a. Monoid a => a
mempty)
      forall (f :: * -> *) a. Applicative f => a -> f a
pure ([SubExp]
ses, forall a b. (a -> b) -> [a] -> [b]
map forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl [TypeBase ExtShape Uniqueness]
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 :: AppRes -> [SubExp] -> InternaliseM ()
bindExtSizes :: AppRes -> [SubExp] -> InternaliseM ()
bindExtSizes (AppRes PatType
ret [VName]
retext) [SubExp]
ses = do
  let ts :: [TypeBase ExtShape Uniqueness]
ts = TypeBase Size () -> [TypeBase ExtShape Uniqueness]
internaliseType forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatType
ret
  [TypeBase Shape NoUniqueness]
ses_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
ses

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

  forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (forall k a. Map k a -> [(k, a)]
M.toList forall a b. (a -> b) -> a -> b
$ forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeBase ExtShape Uniqueness
-> TypeBase Shape NoUniqueness -> Map VName SubExp
combine [TypeBase ExtShape Uniqueness]
ts [TypeBase Shape NoUniqueness]
ses_ts) forall a b. (a -> b) -> a -> b
$ \(VName
v, SubExp
se) ->
    forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
v] forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se

askSafety :: InternaliseM Safety
askSafety :: InternaliseM Safety
askSafety = do
  Bool
check <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks InternaliseEnv -> Bool
envDoBoundsChecks
  forall (f :: * -> *) a. Applicative f => a -> f a
pure 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 SOACS -> [I.VName] -> InternaliseM ([I.SubExp], [I.SubExp])
partitionWithSOACS :: Int -> Lambda SOACS -> [VName] -> InternaliseM ([SubExp], [SubExp])
partitionWithSOACS Int
k Lambda SOACS
lam [VName]
arrs = do
  [TypeBase Shape NoUniqueness]
arr_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
  let w :: SubExp
w = forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [TypeBase Shape NoUniqueness]
arr_ts
  [VName]
classes_and_increments <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [VName]
letTupExp [Char]
"increments" forall a b. (a -> b) -> a -> b
$ forall rep. Op rep -> Exp rep
I.Op forall a b. (a -> b) -> a -> b
$ forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w [VName]
arrs (forall rep. Lambda rep -> ScremaForm rep
mapSOAC Lambda SOACS
lam)
  (VName
classes, [VName]
increments) <- case [VName]
classes_and_increments of
    VName
classes : [VName]
increments -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
classes, forall a. Int -> [a] -> [a]
take Int
k [VName]
increments)
    [VName]
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"partitionWithSOACS"

  [Param (TypeBase Shape NoUniqueness)]
add_lam_x_params <-
    forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
k forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"x" (forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
  [Param (TypeBase Shape NoUniqueness)]
add_lam_y_params <-
    forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
k forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"y" (forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
  Body SOACS
add_lam_body <- forall rep (m :: * -> *) somerep.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
 SameScope somerep rep) =>
Builder rep (Body rep) -> m (Body rep)
runBodyBuilder forall a b. (a -> b) -> a -> b
$
    forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams forall a b. (a -> b) -> a -> b
$ [Param (TypeBase Shape NoUniqueness)]
add_lam_x_params forall a. [a] -> [a] -> [a]
++ [Param (TypeBase Shape NoUniqueness)]
add_lam_y_params) forall a b. (a -> b) -> a -> b
$
      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall rep. Buildable rep => [SubExp] -> Body rep
resultBody forall a b. (a -> b) -> a -> b
$
        forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (forall a b. [a] -> [b] -> [(a, b)]
zip [Param (TypeBase Shape NoUniqueness)]
add_lam_x_params [Param (TypeBase Shape NoUniqueness)]
add_lam_y_params) forall a b. (a -> b) -> a -> b
$ \(Param (TypeBase Shape NoUniqueness)
x, Param (TypeBase Shape NoUniqueness)
y) ->
          forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"z" forall a b. (a -> b) -> a -> b
$
            forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
              BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp
                (IntType -> Overflow -> BinOp
I.Add IntType
Int64 Overflow
I.OverflowUndef)
                (VName -> SubExp
I.Var forall a b. (a -> b) -> a -> b
$ forall dec. Param dec -> VName
I.paramName Param (TypeBase Shape NoUniqueness)
x)
                (VName -> SubExp
I.Var forall a b. (a -> b) -> a -> b
$ forall dec. Param dec -> VName
I.paramName Param (TypeBase Shape NoUniqueness)
y)
  let add_lam :: Lambda SOACS
add_lam =
        I.Lambda
          { lambdaBody :: Body SOACS
I.lambdaBody = Body SOACS
add_lam_body,
            lambdaParams :: [LParam SOACS]
I.lambdaParams = [Param (TypeBase Shape NoUniqueness)]
add_lam_x_params forall a. [a] -> [a] -> [a]
++ [Param (TypeBase Shape NoUniqueness)]
add_lam_y_params,
            lambdaReturnType :: [TypeBase Shape NoUniqueness]
I.lambdaReturnType = forall a. Int -> a -> [a]
replicate Int
k forall a b. (a -> b) -> a -> b
$ forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64
          }
      nes :: [SubExp]
nes = forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
increments) forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0

  ScremaForm SOACS
scan <- forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Scan rep] -> m (ScremaForm rep)
I.scanSOAC [forall rep. Lambda rep -> [SubExp] -> Scan rep
I.Scan Lambda SOACS
add_lam [SubExp]
nes]
  [VName]
all_offsets <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [VName]
letTupExp [Char]
"offsets" forall a b. (a -> b) -> a -> b
$ forall rep. Op rep -> Exp rep
I.Op forall a b. (a -> b) -> a -> b
$ forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w [VName]
increments ScremaForm SOACS
scan

  -- 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 <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"last_index" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
Int64 Overflow
OverflowUndef) SubExp
w forall a b. (a -> b) -> a -> b
$ forall v. IsValue v => v -> SubExp
constant (Int64
1 :: Int64)
  let nonempty_body :: InternaliseM (Body SOACS)
nonempty_body = forall rep (m :: * -> *) somerep.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
 SameScope somerep rep) =>
Builder rep (Body rep) -> m (Body rep)
runBodyBuilder forall a b. (a -> b) -> a -> b
$
        forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall rep. Buildable rep => [SubExp] -> Body rep
resultBody forall a b. (a -> b) -> a -> b
$
          forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
all_offsets forall a b. (a -> b) -> a -> b
$ \VName
offset_array ->
            forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"last_offset" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp
I.Index VName
offset_array forall a b. (a -> b) -> a -> b
$ forall d. [DimIndex d] -> Slice d
Slice [forall d. d -> DimIndex d
I.DimFix SubExp
last_index]
      empty_body :: InternaliseM (Body (Rep InternaliseM))
empty_body = forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM forall a b. (a -> b) -> a -> b
$ forall a. Int -> a -> [a]
replicate Int
k forall a b. (a -> b) -> a -> b
$ forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)
  SubExp
is_empty <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"is_empty" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
CmpEq PrimType
int64) SubExp
w forall a b. (a -> b) -> a -> b
$ forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)
  [VName]
sizes <-
    forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [VName]
letTupExp [Char]
"partition_size" forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf (forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
is_empty) InternaliseM (Body (Rep InternaliseM))
empty_body InternaliseM (Body SOACS)
nonempty_body

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

  -- Create scratch arrays for the result.
  [VName]
blanks <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [TypeBase Shape NoUniqueness]
arr_ts forall a b. (a -> b) -> a -> b
$ \TypeBase Shape NoUniqueness
arr_t ->
    forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
"partition_dest" forall a b. (a -> b) -> a -> b
$
      forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
        PrimType -> [SubExp] -> BasicOp
Scratch (forall shape u. TypeBase shape u -> PrimType
I.elemType TypeBase Shape NoUniqueness
arr_t) (SubExp
w forall a. a -> [a] -> [a]
: forall a. Int -> [a] -> [a]
drop Int
1 (forall u. TypeBase Shape u -> [SubExp]
I.arrayDims TypeBase Shape NoUniqueness
arr_t))

  -- Now write into the result.
  Lambda SOACS
write_lam <- do
    Param (TypeBase Shape NoUniqueness)
c_param <- forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"c" (forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
    [Param (TypeBase Shape NoUniqueness)]
offset_params <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
k forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"offset" (forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
    [Param (TypeBase Shape NoUniqueness)]
value_params <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"v" forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall u. TypeBase Shape u -> TypeBase Shape u
I.rowType) [TypeBase Shape NoUniqueness]
arr_ts
    (SubExp
offset, Stms SOACS
offset_stms) <-
      forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms forall a b. (a -> b) -> a -> b
$
        [SubExp] -> SubExp -> Int -> [LParam SOACS] -> InternaliseM SubExp
mkOffsetLambdaBody
          (forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var [VName]
sizes)
          (VName -> SubExp
I.Var forall a b. (a -> b) -> a -> b
$ forall dec. Param dec -> VName
I.paramName Param (TypeBase Shape NoUniqueness)
c_param)
          Int
0
          [Param (TypeBase Shape NoUniqueness)]
offset_params
    forall (f :: * -> *) a. Applicative f => a -> f a
pure
      I.Lambda
        { lambdaParams :: [LParam SOACS]
I.lambdaParams = Param (TypeBase Shape NoUniqueness)
c_param forall a. a -> [a] -> [a]
: [Param (TypeBase Shape NoUniqueness)]
offset_params forall a. [a] -> [a] -> [a]
++ [Param (TypeBase Shape NoUniqueness)]
value_params,
          lambdaReturnType :: [TypeBase Shape NoUniqueness]
I.lambdaReturnType =
            forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase Shape NoUniqueness]
arr_ts) (forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
              forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map forall u. TypeBase Shape u -> TypeBase Shape u
I.rowType [TypeBase Shape NoUniqueness]
arr_ts,
          lambdaBody :: Body SOACS
I.lambdaBody =
            forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody Stms SOACS
offset_stms forall a b. (a -> b) -> a -> b
$
              forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase Shape NoUniqueness]
arr_ts) (SubExp -> SubExpRes
subExpRes SubExp
offset)
                forall a. [a] -> [a] -> [a]
++ [VName] -> Result
I.varsRes (forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
I.paramName [Param (TypeBase Shape NoUniqueness)]
value_params)
        }
  [VName]
results <-
    forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [VName]
letTupExp [Char]
"partition_res" forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. Op rep -> Exp rep
I.Op forall a b. (a -> b) -> a -> b
$
      forall rep.
SubExp
-> [VName] -> Lambda rep -> [(Shape, Int, VName)] -> SOAC rep
I.Scatter SubExp
w (VName
classes forall a. a -> [a] -> [a]
: [VName]
all_offsets forall a. [a] -> [a] -> [a]
++ [VName]
arrs) Lambda SOACS
write_lam forall a b. (a -> b) -> a -> b
$
        forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 (forall a. a -> [a]
repeat forall a b. (a -> b) -> a -> b
$ forall d. [d] -> ShapeBase d
I.Shape [SubExp
w]) (forall a. a -> [a]
repeat Int
1) [VName]
blanks
  SubExp
sizes' <-
    forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"partition_sizes" forall a b. (a -> b) -> a -> b
$
      forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
        [SubExp] -> TypeBase Shape NoUniqueness -> BasicOp
I.ArrayLit (forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var [VName]
sizes) forall a b. (a -> b) -> a -> b
$
          forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var [VName]
results, [SubExp
sizes'])
  where
    mkOffsetLambdaBody ::
      [SubExp] ->
      SubExp ->
      Int ->
      [I.LParam SOACS] ->
      InternaliseM SubExp
    mkOffsetLambdaBody :: [SubExp] -> SubExp -> Int -> [LParam SOACS] -> InternaliseM SubExp
mkOffsetLambdaBody [SubExp]
_ SubExp
_ Int
_ [] =
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall v. IsValue v => v -> SubExp
constant (-Int64
1 :: Int64)
    mkOffsetLambdaBody [SubExp]
sizes SubExp
c Int
i (LParam SOACS
p : [LParam SOACS]
ps) = do
      SubExp
is_this_one <-
        forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"is_this_one" forall a b. (a -> b) -> a -> b
$
          forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
CmpEq PrimType
int64) SubExp
c forall a b. (a -> b) -> a -> b
$
              IntType -> Integer -> SubExp
intConst IntType
Int64 forall a b. (a -> b) -> a -> b
$
                forall a. Integral a => a -> Integer
toInteger Int
i
      SubExp
next_one <- [SubExp] -> SubExp -> Int -> [LParam SOACS] -> InternaliseM SubExp
mkOffsetLambdaBody [SubExp]
sizes SubExp
c (Int
i forall a. Num a => a -> a -> a
+ Int
1) [LParam SOACS]
ps
      SubExp
this_one <-
        forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"this_offset"
          forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
MonadBuilder m =>
BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
foldBinOp
            (IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
OverflowUndef)
            (forall v. IsValue v => v -> SubExp
constant (-Int64
1 :: Int64))
            (VName -> SubExp
I.Var (forall dec. Param dec -> VName
I.paramName LParam SOACS
p) forall a. a -> [a] -> [a]
: forall a. Int -> [a] -> [a]
take Int
i [SubExp]
sizes)
      forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"total_res"
        forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
          (forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
is_this_one)
          (forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
this_one])
          (forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
next_one])

typeExpForError :: E.TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError :: TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError (E.TEVar QualName VName
qn SrcLoc
_) =
  forall (f :: * -> *) a. Applicative f => a -> f a
pure [forall a. Text -> ErrorMsgPart a
ErrorString forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> Text
prettyText QualName VName
qn]
typeExpForError (E.TEUnique TypeExp VName
te SrcLoc
_) =
  (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.TEDim [VName]
dims TypeExp VName
te SrcLoc
_) =
  (forall a. Text -> ErrorMsgPart a
ErrorString (Text
"?" forall a. Semigroup a => a -> a -> a
<> Text
dims' forall a. Semigroup a => a -> a -> a
<> Text
".") :) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
te
  where
    dims' :: Text
dims' = forall a. Monoid a => [a] -> a
mconcat (forall a b. (a -> b) -> [a] -> [b]
map forall a. Pretty a => a -> Text
onDim [VName]
dims)
    onDim :: a -> Text
onDim a
d = Text
"[" forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText a
d forall a. Semigroup a => a -> a -> a
<> Text
"]"
typeExpForError (E.TEArray SizeExp VName
d TypeExp VName
te SrcLoc
_) = do
  ErrorMsgPart SubExp
d' <- SizeExp VName -> InternaliseM (ErrorMsgPart SubExp)
dimExpForError SizeExp VName
d
  [ErrorMsgPart SubExp]
te' <- TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
te
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp
"[", ErrorMsgPart SubExp
d', ErrorMsgPart SubExp
"]"] forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
te'
typeExpForError (E.TETuple [TypeExp VName]
tes SrcLoc
_) = do
  [[ErrorMsgPart SubExp]]
tes' <- 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
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp
"("] forall a. [a] -> [a] -> [a]
++ forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
", "] [[ErrorMsgPart SubExp]]
tes' forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
")"]
typeExpForError (E.TERecord [(Name, TypeExp VName)]
fields SrcLoc
_) = do
  [[ErrorMsgPart SubExp]]
fields' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall {a}.
Pretty a =>
(a, TypeExp VName) -> InternaliseM [ErrorMsgPart SubExp]
onField [(Name, TypeExp VName)]
fields
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp
"{"] forall a. [a] -> [a] -> [a]
++ forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
", "] [[ErrorMsgPart SubExp]]
fields' forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"}"]
  where
    onField :: (a, TypeExp VName) -> InternaliseM [ErrorMsgPart SubExp]
onField (a
k, TypeExp VName
te) =
      (forall a. Text -> ErrorMsgPart a
ErrorString (forall a. Pretty a => a -> Text
prettyText a
k forall a. Semigroup a => a -> a -> a
<> Text
": ") :) 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
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp]
t1' forall a. [a] -> [a] -> [a]
++ [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 SizeExp VName
d SrcLoc
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SizeExp VName -> InternaliseM (ErrorMsgPart SubExp)
dimExpForError SizeExp VName
d
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp]
t' forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
" "] forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
arg'
typeExpForError (E.TESum [(Name, [TypeExp VName])]
cs SrcLoc
_) = do
  [[ErrorMsgPart SubExp]]
cs' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([TypeExp VName] -> InternaliseM [ErrorMsgPart SubExp]
onClause forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(Name, [TypeExp VName])]
cs
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ 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' <- 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
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
" "] [[ErrorMsgPart SubExp]]
c'

dimExpForError :: E.SizeExp VName -> InternaliseM (ErrorMsgPart SubExp)
dimExpForError :: SizeExp VName -> InternaliseM (ErrorMsgPart SubExp)
dimExpForError (SizeExpNamed QualName VName
d SrcLoc
_) = do
  Maybe [SubExp]
substs <- VName -> InternaliseM (Maybe [SubExp])
lookupSubst forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> vn
E.qualLeaf QualName VName
d
  SubExp
d' <- case Maybe [SubExp]
substs of
    Just [SubExp
v] -> forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
v
    Maybe [SubExp]
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> vn
E.qualLeaf QualName VName
d
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
d'
dimExpForError (SizeExpConst Int
d SrcLoc
_) =
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. Text -> ErrorMsgPart a
ErrorString forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> Text
prettyText Int
d
dimExpForError SizeExp VName
SizeExpAny = forall (f :: * -> *) a. Applicative f => a -> f a
pure ErrorMsgPart SubExp
""

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