{-# 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
import Control.Monad.Reader
import Data.Bifunctor
import Data.Foldable (toList)
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)
import Language.Futhark.TypeChecker.Types qualified as E

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

shiftRetAls :: Int -> RetAls -> RetAls
shiftRetAls :: Int -> RetAls -> RetAls
shiftRetAls Int
d (RetAls [Int]
pals [Int]
rals) = [Int] -> [Int] -> RetAls
RetAls [Int]
pals forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall a. Num a => a -> a -> a
+ Int
d) [Int]
rals

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 Info VName)
_ (Info ResRetType
rettype) [TypeParamBase VName]
tparams [PatBase Info VName ParamType]
params Exp
body Maybe DocComment
_ [AttrInfo VName]
attrs SrcLoc
loc) = do
  forall a.
[TypeParamBase VName]
-> [PatBase Info VName ParamType]
-> ([FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM a)
-> InternaliseM a
bindingFParams [TypeParamBase VName]
tparams [PatBase Info VName ParamType]
params forall a b. (a -> b) -> a -> b
$ \[FParam SOACS]
shapeparams [[Tree (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
        all_params :: [Free [] (Param DeclType)]
all_params = forall a b. (a -> b) -> [a] -> [b]
map forall (f :: * -> *) a. Applicative f => a -> f a
pure [FParam SOACS]
shapeparams forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Tree (FParam SOACS)]]
params'
        msg :: ErrorMsg a
msg = forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [ErrorMsgPart a
"Function return value does not match shape of declared return type."]

    (Body SOACS
body', [(TypeBase ExtShape Uniqueness, RetAls)]
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', [RetAls]
retals) <-
        forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first forall {u}. [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. [(a, b)] -> ([a], [b])
unzip forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall shape u.
[Tree DeclType]
-> ResRetType
-> [TypeBase shape u]
-> [(TypeBase ExtShape Uniqueness, RetAls)]
internaliseReturnType (forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall dec. DeclTyped dec => Param dec -> DeclType
paramDeclType) [Free [] (Param DeclType)]
all_params) ResRetType
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 forall {a}. ErrorMsg a
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
      let num_ctx :: Int
num_ctx = forall (t :: * -> *) a. Foldable t => t a -> Int
length (forall u. [TypeBase ExtShape u] -> Set Int
shapeContext [TypeBase ExtShape Uniqueness]
rettype')
      forall (f :: * -> *) a. Applicative f => a -> f a
pure
        ( Result
body_res',
          forall a. Int -> a -> [a]
replicate Int
num_ctx (forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64, forall a. Monoid a => a
mempty)
            forall a. [a] -> [a] -> [a]
++ forall a b. [a] -> [b] -> [(a, b)]
zip [TypeBase ExtShape Uniqueness]
rettype' (forall a b. (a -> b) -> [a] -> [b]
map (Int -> RetAls -> RetAls
shiftRetAls Int
num_ctx) [RetAls]
retals)
        )

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

    let fd :: FunDef SOACS
fd =
          forall rep.
Maybe EntryPoint
-> Attrs
-> Name
-> [(RetType rep, RetAls)]
-> [FParam rep]
-> Body rep
-> FunDef rep
I.FunDef
            forall a. Maybe a
Nothing
            Attrs
attrs'
            (VName -> Name
internaliseFunName VName
fname)
            [(TypeBase ExtShape Uniqueness, RetAls)]
rettype'
            (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (t :: * -> *) a. Foldable t => t a -> [a]
toList [Free [] (Param DeclType)]
all_params)
            Body SOACS
body'

    if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[Tree (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 :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (t :: * -> *) a. Foldable t => t a -> [a]
toList) [[Tree (FParam SOACS)]]
params',
            forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (t :: * -> *) a. Foldable t => t a -> [a]
toList [Free [] (Param DeclType)]
all_params,
            forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. [a] -> [b] -> [(a, b)]
`zip` forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(TypeBase ExtShape Uniqueness, RetAls)]
rettype')
              forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rt dec.
(IsRetType rt, Typed dec) =>
[rt]
-> [Param dec]
-> [(SubExp, TypeBase Shape NoUniqueness)]
-> Maybe [rt]
applyRetType (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(TypeBase ExtShape Uniqueness, RetAls)]
rettype') (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (t :: * -> *) a. Foldable t => t a -> [a]
toList [Free [] (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 Info VName)
_ (Info ResRetType
rettype) [TypeParamBase VName]
tparams [PatBase Info VName ParamType]
params Exp
_ Maybe DocComment
_ [AttrInfo VName]
attrs SrcLoc
loc) = ValBind
vb
  forall a.
[TypeParamBase VName]
-> [PatBase Info VName ParamType]
-> ([FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM a)
-> InternaliseM a
bindingFParams [TypeParamBase VName]
tparams [PatBase Info VName ParamType]
params forall a b. (a -> b) -> a -> b
$ \[FParam SOACS]
shapeparams [[Tree (FParam SOACS)]]
params' -> do
    let all_params :: [Free [] (Param DeclType)]
all_params = forall a b. (a -> b) -> [a] -> [b]
map forall (f :: * -> *) a. Applicative f => a -> f a
pure [FParam SOACS]
shapeparams forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Tree (FParam SOACS)]]
params'
        ([[TypeBase ExtShape Uniqueness]]
entry_rettype, [[RetAls]]
retals) =
          forall a b. [(a, b)] -> ([a], [b])
unzip forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. [(a, b)] -> ([a], [b])
unzip forall a b. (a -> b) -> a -> b
$ [Tree DeclType]
-> ResRetType -> [[(TypeBase ExtShape Uniqueness, RetAls)]]
internaliseEntryReturnType (forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall dec. DeclTyped dec => Param dec -> DeclType
paramDeclType) [Free [] (Param DeclType)]
all_params) ResRetType
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 forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (t :: * -> *) a. Foldable t => t a -> [a]
toList) [[Tree (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 :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (t :: * -> *) a. Foldable t => t a -> [a]
toList) [[Tree (FParam SOACS)]]
params'

    OpaqueTypes -> InternaliseM ()
addOpaques OpaqueTypes
opaques

    (Body SOACS
entry_body, [(TypeBase ExtShape Uniqueness, RetAls)]
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 ->
          [Char]
-> QualName VName -> [SubExp] -> SrcLoc -> InternaliseM [SubExp]
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, forall a. Monoid a => a
mempty)) [SubExp]
ctx)

    Attrs
attrs' <- [AttrInfo VName] -> InternaliseM Attrs
internaliseAttrs [AttrInfo VName]
attrs
    let num_ctx :: Int
num_ctx = forall (t :: * -> *) a. Foldable t => t a -> Int
length [(TypeBase ExtShape Uniqueness, RetAls)]
ctx_ts
    FunDef SOACS -> InternaliseM ()
addFunDef forall a b. (a -> b) -> a -> b
$
      forall rep.
Maybe EntryPoint
-> Attrs
-> Name
-> [(RetType rep, RetAls)]
-> [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, RetAls)]
ctx_ts
            forall a. [a] -> [a] -> [a]
++ forall a b. [a] -> [b] -> [(a, b)]
zip
              (forall {u}. [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts (forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[TypeBase ExtShape Uniqueness]]
entry_rettype))
              (forall a b. (a -> b) -> [a] -> [b]
map (Int -> RetAls -> RetAls
shiftRetAls Int
num_ctx) forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[RetAls]]
retals)
        )
        ([FParam SOACS]
shapeparams forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (t :: * -> *) a. Foldable t => t a -> [a]
toList) [[Tree (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 -> StructType
E.typeOf Exp
start of
        E.Scalar (E.Prim (E.Unsigned IntType
_)) -> forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntZ IntType
Int64
        StructType
_ -> 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 -> StructType
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)
      StructType
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 StructType
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 StructType
et [VName]
ext) e :: AppExp
e@E.Apply {} =
  case AppExp -> (Function, [(Exp, Maybe VName)])
findFuncall AppExp
e of
    (FunctionHole 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.  One caveat is that we need to replace any
      -- existential sizes, too (with zeroes, because they don't
      -- matter).
      let subst :: [(VName, Subst StructRetType)]
subst = forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
ext forall a b. (a -> b) -> a -> b
$ forall a. a -> [a]
repeat forall a b. (a -> b) -> a -> b
$ forall t. Exp -> Subst t
E.ExpSubst forall a b. (a -> b) -> a -> b
$ Integer -> SrcLoc -> Exp
E.sizeFromInteger Integer
0 forall a. Monoid a => a
mempty
          et' :: StructType
et' = forall a. Substitutable a => TypeSubs -> a -> a
E.applySubst (forall a b. Eq a => a -> [(a, b)] -> Maybe b
`lookup` [(VName, Subst StructRetType)]
subst) StructType
et
      [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc (forall (f :: * -> *) vn. f StructType -> SrcLoc -> ExpBase f vn
E.Hole (forall a. a -> Info a
Info StructType
et') 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
$ StructType -> [VName] -> AppRes
AppRes (forall dim u. ScalarTypeBase dim u -> TypeBase dim u
E.Scalar forall a b. (a -> b) -> a -> b
$ forall dim u. PrimType -> ScalarTypeBase dim u
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
$ StructType -> [VName] -> AppRes
AppRes (forall dim u. ScalarTypeBase dim u -> TypeBase dim u
E.Scalar forall a b. (a -> b) -> a -> b
$ forall dim u. PrimType -> ScalarTypeBase dim u
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 [(StructType, [SubExp])] -> InternaliseM [SubExp]
internalise <- QualName VName
-> [Char]
-> SrcLoc
-> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
isOverloadedFunction QualName VName
qfname [Char]
desc SrcLoc
loc -> do
              let prepareArg :: (Exp, b) -> InternaliseM (StructType, [SubExp])
prepareArg (Exp
arg, b
_) =
                    (forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
E.toStruct (Exp -> StructType
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
              [(StructType, [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 (StructType, [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, RetAls)]
-> (Safety, SrcLoc, [SrcLoc])
-> Exp rep
I.Apply Name
fname [(SubExp, Diet)]
args'' [(forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
rettype, forall a. Monoid a => a
mempty)] (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)
              [Char]
-> QualName VName -> [SubExp] -> SrcLoc -> InternaliseM [SubExp]
funcall [Char]
desc QualName VName
qfname [SubExp]
args' SrcLoc
loc
internaliseAppExp [Char]
desc AppRes
_ (E.LetPat [SizeBinder VName]
sizes PatBase Info VName StructType
pat Exp
e Exp
body SrcLoc
_) =
  forall a.
[Char]
-> [SizeBinder VName]
-> PatBase Info VName StructType
-> Exp
-> InternaliseM a
-> InternaliseM a
internalisePat [Char]
desc [SizeBinder VName]
sizes PatBase Info VName StructType
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 ParamType],
 Maybe (TypeExp Info VName), Info ResRetType, 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 ParamType
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 StructType
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 ParamType
-> [TypeBase Shape NoUniqueness]
-> ([FParam SOACS] -> [FParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' PatBase Info VName ParamType
mergepat [TypeBase Shape NoUniqueness]
ts forall a b. (a -> b) -> a -> b
$ \[FParam SOACS]
shapepat [FParam SOACS]
mergepat' ->
        forall a.
[PatBase Info VName ParamType]
-> [TypeBase Shape NoUniqueness]
-> ([LParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLambdaParams [forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
E.Observe forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PatBase Info VName StructType
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 StructType
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 ParamType
-> [TypeBase Shape NoUniqueness]
-> ([FParam SOACS] -> [FParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' PatBase Info VName ParamType
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 {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
E.identName IdentBase Info VName StructType
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 ParamType
-> [TypeBase Shape NoUniqueness]
-> ([FParam SOACS] -> [FParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' PatBase Info VName ParamType
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

          -- As the condition expression is inserted twice, we have to
          -- avoid shadowing (#1935).
          (Stms SOACS
cond_stms, SubExp
cond') <-
            forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall (m :: * -> *) rep a.
(MonadFreshNames m, Renameable rep, Rename a) =>
Stms rep -> a -> m (Stms rep, a)
renameStmsWith)
              forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms ([Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"loop_cond" Exp
cond)
          forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
addStms Stms SOACS
cond_stms
          forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
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 StructType
name IdentBase Info VName StructType
src SliceBase Info VName
idxs Exp
ve Exp
body SrcLoc
loc) = do
  let pat :: PatBase Info VName StructType
pat = forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
E.Id (forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
E.identName IdentBase Info VName StructType
name) (forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> f t
E.identType IdentBase Info VName StructType
name) SrcLoc
loc
      src_t :: Info StructType
src_t = forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> f t
E.identType IdentBase Info VName StructType
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 StructType -> SrcLoc -> ExpBase f vn
E.Var (forall v. v -> QualName v
E.qualName forall a b. (a -> b) -> a -> b
$ forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
E.identName IdentBase Info VName StructType
src) Info StructType
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 StructType
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
E.LetPat [] PatBase Info VName StructType
pat Exp
e Exp
body SrcLoc
loc)
      (forall a. a -> Info a
Info (StructType -> [VName] -> AppRes
AppRes (Exp -> StructType
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 StructType
p Exp
case_e SrcLoc
_) = do
      ([Maybe PrimValue]
cmps, [SubExp]
pertinent) <- PatBase Info VName StructType
-> [SubExp] -> InternaliseM ([Maybe PrimValue], [SubExp])
generateCond PatBase Info VName StructType
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 StructType
-> [SubExp]
-> InternaliseM a
-> InternaliseM a
internalisePat' [] PatBase Info VName StructType
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 StructType
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 StructType
t)
      ts :: [TypeBase ExtShape Uniqueness]
ts = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (t :: * -> *) a. Foldable t => t a -> [a]
toList forall a b. (a -> b) -> a -> b
$ StructType -> [Tree (TypeBase ExtShape Uniqueness)]
internaliseType (forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
E.toStruct StructType
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 StructType
_ 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 StructType
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 StructType -> SrcLoc -> ExpBase f vn
E.Var (forall v. v -> QualName v
E.qualName VName
name) Info StructType
t SrcLoc
loc)
          SrcLoc
loc
internaliseExp [Char]
desc (E.ArrayLit [Exp]
es (Info StructType
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 StructType
basetype <- forall dim u. Int -> TypeBase dim u -> Maybe (TypeBase dim u)
E.peelArray (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
eshape) StructType
arr_t = do
      let flat_lit :: Exp
flat_lit = forall (f :: * -> *) vn.
[ExpBase f vn] -> f StructType -> 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 StructType
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 = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (t :: * -> *) a. Foldable t => t a -> [a]
toList forall a b. (a -> b) -> a -> b
$ StructType -> [Tree (TypeBase ExtShape Uniqueness)]
internaliseType forall a b. (a -> b) -> a -> b
$ forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
E.toStruct StructType
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 StructType
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 StructType
_ 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 Info VName
_ SrcLoc
_) =
  [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
internaliseExp [Char]
desc (E.Coerce Exp
e TypeExp Info VName
_ (Info StructType
et) SrcLoc
loc) = do
  [SubExp]
ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
  [TypeBase ExtShape Uniqueness]
ts <- forall shape u.
StructType -> [TypeBase shape u] -> [TypeBase ExtShape Uniqueness]
internaliseCoerceType (forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
E.toStruct StructType
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' <- forall u. TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError forall a b. (a -> b) -> a -> b
$ forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct StructType
et
  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'
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 StructType
_ 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 Exp als -> [Name] -> [a] -> [a] -> m [a]
replace (Exp -> StructType
E.typeOf Exp
src) [Name]
fields [SubExp]
ve' [SubExp]
src'
  where
    replace :: TypeBase Exp als -> [Name] -> [a] -> [a] -> m [a]
replace (E.Scalar (E.Record Map Name (TypeBase Exp als)
m)) (Name
f : [Name]
fs) [a]
ve' [a]
src'
      | Just TypeBase Exp als
t <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f Map Name (TypeBase Exp 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 Exp 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 Exp als)
m
              k :: Int
k = forall als. TypeBase Exp als -> Int
internalisedTypeSize TypeBase Exp 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 Exp als -> [Name] -> [a] -> [a] -> m [a]
replace TypeBase Exp 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 Exp 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 [StructType]
fs))) SrcLoc
_) = do
  ([TypeBase ExtShape Uniqueness]
ts, Map Name (Int, [Int])
constr_map) <- Map Name [StructType]
-> 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 u. TypeBase dim u -> TypeBase dim NoUniqueness
E.toStruct) Map Name [StructType]
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 StructType
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 StructType
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 StructType
t) SrcLoc
_) =
  case StructType
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]
    StructType
_ -> 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 StructType
t
internaliseExp [Char]
_ (E.FloatLit Double
v (Info StructType
t) SrcLoc
_) =
  case StructType
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]
    StructType
_ -> 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 StructType
t
-- Builtin operators are handled specially because they are
-- overloaded.
internaliseExp [Char]
desc (E.Project Name
k Exp
e (Info StructType
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 Exp als -> Int
internalisedTypeSize forall a b. (a -> b) -> a -> b
$
        case Exp -> StructType
E.typeOf Exp
e of
          E.Scalar (Record Map Name StructType
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 StructType
fs
          StructType
t -> [StructType
t]
  forall a. Int -> [a] -> [a]
take (forall als. TypeBase Exp als -> Int
internalisedTypeSize StructType
rt) 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.StructType -> I.PrimValue
internalisePatLit :: PatLit -> StructType -> PrimValue
internalisePatLit (E.PatLitPrim PrimValue
v) StructType
_ =
  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 StructType
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, StructType
t)

generateCond ::
  E.Pat StructType ->
  [I.SubExp] ->
  InternaliseM ([Maybe I.PrimValue], [I.SubExp])
generateCond :: PatBase Info VName StructType
-> [SubExp] -> InternaliseM ([Maybe PrimValue], [SubExp])
generateCond PatBase Info VName StructType
orig_p [SubExp]
orig_ses = do
  ([Maybe PrimValue]
cmps, [SubExp]
pertinent, [SubExp]
_) <- forall {vn} {a}.
(Eq vn, IsName vn) =>
PatBase Info vn StructType
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares PatBase Info VName StructType
orig_p [SubExp]
orig_ses
  forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Maybe PrimValue]
cmps, [SubExp]
pertinent)
  where
    compares :: PatBase Info vn StructType
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares (E.PatLit PatLit
l (Info StructType
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 -> StructType -> PrimValue
internalisePatLit PatLit
l StructType
t], [a
se], [a]
ses)
    compares (E.PatConstr Name
c (Info (E.Scalar (E.Sum Map Name [StructType]
fs))) [PatBase Info vn StructType]
pats SrcLoc
_) (a
_ : [a]
ses) = do
      ([TypeBase ExtShape Uniqueness]
payload_ts, Map Name (Int, [Int])
m) <- Map Name [StructType]
-> 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 u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct) Map Name [StructType]
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 StructType]
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
comparesMany [PatBase Info vn StructType]
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 StructType
t) [PatBase Info vn StructType]
_ 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 StructType
t
    compares (E.Id vn
_ Info StructType
t SrcLoc
loc) [a]
ses =
      PatBase Info vn StructType
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares (forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
E.Wildcard Info StructType
t SrcLoc
loc) [a]
ses
    compares (E.Wildcard (Info StructType
t) SrcLoc
_) [a]
ses = do
      let ([a]
id_ses, [a]
rest_ses) = forall a. Int -> [a] -> ([a], [a])
splitAt (forall als. TypeBase Exp als -> Int
internalisedTypeSize forall a b. (a -> b) -> a -> b
$ forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
E.toStruct StructType
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 StructType
pat SrcLoc
_) [a]
ses =
      PatBase Info vn StructType
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares PatBase Info vn StructType
pat [a]
ses
    compares (E.PatAttr AttrInfo vn
_ PatBase Info vn StructType
pat SrcLoc
_) [a]
ses =
      PatBase Info vn StructType
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares PatBase Info vn StructType
pat [a]
ses
    compares (E.TuplePat [] SrcLoc
loc) [a]
ses =
      PatBase Info vn StructType
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares (forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
E.Wildcard (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ forall dim u. ScalarTypeBase dim u -> TypeBase dim u
E.Scalar forall a b. (a -> b) -> a -> b
$ forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
E.Record forall a. Monoid a => a
mempty) SrcLoc
loc) [a]
ses
    compares (E.RecordPat [] SrcLoc
loc) [a]
ses =
      PatBase Info vn StructType
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares (forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
E.Wildcard (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ forall dim u. ScalarTypeBase dim u -> TypeBase dim u
E.Scalar forall a b. (a -> b) -> a -> b
$ forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
E.Record forall a. Monoid a => a
mempty) SrcLoc
loc) [a]
ses
    compares (E.TuplePat [PatBase Info vn StructType]
pats SrcLoc
_) [a]
ses =
      [PatBase Info vn StructType]
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
comparesMany [PatBase Info vn StructType]
pats [a]
ses
    compares (E.RecordPat [(Name, PatBase Info vn StructType)]
fs SrcLoc
_) [a]
ses =
      [PatBase Info vn StructType]
-> [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 StructType)]
fs) [a]
ses
    compares (E.PatAscription PatBase Info vn StructType
pat TypeExp Info vn
_ SrcLoc
_) [a]
ses =
      PatBase Info vn StructType
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares PatBase Info vn StructType
pat [a]
ses
    compares PatBase Info vn StructType
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 StructType
pat

    comparesMany :: [PatBase Info vn StructType]
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
comparesMany [] [a]
ses = forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], [], [a]
ses)
    comparesMany (PatBase Info vn StructType
pat : [PatBase Info vn StructType]
pats) [a]
ses = do
      ([Maybe PrimValue]
cmps1, [a]
pertinent1, [a]
ses') <- PatBase Info vn StructType
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares PatBase Info vn StructType
pat [a]
ses
      ([Maybe PrimValue]
cmps2, [a]
pertinent2, [a]
ses'') <- [PatBase Info vn StructType]
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
comparesMany [PatBase Info vn StructType]
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 StructType ->
  E.Exp ->
  InternaliseM a ->
  InternaliseM a
internalisePat :: forall a.
[Char]
-> [SizeBinder VName]
-> PatBase Info VName StructType
-> Exp
-> InternaliseM a
-> InternaliseM a
internalisePat [Char]
desc [SizeBinder VName]
sizes PatBase Info VName StructType
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 StructType
-> [SubExp]
-> InternaliseM a
-> InternaliseM a
internalisePat' [SizeBinder VName]
sizes PatBase Info VName StructType
p [SubExp]
ses InternaliseM a
m
  where
    desc' :: [Char]
desc' = case forall t. Pat t -> [Ident t]
E.patIdents PatBase Info VName StructType
p of
      [IdentBase Info VName StructType
v] -> VName -> [Char]
baseString forall a b. (a -> b) -> a -> b
$ forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
E.identName IdentBase Info VName StructType
v
      [IdentBase Info VName StructType]
_ -> [Char]
desc

internalisePat' ::
  [E.SizeBinder VName] ->
  E.Pat StructType ->
  [I.SubExp] ->
  InternaliseM a ->
  InternaliseM a
internalisePat' :: forall a.
[SizeBinder VName]
-> PatBase Info VName StructType
-> [SubExp]
-> InternaliseM a
-> InternaliseM a
internalisePat' [SizeBinder VName]
sizes PatBase Info VName StructType
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 ParamType
-> [TypeBase Shape NoUniqueness]
-> ([VName] -> InternaliseM a)
-> InternaliseM a
stmPat (forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
E.Observe forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PatBase Info VName StructType
p) [TypeBase Shape NoUniqueness]
ses_ts forall a b. (a -> b) -> a -> b
$ \[VName]
pat_names -> do
    AppRes -> [SubExp] -> InternaliseM ()
bindExtSizes (StructType -> [VName] -> AppRes
AppRes (forall d u. Pat (TypeBase d u) -> TypeBase d u
E.patternType PatBase Info VName StructType
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 StructType
_ 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 :: m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
divRounding m (Exp (Rep m))
x m (Exp (Rep m))
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)
              m (Exp (Rep m))
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) m (Exp (Rep m))
y (forall (m :: * -> *).
MonadBuilder m =>
m (Exp (Rep m)) -> m (Exp (Rep m))
eSignum m (Exp (Rep m))
y))
          )
          m (Exp (Rep m))
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
=<< forall {m :: * -> *}.
MonadBuilder m =>
m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
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
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
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 -> StructType
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'
    StructType
_ -> 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 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 NonEmpty (Info (Diet, Maybe VName), Exp)
args SrcLoc
_)
  | E.Var QualName VName
fname Info StructType
_ SrcLoc
_ <- Exp
f =
      (QualName VName -> Function
FunctionName QualName VName
fname, forall a b. (a -> b) -> [a] -> [b]
map forall {a} {b} {a}. (Info (a, b), a) -> (a, b)
onArg forall a b. (a -> b) -> a -> b
$ forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Info (Diet, Maybe VName), Exp)
args)
  | E.Hole (Info StructType
_) SrcLoc
loc <- Exp
f =
      (SrcLoc -> Function
FunctionHole SrcLoc
loc, forall a b. (a -> b) -> [a] -> [b]
map forall {a} {b} {a}. (Info (a, b), a) -> (a, b)
onArg forall a b. (a -> b) -> a -> b
$ forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Info (Diet, Maybe VName), Exp)
args)
  where
    onArg :: (Info (a, b), a) -> (a, b)
onArg (Info (a
_, b
argext), a
e) = (a
e, b
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 ParamType]
params Exp
body Maybe (TypeExp Info VName)
_ (Info (RetType [VName]
_ TypeBase Exp Uniqueness
rettype)) SrcLoc
_) [TypeBase Shape NoUniqueness]
rowtypes =
  forall a.
[PatBase Info VName ParamType]
-> [TypeBase Shape NoUniqueness]
-> ([LParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLambdaParams [PatBase Info VName ParamType]
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 Exp Uniqueness
-> [TypeBase shape u] -> InternaliseM [TypeBase Shape NoUniqueness]
internaliseLambdaReturnType TypeBase Exp Uniqueness
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 ([(StructType, [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 ([(StructType, [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 ([(StructType, [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
$ \[(StructType
_, [SubExp]
xe'), (StructType
_, [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
$ \[(StructType
x_t, [SubExp
x']), (StructType
y_t, [SubExp
y'])] ->
            case (StructType
x_t, StructType
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
              (StructType, StructType)
_ -> 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 [Exp
lam, Exp
arr] [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 [Exp
k, Exp
lam, Exp
arr] [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 [Exp
lam, Exp
ne, Exp
arr] [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 [Exp
lam, Exp
ne, Exp
arr] [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 [Exp
lam, Exp
ne, Exp
arr] [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 [Exp
rf, Exp
dest, Exp
op, Exp
ne, Exp
buckets, Exp
img] [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 [Exp
rf, Exp
dest, Exp
op, Exp
ne, Exp
buckets, Exp
img] [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 [Exp
rf, Exp
dest, Exp
op, Exp
ne, Exp
buckets, Exp
img] [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 [Exp
dest, Exp
f, Exp
bs] 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 [Exp
dest, Exp
op, Exp
ne, Exp
f, Exp
bs] 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 [Exp
acc, Exp
i, Exp
v] 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 [Exp
f, Exp
x, Exp
v] 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 [Exp
a, Exp
si, Exp
v] [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 [Exp
a, Exp
si, Exp
v] [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 [Exp
a, Exp
si, Exp
v] [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 [Exp
n, Exp
m, Exp
arr] [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
      -- Each dimension must be nonnegative, and 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 :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
toExp forall a b. (a -> b) -> a -> b
$
          SubExp -> TPrimExp Int64 VName
pe64 SubExp
old_dim forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.==. SubExp -> TPrimExp Int64 VName
pe64 SubExp
n' forall a. Num a => a -> a -> a
* SubExp -> TPrimExp Int64 VName
pe64 SubExp
m'
            forall v.
Eq v =>
TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v
.&&. SubExp -> TPrimExp Int64 VName
pe64 SubExp
n' forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.>=. TPrimExp Int64 VName
0
            forall v.
Eq v =>
TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v
.&&. SubExp -> TPrimExp Int64 VName
pe64 SubExp
m' forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.>=. TPrimExp Int64 VName
0
      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 [Exp
x, Exp
y] [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 [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 [Exp
x, Exp
y] [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
. Shape -> SubExp -> BasicOp
Replicate forall a. Monoid a => a
mempty forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
I.Var)
        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
$ \[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]
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
. Shape -> SubExp -> BasicOp
Replicate forall a. Monoid a => a
mempty forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
I.Var)
        forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
desc Exp
x
    handleRest [Exp
arr, Exp
offset, Exp
n1, Exp
s1, Exp
n2, Exp
s2] [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 [Exp
arr1, Exp
offset, Exp
s1, Exp
s2, Exp
arr2] [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 [Exp
arr, Exp
offset, Exp
n1, Exp
s1, Exp
n2, Exp
s2, Exp
n3, Exp
s3] [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 [Exp
arr1, Exp
offset, Exp
s1, Exp
s2, Exp
s3, Exp
arr2] [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 [Exp
arr, Exp
offset, Exp
n1, Exp
s1, Exp
n2, Exp
s2, Exp
n3, Exp
s3, Exp
n4, Exp
s4] [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 [Exp
arr1, Exp
offset, Exp
s1, Exp
s2, Exp
s3, Exp
s4, Exp
arr2] [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 -> StructType
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'
        StructType
_ -> 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 -> StructType
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'
        StructType
_ -> 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 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]
funcall :: [Char]
-> QualName VName -> [SubExp] -> SrcLoc -> InternaliseM [SubExp]
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, RetAls)]
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, RetAls)]
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, RetAls)]
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, RetAls)]
ts -> do
      Safety
safety <- InternaliseM Safety
askSafety
      Attrs
attrs <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks InternaliseEnv -> Attrs
envAttrs
      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, RetAls)]
-> (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, RetAls)]
ts (Safety
safety, SrcLoc
loc, forall a. Monoid a => a
mempty)

-- 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 StructType
ret [VName]
retext) [SubExp]
ses = do
  let ts :: [TypeBase ExtShape Uniqueness]
ts = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (t :: * -> *) a. Foldable t => t a -> [a]
toList forall a b. (a -> b) -> a -> b
$ StructType -> [Tree (TypeBase ExtShape Uniqueness)]
internaliseType forall a b. (a -> b) -> a -> b
$ forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
E.toStruct StructType
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])

sizeExpForError :: E.Size -> InternaliseM [ErrorMsgPart SubExp]
sizeExpForError :: Exp -> InternaliseM [ErrorMsgPart SubExp]
sizeExpForError Exp
e = do
  SubExp
e' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"size" Exp
e
  forall (f :: * -> *) a. Applicative f => a -> f a
pure [ErrorMsgPart SubExp
"[", forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
e', ErrorMsgPart SubExp
"]"]

typeExpForError :: E.TypeBase Size u -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError :: forall u. TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError (E.Scalar (E.Prim PrimType
t)) = 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 PrimType
t]
typeExpForError (E.Scalar (E.TypeVar u
_ QualName VName
v [TypeArg Exp]
args)) = do
  [ErrorMsgPart SubExp]
args' <- 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 TypeArg Exp -> InternaliseM [ErrorMsgPart SubExp]
onArg [TypeArg Exp]
args
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> [a] -> [a]
intersperse ErrorMsgPart SubExp
" " forall a b. (a -> b) -> a -> b
$ forall a. Text -> ErrorMsgPart a
ErrorString (forall a. Pretty a => a -> Text
prettyText QualName VName
v) forall a. a -> [a] -> [a]
: [ErrorMsgPart SubExp]
args'
  where
    onArg :: TypeArg Exp -> InternaliseM [ErrorMsgPart SubExp]
onArg (TypeArgDim Exp
d) = Exp -> InternaliseM [ErrorMsgPart SubExp]
sizeExpForError Exp
d
    onArg (TypeArgType StructType
t) = forall u. TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError StructType
t
typeExpForError (E.Scalar (E.Record Map Name (TypeBase Exp u)
fs))
  | Just [TypeBase Exp u]
ts <- forall a. Map Name a -> Maybe [a]
E.areTupleFields Map Name (TypeBase Exp u)
fs = do
      [[ErrorMsgPart SubExp]]
ts' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall u. TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError [TypeBase Exp u]
ts
      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]]
ts' forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
")"]
  | Bool
otherwise = do
      [[ErrorMsgPart SubExp]]
fs' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall {a} {u}.
Pretty a =>
(a, TypeBase Exp u) -> InternaliseM [ErrorMsgPart SubExp]
onField forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
M.toList Map Name (TypeBase Exp u)
fs
      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]]
fs' forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"}"]
  where
    onField :: (a, TypeBase Exp u) -> InternaliseM [ErrorMsgPart SubExp]
onField (a
k, TypeBase Exp u
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
<$> forall u. TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeBase Exp u
te
typeExpForError (E.Array u
_ Shape Exp
shape ScalarTypeBase Exp NoUniqueness
et) = do
  [ErrorMsgPart SubExp]
shape' <- forall a. Monoid a => [a] -> a
mconcat 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 Exp -> InternaliseM [ErrorMsgPart SubExp]
sizeExpForError (forall dim. Shape dim -> [dim]
E.shapeDims Shape Exp
shape)
  [ErrorMsgPart SubExp]
et' <- forall u. TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError forall a b. (a -> b) -> a -> b
$ forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar ScalarTypeBase Exp NoUniqueness
et
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp]
shape' forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
et'
typeExpForError (E.Scalar (E.Sum Map Name [TypeBase Exp u]
cs)) = do
  [[ErrorMsgPart SubExp]]
cs' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall {a} {u}.
Pretty a =>
(a, [TypeBase Exp u]) -> InternaliseM [ErrorMsgPart SubExp]
onConstructor forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
M.toList Map Name [TypeBase Exp u]
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
    onConstructor :: (a, [TypeBase Exp u]) -> InternaliseM [ErrorMsgPart SubExp]
onConstructor (a
c, [TypeBase Exp u]
ts) = do
      [[ErrorMsgPart SubExp]]
ts' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall u. TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError [TypeBase Exp u]
ts
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. Text -> ErrorMsgPart a
ErrorString (Text
"#" forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText a
c forall a. Semigroup a => a -> a -> a
<> Text
" ") forall a. a -> [a] -> [a]
: forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
" "] [[ErrorMsgPart SubExp]]
ts'
typeExpForError (E.Scalar Arrow {}) = forall (f :: * -> *) a. Applicative f => a -> f a
pure [ErrorMsgPart SubExp
"#<fun>"]

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