{-# 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 (lookupWithIndex, 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) <-
    Bool
-> InternaliseM () -> m (OpaqueTypes, Stms SOACS, [FunDef SOACS])
forall (m :: * -> *).
MonadFreshNames m =>
Bool
-> InternaliseM () -> m (OpaqueTypes, Stms SOACS, [FunDef SOACS])
runInternaliseM Bool
always_safe (VisibleTypes -> [ValBind] -> InternaliseM ()
internaliseValBinds VisibleTypes
types [ValBind]
vbinds)
  Prog SOACS -> m (Prog SOACS)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Prog rep -> m (Prog rep)
I.renameProg (Prog SOACS -> m (Prog SOACS)) -> Prog SOACS -> m (Prog SOACS)
forall a b. (a -> b) -> a -> b
$ OpaqueTypes -> Stms SOACS -> [FunDef SOACS] -> Prog SOACS
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 = (ValBind -> InternaliseM ()) -> [ValBind] -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((ValBind -> InternaliseM ()) -> [ValBind] -> InternaliseM ())
-> (ValBind -> InternaliseM ()) -> [ValBind] -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ VisibleTypes -> ValBind -> InternaliseM ()
internaliseValBind VisibleTypes
types

internaliseFunName :: VName -> Name
internaliseFunName :: VName -> Name
internaliseFunName = [Char] -> Name
nameFromString ([Char] -> Name) -> (VName -> [Char]) -> VName -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> [Char]
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 ([Int] -> RetAls) -> [Int] -> RetAls
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> Int
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 Exp VName)
_ (Info ResRetType
rettype) [TypeParamBase VName]
tparams [PatBase Info VName ParamType]
params Exp
body Maybe DocComment
_ [AttrInfo VName]
attrs SrcLoc
loc) = do
  [TypeParamBase VName]
-> [PatBase Info VName ParamType]
-> ([FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM ())
-> InternaliseM ()
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 (([FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM ())
 -> InternaliseM ())
-> ([FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM ())
-> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \[FParam SOACS]
shapeparams [[Tree (FParam SOACS)]]
params' -> do
    let shapenames :: [VName]
shapenames = (Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
[FParam SOACS]
shapeparams
        all_params :: [Free [] (Param DeclType)]
all_params = (Param DeclType -> Free [] (Param DeclType))
-> [Param DeclType] -> [Free [] (Param DeclType)]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> Free [] (Param DeclType)
forall a. a -> Free [] a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Param DeclType]
[FParam SOACS]
shapeparams [Free [] (Param DeclType)]
-> [Free [] (Param DeclType)] -> [Free [] (Param DeclType)]
forall a. [a] -> [a] -> [a]
++ [[Free [] (Param DeclType)]] -> [Free [] (Param DeclType)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Free [] (Param DeclType)]]
[[Tree (FParam SOACS)]]
params'
        msg :: ErrorMsg SubExp
msg =
          [ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg
            [ ErrorMsgPart SubExp
"Internal runtime error.\n",
              ErrorMsgPart SubExp
"Return value of ",
              Text -> ErrorMsgPart SubExp
forall a. Text -> ErrorMsgPart a
ErrorString (VName -> Text
forall a. Pretty a => a -> Text
prettyText VName
fname),
              ErrorMsgPart SubExp
" does not match type shape.\n",
              ErrorMsgPart SubExp
"This is a bug in the Futhark compiler. Please report this:\n",
              ErrorMsgPart SubExp
"  https://github.com/diku-dk/futhark/issues"
            ]

    (Body SOACS
body', [(TypeBase ExtShape Uniqueness, RetAls)]
rettype') <- InternaliseM (Result, [(TypeBase ExtShape Uniqueness, RetAls)])
-> InternaliseM
     (Body (Rep InternaliseM), [(TypeBase ExtShape Uniqueness, RetAls)])
forall (m :: * -> *) a.
MonadBuilder m =>
m (Result, a) -> m (Body (Rep m), a)
buildBody (InternaliseM (Result, [(TypeBase ExtShape Uniqueness, RetAls)])
 -> InternaliseM
      (Body (Rep InternaliseM),
       [(TypeBase ExtShape Uniqueness, RetAls)]))
-> InternaliseM (Result, [(TypeBase ExtShape Uniqueness, RetAls)])
-> InternaliseM
     (Body (Rep InternaliseM), [(TypeBase ExtShape Uniqueness, RetAls)])
forall a b. (a -> b) -> a -> b
$ do
      [SubExp]
body_res <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp (VName -> [Char]
baseString VName
fname [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"_res") Exp
body
      ([TypeBase ExtShape Uniqueness]
rettype', [RetAls]
retals) <-
        ([TypeBase ExtShape Uniqueness] -> [TypeBase ExtShape Uniqueness])
-> ([TypeBase ExtShape Uniqueness], [RetAls])
-> ([TypeBase ExtShape Uniqueness], [RetAls])
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first [TypeBase ExtShape Uniqueness] -> [TypeBase ExtShape Uniqueness]
forall {u}. [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts (([TypeBase ExtShape Uniqueness], [RetAls])
 -> ([TypeBase ExtShape Uniqueness], [RetAls]))
-> ([TypeBase Shape NoUniqueness]
    -> ([TypeBase ExtShape Uniqueness], [RetAls]))
-> [TypeBase Shape NoUniqueness]
-> ([TypeBase ExtShape Uniqueness], [RetAls])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(TypeBase ExtShape Uniqueness, RetAls)]
-> ([TypeBase ExtShape Uniqueness], [RetAls])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(TypeBase ExtShape Uniqueness, RetAls)]
 -> ([TypeBase ExtShape Uniqueness], [RetAls]))
-> ([TypeBase Shape NoUniqueness]
    -> [(TypeBase ExtShape Uniqueness, RetAls)])
-> [TypeBase Shape NoUniqueness]
-> ([TypeBase ExtShape Uniqueness], [RetAls])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Tree DeclType]
-> ResRetType
-> [TypeBase Shape NoUniqueness]
-> [(TypeBase ExtShape Uniqueness, RetAls)]
forall shape u.
[Tree DeclType]
-> ResRetType
-> [TypeBase shape u]
-> [(TypeBase ExtShape Uniqueness, RetAls)]
internaliseReturnType ((Free [] (Param DeclType) -> Tree DeclType)
-> [Free [] (Param DeclType)] -> [Tree DeclType]
forall a b. (a -> b) -> [a] -> [b]
map ((Param DeclType -> DeclType)
-> Free [] (Param DeclType) -> Tree DeclType
forall a b. (a -> b) -> Free [] a -> Free [] b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Param DeclType -> DeclType
forall dec. DeclTyped dec => Param dec -> DeclType
paramDeclType) [Free [] (Param DeclType)]
all_params) ResRetType
rettype
          ([TypeBase Shape NoUniqueness]
 -> ([TypeBase ExtShape Uniqueness], [RetAls]))
-> InternaliseM [TypeBase Shape NoUniqueness]
-> InternaliseM ([TypeBase ExtShape Uniqueness], [RetAls])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
body_res

      Bool -> InternaliseM () -> InternaliseM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ([[Free [] (Param DeclType)]] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[Free [] (Param DeclType)]]
[[Tree (FParam SOACS)]]
params') (InternaliseM () -> InternaliseM ())
-> InternaliseM () -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
        AppRes -> [SubExp] -> InternaliseM ()
bindExtSizes (StructType -> [VName] -> AppRes
E.AppRes (TypeBase Exp Uniqueness -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
E.toStruct (TypeBase Exp Uniqueness -> StructType)
-> TypeBase Exp Uniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ ResRetType -> TypeBase Exp Uniqueness
forall dim as. RetTypeBase dim as -> TypeBase dim as
E.retType ResRetType
rettype) (ResRetType -> [VName]
forall dim as. RetTypeBase dim as -> [VName]
E.retDims ResRetType
rettype)) [SubExp]
body_res

      Result
body_res' <-
        ErrorMsg SubExp
-> SrcLoc -> [ExtType] -> Result -> InternaliseM Result
ensureResultExtShape ErrorMsg SubExp
msg SrcLoc
loc ((TypeBase ExtShape Uniqueness -> ExtType)
-> [TypeBase ExtShape Uniqueness] -> [ExtType]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase ExtShape Uniqueness -> ExtType
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl [TypeBase ExtShape Uniqueness]
rettype') (Result -> InternaliseM Result) -> Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Result
subExpsRes [SubExp]
body_res
      let num_ctx :: Int
num_ctx = Set Int -> Int
forall a. Set a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([TypeBase ExtShape Uniqueness] -> Set Int
forall u. [TypeBase ExtShape u] -> Set Int
shapeContext [TypeBase ExtShape Uniqueness]
rettype')
      (Result, [(TypeBase ExtShape Uniqueness, RetAls)])
-> InternaliseM (Result, [(TypeBase ExtShape Uniqueness, RetAls)])
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        ( Result
body_res',
          Int
-> (TypeBase ExtShape Uniqueness, RetAls)
-> [(TypeBase ExtShape Uniqueness, RetAls)]
forall a. Int -> a -> [a]
replicate Int
num_ctx (PrimType -> TypeBase ExtShape Uniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64, RetAls
forall a. Monoid a => a
mempty)
            [(TypeBase ExtShape Uniqueness, RetAls)]
-> [(TypeBase ExtShape Uniqueness, RetAls)]
-> [(TypeBase ExtShape Uniqueness, RetAls)]
forall a. [a] -> [a] -> [a]
++ [TypeBase ExtShape Uniqueness]
-> [RetAls] -> [(TypeBase ExtShape Uniqueness, RetAls)]
forall a b. [a] -> [b] -> [(a, b)]
zip [TypeBase ExtShape Uniqueness]
rettype' ((RetAls -> RetAls) -> [RetAls] -> [RetAls]
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 =
          Maybe EntryPoint
-> Attrs
-> Name
-> [(RetType SOACS, RetAls)]
-> [FParam SOACS]
-> Body SOACS
-> FunDef SOACS
forall rep.
Maybe EntryPoint
-> Attrs
-> Name
-> [(RetType rep, RetAls)]
-> [FParam rep]
-> Body rep
-> FunDef rep
I.FunDef
            Maybe EntryPoint
forall a. Maybe a
Nothing
            Attrs
attrs'
            (VName -> Name
internaliseFunName VName
fname)
            [(TypeBase ExtShape Uniqueness, RetAls)]
[(RetType SOACS, RetAls)]
rettype'
            ((Free [] (Param DeclType) -> [Param DeclType])
-> [Free [] (Param DeclType)] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Free [] (Param DeclType) -> [Param DeclType]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList [Free [] (Param DeclType)]
all_params)
            Body SOACS
body'

    if [[Free [] (Param DeclType)]] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[Free [] (Param DeclType)]]
[[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,
            (Param DeclType -> DeclType) -> [Param DeclType] -> [DeclType]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> DeclType
forall t. DeclTyped t => t -> DeclType
declTypeOf ([Param DeclType] -> [DeclType]) -> [Param DeclType] -> [DeclType]
forall a b. (a -> b) -> a -> b
$ ([Free [] (Param DeclType)] -> [Param DeclType])
-> [[Free [] (Param DeclType)]] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((Free [] (Param DeclType) -> [Param DeclType])
-> [Free [] (Param DeclType)] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Free [] (Param DeclType) -> [Param DeclType]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList) [[Free [] (Param DeclType)]]
[[Tree (FParam SOACS)]]
params',
            (Free [] (Param DeclType) -> [Param DeclType])
-> [Free [] (Param DeclType)] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Free [] (Param DeclType) -> [Param DeclType]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList [Free [] (Param DeclType)]
all_params,
            ([TypeBase ExtShape Uniqueness]
 -> [(TypeBase ExtShape Uniqueness, RetAls)])
-> Maybe [TypeBase ExtShape Uniqueness]
-> Maybe [(TypeBase ExtShape Uniqueness, RetAls)]
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([TypeBase ExtShape Uniqueness]
-> [RetAls] -> [(TypeBase ExtShape Uniqueness, RetAls)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` ((TypeBase ExtShape Uniqueness, RetAls) -> RetAls)
-> [(TypeBase ExtShape Uniqueness, RetAls)] -> [RetAls]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase ExtShape Uniqueness, RetAls) -> RetAls
forall a b. (a, b) -> b
snd [(TypeBase ExtShape Uniqueness, RetAls)]
rettype')
              (Maybe [TypeBase ExtShape Uniqueness]
 -> Maybe [(TypeBase ExtShape Uniqueness, RetAls)])
-> ([(SubExp, TypeBase Shape NoUniqueness)]
    -> Maybe [TypeBase ExtShape Uniqueness])
-> [(SubExp, TypeBase Shape NoUniqueness)]
-> Maybe [(TypeBase ExtShape Uniqueness, RetAls)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TypeBase ExtShape Uniqueness]
-> [Param DeclType]
-> [(SubExp, TypeBase Shape NoUniqueness)]
-> Maybe [TypeBase ExtShape Uniqueness]
forall dec.
Typed dec =>
[TypeBase ExtShape Uniqueness]
-> [Param dec]
-> [(SubExp, TypeBase Shape NoUniqueness)]
-> Maybe [TypeBase ExtShape Uniqueness]
forall rt dec.
(IsRetType rt, Typed dec) =>
[rt]
-> [Param dec]
-> [(SubExp, TypeBase Shape NoUniqueness)]
-> Maybe [rt]
applyRetType (((TypeBase ExtShape Uniqueness, RetAls)
 -> TypeBase ExtShape Uniqueness)
-> [(TypeBase ExtShape Uniqueness, RetAls)]
-> [TypeBase ExtShape Uniqueness]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase ExtShape Uniqueness, RetAls)
-> TypeBase ExtShape Uniqueness
forall a b. (a, b) -> a
fst [(TypeBase ExtShape Uniqueness, RetAls)]
rettype') ((Free [] (Param DeclType) -> [Param DeclType])
-> [Free [] (Param DeclType)] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Free [] (Param DeclType) -> [Param DeclType]
forall a. Free [] a -> [a]
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 -> () -> InternaliseM ()
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  where
    zeroExts :: [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts [TypeBase ExtShape u]
ts = [TypeBase ExtShape u]
-> [TypeBase ExtShape u] -> [TypeBase ExtShape u]
forall u.
[TypeBase ExtShape u]
-> [TypeBase ExtShape u] -> [TypeBase ExtShape u]
generaliseExtTypes [TypeBase ExtShape u]
ts [TypeBase ExtShape u]
ts

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 Exp VName)
_ (Info ResRetType
rettype) [TypeParamBase VName]
tparams [PatBase Info VName ParamType]
params Exp
_ Maybe DocComment
_ [AttrInfo VName]
attrs SrcLoc
loc) = ValBind
vb
  [TypeParamBase VName]
-> [PatBase Info VName ParamType]
-> ([FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM ())
-> InternaliseM ()
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 (([FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM ())
 -> InternaliseM ())
-> ([FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM ())
-> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \[FParam SOACS]
shapeparams [[Tree (FParam SOACS)]]
params' -> do
    let all_params :: [Free [] (Param DeclType)]
all_params = (Param DeclType -> Free [] (Param DeclType))
-> [Param DeclType] -> [Free [] (Param DeclType)]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> Free [] (Param DeclType)
forall a. a -> Free [] a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Param DeclType]
[FParam SOACS]
shapeparams [Free [] (Param DeclType)]
-> [Free [] (Param DeclType)] -> [Free [] (Param DeclType)]
forall a. [a] -> [a] -> [a]
++ [[Free [] (Param DeclType)]] -> [Free [] (Param DeclType)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Free [] (Param DeclType)]]
[[Tree (FParam SOACS)]]
params'
        ([[TypeBase ExtShape Uniqueness]]
entry_rettype, [[RetAls]]
retals) =
          [([TypeBase ExtShape Uniqueness], [RetAls])]
-> ([[TypeBase ExtShape Uniqueness]], [[RetAls]])
forall a b. [(a, b)] -> ([a], [b])
unzip ([([TypeBase ExtShape Uniqueness], [RetAls])]
 -> ([[TypeBase ExtShape Uniqueness]], [[RetAls]]))
-> [([TypeBase ExtShape Uniqueness], [RetAls])]
-> ([[TypeBase ExtShape Uniqueness]], [[RetAls]])
forall a b. (a -> b) -> a -> b
$ ([(TypeBase ExtShape Uniqueness, RetAls)]
 -> ([TypeBase ExtShape Uniqueness], [RetAls]))
-> [[(TypeBase ExtShape Uniqueness, RetAls)]]
-> [([TypeBase ExtShape Uniqueness], [RetAls])]
forall a b. (a -> b) -> [a] -> [b]
map [(TypeBase ExtShape Uniqueness, RetAls)]
-> ([TypeBase ExtShape Uniqueness], [RetAls])
forall a b. [(a, b)] -> ([a], [b])
unzip ([[(TypeBase ExtShape Uniqueness, RetAls)]]
 -> [([TypeBase ExtShape Uniqueness], [RetAls])])
-> [[(TypeBase ExtShape Uniqueness, RetAls)]]
-> [([TypeBase ExtShape Uniqueness], [RetAls])]
forall a b. (a -> b) -> a -> b
$ [Tree DeclType]
-> ResRetType -> [[(TypeBase ExtShape Uniqueness, RetAls)]]
internaliseEntryReturnType ((Free [] (Param DeclType) -> Tree DeclType)
-> [Free [] (Param DeclType)] -> [Tree DeclType]
forall a b. (a -> b) -> [a] -> [b]
map ((Param DeclType -> DeclType)
-> Free [] (Param DeclType) -> Tree DeclType
forall a b. (a -> b) -> Free [] a -> Free [] b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Param DeclType -> DeclType
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)
            ([EntryParam]
-> [[Param DeclType]] -> [(EntryParam, [Param DeclType])]
forall a b. [a] -> [b] -> [(a, b)]
zip [EntryParam]
e_params ([[Param DeclType]] -> [(EntryParam, [Param DeclType])])
-> [[Param DeclType]] -> [(EntryParam, [Param DeclType])]
forall a b. (a -> b) -> a -> b
$ ([Free [] (Param DeclType)] -> [Param DeclType])
-> [[Free [] (Param DeclType)]] -> [[Param DeclType]]
forall a b. (a -> b) -> [a] -> [b]
map ((Free [] (Param DeclType) -> [Param DeclType])
-> [Free [] (Param DeclType)] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Free [] (Param DeclType) -> [Param DeclType]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList) [[Free [] (Param DeclType)]]
[[Tree (FParam SOACS)]]
params')
            (EntryType
e_rettype, ([TypeBase ExtShape Uniqueness] -> [TypeBase Rank Uniqueness])
-> [[TypeBase ExtShape Uniqueness]] -> [[TypeBase Rank Uniqueness]]
forall a b. (a -> b) -> [a] -> [b]
map ((TypeBase ExtShape Uniqueness -> TypeBase Rank Uniqueness)
-> [TypeBase ExtShape Uniqueness] -> [TypeBase Rank Uniqueness]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase ExtShape Uniqueness -> TypeBase Rank Uniqueness
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase Rank u
I.rankShaped) [[TypeBase ExtShape Uniqueness]]
entry_rettype)
        args :: [SubExp]
args = (Param DeclType -> SubExp) -> [Param DeclType] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
I.Var (VName -> SubExp)
-> (Param DeclType -> VName) -> Param DeclType -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName) ([Param DeclType] -> [SubExp]) -> [Param DeclType] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ ([Free [] (Param DeclType)] -> [Param DeclType])
-> [[Free [] (Param DeclType)]] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((Free [] (Param DeclType) -> [Param DeclType])
-> [Free [] (Param DeclType)] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Free [] (Param DeclType) -> [Param DeclType]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList) [[Free [] (Param DeclType)]]
[[Tree (FParam SOACS)]]
params'

    OpaqueTypes -> InternaliseM ()
addOpaques OpaqueTypes
opaques

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

    Attrs
attrs' <- [AttrInfo VName] -> InternaliseM Attrs
internaliseAttrs [AttrInfo VName]
attrs
    let num_ctx :: Int
num_ctx = [(TypeBase ExtShape Uniqueness, RetAls)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(TypeBase ExtShape Uniqueness, RetAls)]
ctx_ts
    FunDef SOACS -> InternaliseM ()
addFunDef (FunDef SOACS -> InternaliseM ())
-> FunDef SOACS -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
      Maybe EntryPoint
-> Attrs
-> Name
-> [(RetType SOACS, RetAls)]
-> [FParam SOACS]
-> Body SOACS
-> FunDef SOACS
forall rep.
Maybe EntryPoint
-> Attrs
-> Name
-> [(RetType rep, RetAls)]
-> [FParam rep]
-> Body rep
-> FunDef rep
I.FunDef
        (EntryPoint -> Maybe EntryPoint
forall a. a -> Maybe a
Just EntryPoint
entry')
        Attrs
attrs'
        (Name
"entry_" Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> VName -> Name
baseName VName
ofname)
        ( [(TypeBase ExtShape Uniqueness, RetAls)]
ctx_ts
            [(TypeBase ExtShape Uniqueness, RetAls)]
-> [(TypeBase ExtShape Uniqueness, RetAls)]
-> [(TypeBase ExtShape Uniqueness, RetAls)]
forall a. [a] -> [a] -> [a]
++ [TypeBase ExtShape Uniqueness]
-> [RetAls] -> [(TypeBase ExtShape Uniqueness, RetAls)]
forall a b. [a] -> [b] -> [(a, b)]
zip
              ([TypeBase ExtShape Uniqueness] -> [TypeBase ExtShape Uniqueness]
forall {u}. [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts ([[TypeBase ExtShape Uniqueness]] -> [TypeBase ExtShape Uniqueness]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[TypeBase ExtShape Uniqueness]]
entry_rettype))
              ((RetAls -> RetAls) -> [RetAls] -> [RetAls]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> RetAls -> RetAls
shiftRetAls Int
num_ctx) ([RetAls] -> [RetAls]) -> [RetAls] -> [RetAls]
forall a b. (a -> b) -> a -> b
$ [[RetAls]] -> [RetAls]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[RetAls]]
retals)
        )
        ([Param DeclType]
[FParam SOACS]
shapeparams [Param DeclType] -> [Param DeclType] -> [Param DeclType]
forall a. [a] -> [a] -> [a]
++ ([Free [] (Param DeclType)] -> [Param DeclType])
-> [[Free [] (Param DeclType)]] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((Free [] (Param DeclType) -> [Param DeclType])
-> [Free [] (Param DeclType)] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Free [] (Param DeclType) -> [Param DeclType]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList) [[Free [] (Param DeclType)]]
[[Tree (FParam SOACS)]]
params')
        Body SOACS
entry_body
  where
    zeroExts :: [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts [TypeBase ExtShape u]
ts = [TypeBase ExtShape u]
-> [TypeBase ExtShape u] -> [TypeBase ExtShape u]
forall u.
[TypeBase ExtShape u]
-> [TypeBase ExtShape u] -> [TypeBase ExtShape u]
generaliseExtTypes [TypeBase ExtShape u]
ts [TypeBase ExtShape u]
ts

internaliseBody :: String -> E.Exp -> InternaliseM (Body SOACS)
internaliseBody :: [Char] -> Exp -> InternaliseM (Body SOACS)
internaliseBody [Char]
desc Exp
e =
  InternaliseM Result -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
m Result -> m (Body (Rep m))
buildBody_ (InternaliseM Result -> InternaliseM (Body (Rep InternaliseM)))
-> InternaliseM Result -> InternaliseM (Body (Rep InternaliseM))
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Result
subExpsRes ([SubExp] -> Result)
-> InternaliseM [SubExp] -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp ([Char]
desc [Char] -> [Char] -> [Char]
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) <- InternaliseM (Result, a)
-> InternaliseM ((Result, a), Stms (Rep InternaliseM))
forall a.
InternaliseM a -> InternaliseM (a, Stms (Rep InternaliseM))
forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms InternaliseM (Result, a)
m
  (,a
a) (Body SOACS -> (Body SOACS, a))
-> InternaliseM (Body SOACS) -> InternaliseM (Body SOACS, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Stms (Rep InternaliseM)
-> Result -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
Stms (Rep m) -> Result -> m (Body (Rep m))
mkBodyM Stms (Rep InternaliseM)
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 <- Exp SOACS -> InternaliseM [ExtType]
forall rep (m :: * -> *).
(HasScope rep m, TypedOp (OpC rep)) =>
Exp rep -> m [ExtType]
expExtType Exp SOACS
e
  [VName]
names <- Int -> InternaliseM VName -> InternaliseM [VName]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM ([ExtType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ExtType]
e_t) (InternaliseM VName -> InternaliseM [VName])
-> InternaliseM VName -> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$ [Char] -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
name
  [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName]
names Exp (Rep InternaliseM)
Exp SOACS
e
  let ctx :: Set Int
ctx = [ExtType] -> Set Int
forall u. [TypeBase ExtShape u] -> Set Int
shapeContext [ExtType]
e_t
  [VName] -> InternaliseM [VName]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([VName] -> InternaliseM [VName])
-> [VName] -> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$ ((VName, Int) -> VName) -> [(VName, Int)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, Int) -> VName
forall a b. (a, b) -> a
fst ([(VName, Int)] -> [VName]) -> [(VName, Int)] -> [VName]
forall a b. (a -> b) -> a -> b
$ ((VName, Int) -> Bool) -> [(VName, Int)] -> [(VName, Int)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Int -> Set Int -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set Int
ctx) (Int -> Bool) -> ((VName, Int) -> Int) -> (VName, Int) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, Int) -> Int
forall a b. (a, b) -> b
snd) ([(VName, Int)] -> [(VName, Int)])
-> [(VName, Int)] -> [(VName, Int)]
forall a b. (a -> b) -> a -> b
$ [VName] -> [Int] -> [(VName, Int)]
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)) = [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp
se]
letValExp' [Char]
name Exp SOACS
ses = (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var ([VName] -> [SubExp])
-> InternaliseM [VName] -> InternaliseM [SubExp]
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
    [] -> [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [] -- Will this happen?
    VName
v : [VName]
_ -> TypeBase Shape NoUniqueness -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
I.arrayDims (TypeBase Shape NoUniqueness -> [SubExp])
-> InternaliseM (TypeBase Shape NoUniqueness)
-> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM (TypeBase Shape NoUniqueness)
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 <- VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
v
        Exp SOACS -> InternaliseM (Exp SOACS)
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp SOACS -> InternaliseM (Exp SOACS))
-> Exp SOACS -> InternaliseM (Exp SOACS)
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp
I.Index VName
v (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoUniqueness -> [DimIndex SubExp] -> Slice SubExp
fullSlice TypeBase Shape NoUniqueness
v_t [DimIndex SubExp]
idxs'
  Certs -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a. Certs -> InternaliseM a -> InternaliseM a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ (VName -> InternaliseM SubExp) -> [VName] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc (Exp SOACS -> InternaliseM SubExp)
-> (VName -> InternaliseM (Exp SOACS))
-> VName
-> InternaliseM SubExp
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" (Exp -> InternaliseM SubExp) -> Exp -> InternaliseM SubExp
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' <-
    (Exp -> InternaliseM SubExp)
-> Maybe Exp -> InternaliseM (Maybe SubExp)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe 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
_)) -> IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntZ IntType
Int64
        StructType
_ -> IntType -> SubExp -> InternaliseM SubExp
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 <- (SubExp -> InternaliseM SubExp)
-> Maybe SubExp -> InternaliseM (Maybe SubExp)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse SubExp -> InternaliseM SubExp
conv Maybe SubExp
maybe_second'
  let errmsg :: ErrorMsg SubExp
errmsg =
        [ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg ([ErrorMsgPart SubExp] -> ErrorMsg SubExp)
-> [ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a b. (a -> b) -> a -> b
$
          [ErrorMsgPart SubExp
"Range "]
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
start'_i64]
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ ( case Maybe SubExp
maybe_second'_i64 of
                   Maybe SubExp
Nothing -> []
                   Just SubExp
second_i64 -> [ErrorMsgPart SubExp
"..", PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
second_i64]
               )
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ ( case Inclusiveness Exp
end of
                   DownToExclusive {} -> [ErrorMsgPart SubExp
"..>"]
                   ToInclusive {} -> [ErrorMsgPart SubExp
"..."]
                   UpToExclusive {} -> [ErrorMsgPart SubExp
"..<"]
               )
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
end'_i64, ErrorMsgPart SubExp
" is invalid."]

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

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

  SubExp
bounds_invalid_downwards <-
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"bounds_invalid_downwards" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
        CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp CmpOp
lt_op SubExp
start' SubExp
end'
  SubExp
bounds_invalid_upwards <-
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"bounds_invalid_upwards" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"step_wrong_dir" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
one
      SubExp
distance <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"distance" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
distance
      (SubExp, SubExp, SubExp) -> InternaliseM (SubExp, SubExp, SubExp)
forall a. a -> InternaliseM a
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 <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"step_wrong_dir" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
negone
      SubExp
distance <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"distance" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
distance
      (SubExp, SubExp, SubExp) -> InternaliseM (SubExp, SubExp, SubExp)
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
distance_i64, SubExp
step_wrong_dir, SubExp
bounds_invalid_upwards)
    ToInclusive {} -> do
      SubExp
downwards <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"downwards" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
negone
      SubExp
distance_downwards_exclusive <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"distance_downwards_exclusive" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"distance_upwards_exclusive" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"bounds_invalid"
          (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
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
            (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
downwards)
            ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
bounds_invalid_downwards])
            ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
bounds_invalid_upwards])
      SubExp
distance_exclusive <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"distance_exclusive"
          (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
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
            (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
downwards)
            ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
distance_downwards_exclusive])
            ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
distance_upwards_exclusive])
      SubExp
distance_exclusive_i64 <- IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
distance_exclusive
      SubExp
distance <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"distance" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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)
      (SubExp, SubExp, SubExp) -> InternaliseM (SubExp, SubExp, SubExp)
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
distance, Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
False, SubExp
bounds_invalid)

  SubExp
step_invalid <-
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"step_invalid" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <-
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"range_invalid" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
        BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogOr SubExp
step_invalid SubExp
bounds_invalid
  SubExp
valid <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"valid" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (PrimType -> UnOp
I.Neg PrimType
I.Bool) 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 <- IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
step
  SubExp
pos_step <-
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"pos_step" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <-
    Certs -> InternaliseM SubExp -> InternaliseM SubExp
forall a. Certs -> InternaliseM a -> InternaliseM a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs (InternaliseM SubExp -> InternaliseM SubExp)
-> InternaliseM SubExp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"num_elems" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
        BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc (BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ SubExp -> SubExp -> SubExp -> IntType -> BasicOp
I.Iota SubExp
num_elems SubExp
start' SubExp
step IntType
it)
  [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
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 = (VName -> (VName, Subst StructRetType))
-> [VName] -> [(VName, Subst StructRetType)]
forall a b. (a -> b) -> [a] -> [b]
map (,Exp -> Subst StructRetType
forall t. Exp -> Subst t
E.ExpSubst (Integer -> SrcLoc -> Exp
E.sizeFromInteger Integer
0 SrcLoc
forall a. Monoid a => a
mempty)) [VName]
ext
          et' :: StructType
et' = TypeSubs -> StructType -> StructType
forall a. Substitutable a => TypeSubs -> a -> a
E.applySubst (VName
-> [(VName, Subst StructRetType)] -> Maybe (Subst StructRetType)
forall a b. Eq a => a -> [(a, b)] -> Maybe b
`lookup` [(VName, Subst StructRetType)]
subst) StructType
et
      [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc (Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn. f StructType -> SrcLoc -> ExpBase f vn
E.Hole (StructType -> Info StructType
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 ([Char] -> Name) -> [Char] -> Name
forall a b. (a -> b) -> a -> b
$ Name -> [Char]
forall a. Pretty a => a -> [Char]
prettyString (Name -> [Char]) -> Name -> [Char]
forall a b. (a -> b) -> a -> b
$ VName -> Name
baseName (VName -> Name) -> VName -> Name
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname
          loc :: SrcLoc
loc = AppExp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf AppExp
e
          arg_desc :: [Char]
arg_desc = Name -> [Char]
nameToString Name
fname [Char] -> [Char] -> [Char]
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 (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag,
            VName -> [Char]
baseString (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) [Char] -> [Char] -> Bool
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 (Exp -> InternaliseM [SubExp]) -> Exp -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
                AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
E.AppExp
                  (Exp -> Exp -> Exp -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
E.If Exp
x Exp
y (PrimValue -> SrcLoc -> Exp
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
E.Literal (Bool -> PrimValue
E.BoolValue Bool
False) SrcLoc
forall a. Monoid a => a
mempty) SrcLoc
forall a. Monoid a => a
mempty)
                  (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes (ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
E.Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
-> ScalarTypeBase Exp NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp NoUniqueness
forall dim u. PrimType -> ScalarTypeBase dim u
E.Prim PrimType
E.Bool) [])
          | VName -> Int
baseTag (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag,
            VName -> [Char]
baseString (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) [Char] -> [Char] -> Bool
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 (Exp -> InternaliseM [SubExp]) -> Exp -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
                AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
E.AppExp
                  (Exp -> Exp -> Exp -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
E.If Exp
x (PrimValue -> SrcLoc -> Exp
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
E.Literal (Bool -> PrimValue
E.BoolValue Bool
True) SrcLoc
forall a. Monoid a => a
mempty) Exp
y SrcLoc
forall a. Monoid a => a
mempty)
                  (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes (ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
E.Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
-> ScalarTypeBase Exp NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp NoUniqueness
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
_) =
                    (StructType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
E.toStruct (Exp -> StructType
E.typeOf Exp
arg),) ([SubExp] -> (StructType, [SubExp]))
-> InternaliseM [SubExp] -> InternaliseM (StructType, [SubExp])
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 ([(StructType, [SubExp])] -> InternaliseM [SubExp])
-> InternaliseM [(StructType, [SubExp])] -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ((Exp, Maybe VName) -> InternaliseM (StructType, [SubExp]))
-> [(Exp, Maybe VName)] -> InternaliseM [(StructType, [SubExp])]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Exp, Maybe VName) -> InternaliseM (StructType, [SubExp])
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 (((Exp, Maybe VName) -> Exp) -> [(Exp, Maybe VName)] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (Exp, Maybe VName) -> Exp
forall a b. (a, b) -> a
fst [(Exp, Maybe VName)]
args) SrcLoc
loc ->
              [Char] -> InternaliseM [SubExp]
internalise [Char]
desc
          | VName -> Int
baseTag (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag,
            Just (PrimType
rettype, [PrimType]
_) <- Name
-> Map Name (PrimType, [PrimType]) -> Maybe (PrimType, [PrimType])
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
fname Map Name (PrimType, [PrimType])
I.builtInFunctions -> do
              let tag :: [a] -> [(a, Diet)]
tag [a]
ses = [(a
se, Diet
I.Observe) | a
se <- [a]
ses]
              [[SubExp]]
args' <- [[SubExp]] -> [[SubExp]]
forall a. [a] -> [a]
reverse ([[SubExp]] -> [[SubExp]])
-> InternaliseM [[SubExp]] -> InternaliseM [[SubExp]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Exp, Maybe VName) -> InternaliseM [SubExp])
-> [(Exp, Maybe VName)] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Char] -> (Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg [Char]
arg_desc) ([(Exp, Maybe VName)] -> [(Exp, Maybe VName)]
forall a. [a] -> [a]
reverse [(Exp, Maybe VName)]
args)
              let args'' :: [(SubExp, Diet)]
args'' = ([SubExp] -> [(SubExp, Diet)]) -> [[SubExp]] -> [(SubExp, Diet)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap [SubExp] -> [(SubExp, Diet)]
forall {a}. [a] -> [(a, Diet)]
tag [[SubExp]]
args'
              [Char] -> Exp SOACS -> InternaliseM [SubExp]
letValExp' [Char]
desc (Exp SOACS -> InternaliseM [SubExp])
-> Exp SOACS -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Name
-> [(SubExp, Diet)]
-> [(RetType SOACS, RetAls)]
-> (Safety, SrcLoc, [SrcLoc])
-> Exp SOACS
forall rep.
Name
-> [(SubExp, Diet)]
-> [(RetType rep, RetAls)]
-> (Safety, SrcLoc, [SrcLoc])
-> Exp rep
I.Apply Name
fname [(SubExp, Diet)]
args'' [(PrimType -> TypeBase ExtShape Uniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
rettype, RetAls
forall a. Monoid a => a
mempty)] (Safety
Safe, SrcLoc
loc, [])
          | Bool
otherwise -> do
              [SubExp]
args' <- [[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[SubExp]] -> [SubExp])
-> ([[SubExp]] -> [[SubExp]]) -> [[SubExp]] -> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[SubExp]] -> [[SubExp]]
forall a. [a] -> [a]
reverse ([[SubExp]] -> [SubExp])
-> InternaliseM [[SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Exp, Maybe VName) -> InternaliseM [SubExp])
-> [(Exp, Maybe VName)] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Char] -> (Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg [Char]
arg_desc) ([(Exp, Maybe VName)] -> [(Exp, Maybe VName)]
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
_) =
  [Char]
-> [SizeBinder VName]
-> PatBase Info VName StructType
-> Exp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
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 (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
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 Exp VName), Info ResRetType, Exp)
_ Exp
_ SrcLoc
_) =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"Unexpected LetFun " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ VName -> [Char]
forall a. Pretty a => a -> [Char]
prettyString VName
ofname
internaliseAppExp [Char]
desc AppRes
_ (E.Loop [VName]
sparams PatBase Info VName ParamType
mergepat LoopInitBase Info VName
loopinit LoopFormBase Info VName
form Exp
loopbody SrcLoc
loc) = do
  [SubExp]
ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"loop_init" (Exp -> InternaliseM [SubExp]) -> Exp -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ LoopInitBase Info VName -> Exp
loopInitExp LoopInitBase Info VName
loopinit
  ((Body SOACS
loopbody', (LoopForm
form', [Param DeclType]
shapepat, [Param DeclType]
mergepat', [SubExp]
mergeinit')), Stms SOACS
initstms) <-
    InternaliseM
  (Body SOACS,
   (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     ((Body SOACS,
       (LoopForm, [Param DeclType], [Param DeclType], [SubExp])),
      Stms (Rep InternaliseM))
forall a.
InternaliseM a -> InternaliseM (a, Stms (Rep InternaliseM))
forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms (InternaliseM
   (Body SOACS,
    (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
 -> InternaliseM
      ((Body SOACS,
        (LoopForm, [Param DeclType], [Param DeclType], [SubExp])),
       Stms (Rep InternaliseM)))
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     ((Body SOACS,
       (LoopForm, [Param DeclType], [Param DeclType], [SubExp])),
      Stms (Rep InternaliseM))
forall a b. (a -> b) -> a -> b
$ [SubExp]
-> LoopFormBase Info VName
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
handleForm [SubExp]
ses LoopFormBase Info VName
form

  Stms (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
addStms Stms (Rep InternaliseM)
Stms SOACS
initstms
  [TypeBase Shape NoUniqueness]
mergeinit_ts' <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
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 ((Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
shapepat) [Param DeclType]
[FParam SOACS]
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 [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
mergeinit'
  [SubExp]
args' <-
    ErrorMsg SubExp
-> SrcLoc
-> [VName]
-> [TypeBase Shape NoUniqueness]
-> [SubExp]
-> InternaliseM [SubExp]
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
      ((Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
shapepat)
      ((Param DeclType -> TypeBase Shape NoUniqueness)
-> [Param DeclType] -> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> TypeBase Shape NoUniqueness
forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
paramType ([Param DeclType] -> [TypeBase Shape NoUniqueness])
-> [Param DeclType] -> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> a -> b
$ [Param DeclType]
shapepat [Param DeclType] -> [Param DeclType] -> [Param DeclType]
forall a. [a] -> [a] -> [a]
++ [Param DeclType]
mergepat')
      [SubExp]
args

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

  -- As above, ensure that the result has the right shape.
  let merge :: [(Param DeclType, SubExp)]
merge = [Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([Param DeclType]
shapepat [Param DeclType] -> [Param DeclType] -> [Param DeclType]
forall a. [a] -> [a] -> [a]
++ [Param DeclType]
mergepat') [SubExp]
args'
      merge_ts :: [TypeBase Shape NoUniqueness]
merge_ts = ((Param DeclType, SubExp) -> TypeBase Shape NoUniqueness)
-> [(Param DeclType, SubExp)] -> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map (Param DeclType -> TypeBase Shape NoUniqueness
forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
I.paramType (Param DeclType -> TypeBase Shape NoUniqueness)
-> ((Param DeclType, SubExp) -> Param DeclType)
-> (Param DeclType, SubExp)
-> TypeBase Shape NoUniqueness
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param DeclType, SubExp) -> Param DeclType
forall a b. (a, b) -> a
fst) [(Param DeclType, SubExp)]
merge
  Body SOACS
loopbody'' <-
    Scope SOACS
-> InternaliseM (Body SOACS) -> InternaliseM (Body SOACS)
forall a. Scope SOACS -> InternaliseM a -> InternaliseM a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope ([Param DeclType] -> Scope SOACS
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams (((Param DeclType, SubExp) -> Param DeclType)
-> [(Param DeclType, SubExp)] -> [Param DeclType]
forall a b. (a -> b) -> [a] -> [b]
map (Param DeclType, SubExp) -> Param DeclType
forall a b. (a, b) -> a
fst [(Param DeclType, SubExp)]
merge) Scope SOACS -> Scope SOACS -> Scope SOACS
forall a. Semigroup a => a -> a -> a
<> LoopForm -> Scope SOACS
forall rep. LoopForm -> Scope rep
scopeOfLoopForm LoopForm
form') (InternaliseM (Body SOACS) -> InternaliseM (Body SOACS))
-> (InternaliseM Result -> InternaliseM (Body SOACS))
-> InternaliseM Result
-> InternaliseM (Body SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InternaliseM Result -> InternaliseM (Body (Rep InternaliseM))
InternaliseM Result -> InternaliseM (Body SOACS)
forall (m :: * -> *).
MonadBuilder m =>
m Result -> m (Body (Rep m))
buildBody_ (InternaliseM Result -> InternaliseM (Body SOACS))
-> InternaliseM Result -> InternaliseM (Body SOACS)
forall a b. (a -> b) -> a -> b
$
      ([SubExp] -> Result)
-> InternaliseM [SubExp] -> InternaliseM Result
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [SubExp] -> Result
subExpsRes
        (InternaliseM [SubExp] -> InternaliseM Result)
-> (Result -> InternaliseM [SubExp])
-> Result
-> InternaliseM Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ErrorMsg SubExp
-> SrcLoc
-> [VName]
-> [TypeBase Shape NoUniqueness]
-> [SubExp]
-> InternaliseM [SubExp]
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
          (((Param DeclType, SubExp) -> VName)
-> [(Param DeclType, SubExp)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName (Param DeclType -> VName)
-> ((Param DeclType, SubExp) -> Param DeclType)
-> (Param DeclType, SubExp)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param DeclType, SubExp) -> Param DeclType
forall a b. (a, b) -> a
fst) [(Param DeclType, SubExp)]
merge)
          [TypeBase Shape NoUniqueness]
merge_ts
        ([SubExp] -> InternaliseM [SubExp])
-> (Result -> [SubExp]) -> Result -> InternaliseM [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SubExpRes -> SubExp) -> Result -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp
        (Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Body (Rep InternaliseM) -> InternaliseM Result
forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind Body (Rep InternaliseM)
Body SOACS
loopbody'

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

    -- Attributes that apply to loops.
    loopAttrs :: Attrs
loopAttrs = Attr -> Attrs
oneAttr Attr
"unroll"
    -- Remove those attributes from the attribute set that apply to
    -- the loop itself.
    noLoopAttrs :: InternaliseEnv -> InternaliseEnv
noLoopAttrs InternaliseEnv
env = InternaliseEnv
env {envAttrs = envAttrs env `withoutAttrs` loopAttrs}

    loopBody :: InternaliseM [SubExp]
loopBody = (InternaliseEnv -> InternaliseEnv)
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a.
(InternaliseEnv -> InternaliseEnv)
-> InternaliseM a -> InternaliseM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local InternaliseEnv -> InternaliseEnv
noLoopAttrs (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"loopres" Exp
loopbody

    forLoop :: [Param DeclType]
-> [Param DeclType]
-> [SubExp]
-> VName
-> [(Param (TypeBase Shape NoUniqueness), VName)]
-> LoopForm
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forLoop [Param DeclType]
mergepat' [Param DeclType]
shapepat [SubExp]
mergeinit VName
i [(Param (TypeBase Shape NoUniqueness), VName)]
loopvars LoopForm
form' =
      InternaliseM
  (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a. InternaliseM (Result, a) -> InternaliseM (Body SOACS, a)
bodyFromStms (InternaliseM
   (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
 -> InternaliseM
      (Body SOACS,
       (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> (InternaliseM
      (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
    -> InternaliseM
         (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Scope SOACS
-> InternaliseM
     (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a. Scope SOACS -> InternaliseM a -> InternaliseM a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (LoopForm -> Scope SOACS
forall rep. LoopForm -> Scope rep
scopeOfLoopForm LoopForm
form') (InternaliseM
   (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
 -> InternaliseM
      (Body SOACS,
       (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ do
        [(Param (TypeBase Shape NoUniqueness), VName)]
-> ((Param (TypeBase Shape NoUniqueness), VName)
    -> InternaliseM ())
-> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(Param (TypeBase Shape NoUniqueness), VName)]
loopvars (((Param (TypeBase Shape NoUniqueness), VName) -> InternaliseM ())
 -> InternaliseM ())
-> ((Param (TypeBase Shape NoUniqueness), VName)
    -> InternaliseM ())
-> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \(Param (TypeBase Shape NoUniqueness)
p, VName
arr) ->
          [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param (TypeBase Shape NoUniqueness) -> VName
forall dec. Param dec -> VName
I.paramName Param (TypeBase Shape NoUniqueness)
p] (Exp SOACS -> InternaliseM ())
-> InternaliseM (Exp SOACS) -> InternaliseM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName
-> [InternaliseM (Exp (Rep InternaliseM))]
-> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
arr [SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp (VName -> SubExp
I.Var VName
i)]
        [SubExp]
ses <- InternaliseM [SubExp]
loopBody
        [TypeBase Shape NoUniqueness]
sets <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
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 ((Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
shapepat) [Param DeclType]
[FParam SOACS]
mergepat' [TypeBase Shape NoUniqueness]
sets
        (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
          ( [SubExp] -> Result
subExpsRes ([SubExp] -> Result) -> [SubExp] -> Result
forall a b. (a -> b) -> a -> b
$ [SubExp]
shapeargs [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
ses,
            ( LoopForm
form',
              [Param DeclType]
shapepat,
              [Param DeclType]
mergepat',
              [SubExp]
mergeinit
            )
          )

    handleForm :: [SubExp]
-> LoopFormBase Info VName
-> InternaliseM
     (Body SOACS,
      (LoopForm, [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 <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arr'
      let w :: SubExp
w = Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [TypeBase Shape NoUniqueness]
arr_ts

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

      [TypeBase Shape NoUniqueness]
ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
mergeinit
      [TypeParamBase VName]
-> PatBase Info VName ParamType
-> [TypeBase Shape NoUniqueness]
-> ([FParam SOACS]
    -> [FParam SOACS]
    -> InternaliseM
         (Body SOACS,
          (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
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 (([FParam SOACS]
  -> [FParam SOACS]
  -> InternaliseM
       (Body SOACS,
        (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
 -> InternaliseM
      (Body SOACS,
       (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> ([FParam SOACS]
    -> [FParam SOACS]
    -> InternaliseM
         (Body SOACS,
          (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ \[FParam SOACS]
shapepat [FParam SOACS]
mergepat' ->
        [PatBase Info VName ParamType]
-> [TypeBase Shape NoUniqueness]
-> ([LParam SOACS]
    -> InternaliseM
         (Body SOACS,
          (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a.
[PatBase Info VName ParamType]
-> [TypeBase Shape NoUniqueness]
-> ([LParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLambdaParams [Diet -> StructType -> ParamType
forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
E.Observe (StructType -> ParamType)
-> PatBase Info VName StructType -> PatBase Info VName ParamType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PatBase Info VName StructType
x] ((TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
-> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall u. TypeBase Shape u -> TypeBase Shape u
rowType [TypeBase Shape NoUniqueness]
arr_ts) (([LParam SOACS]
  -> InternaliseM
       (Body SOACS,
        (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
 -> InternaliseM
      (Body SOACS,
       (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> ([LParam SOACS]
    -> InternaliseM
         (Body SOACS,
          (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ \[LParam SOACS]
x_params -> do
          let loopvars :: [(Param (TypeBase Shape NoUniqueness), VName)]
loopvars = [Param (TypeBase Shape NoUniqueness)]
-> [VName] -> [(Param (TypeBase Shape NoUniqueness), VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param (TypeBase Shape NoUniqueness)]
[LParam SOACS]
x_params [VName]
arr'
          [Param DeclType]
-> [Param DeclType]
-> [SubExp]
-> VName
-> [(Param (TypeBase Shape NoUniqueness), VName)]
-> LoopForm
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forLoop [Param DeclType]
[FParam SOACS]
mergepat' [Param DeclType]
[FParam SOACS]
shapepat [SubExp]
mergeinit VName
i [(Param (TypeBase Shape NoUniqueness), VName)]
loopvars (LoopForm
 -> InternaliseM
      (Body SOACS,
       (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> LoopForm
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ VName -> IntType -> SubExp -> LoopForm
I.ForLoop VName
i IntType
Int64 SubExp
w
    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 <- SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
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) -> IntType -> InternaliseM IntType
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure IntType
it
        TypeBase Shape NoUniqueness
_ -> [Char] -> InternaliseM IntType
forall a. HasCallStack => [Char] -> a
error [Char]
"internaliseExp Loop: invalid type"

      [TypeBase Shape NoUniqueness]
ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
mergeinit
      [TypeParamBase VName]
-> PatBase Info VName ParamType
-> [TypeBase Shape NoUniqueness]
-> ([FParam SOACS]
    -> [FParam SOACS]
    -> InternaliseM
         (Body SOACS,
          (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
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 (([FParam SOACS]
  -> [FParam SOACS]
  -> InternaliseM
       (Body SOACS,
        (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
 -> InternaliseM
      (Body SOACS,
       (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> ([FParam SOACS]
    -> [FParam SOACS]
    -> InternaliseM
         (Body SOACS,
          (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ \[FParam SOACS]
shapepat [FParam SOACS]
mergepat' ->
        [Param DeclType]
-> [Param DeclType]
-> [SubExp]
-> VName
-> [(Param (TypeBase Shape NoUniqueness), VName)]
-> LoopForm
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forLoop [Param DeclType]
[FParam SOACS]
mergepat' [Param DeclType]
[FParam SOACS]
shapepat [SubExp]
mergeinit (IdentBase Info VName StructType -> VName
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
E.identName IdentBase Info VName StructType
i) [] (LoopForm
 -> InternaliseM
      (Body SOACS,
       (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> LoopForm
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$
          VName -> IntType -> SubExp -> LoopForm
I.ForLoop (IdentBase Info VName StructType -> VName
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 <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
mergeinit
      [TypeParamBase VName]
-> PatBase Info VName ParamType
-> [TypeBase Shape NoUniqueness]
-> ([FParam SOACS]
    -> [FParam SOACS]
    -> InternaliseM
         (Body SOACS,
          (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
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 (([FParam SOACS]
  -> [FParam SOACS]
  -> InternaliseM
       (Body SOACS,
        (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
 -> InternaliseM
      (Body SOACS,
       (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> ([FParam SOACS]
    -> [FParam SOACS]
    -> InternaliseM
         (Body SOACS,
          (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ \[FParam SOACS]
shapepat [FParam SOACS]
mergepat' -> do
        [TypeBase Shape NoUniqueness]
mergeinit_ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
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 ((Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
[FParam SOACS]
shapepat) [FParam SOACS]
mergepat' [TypeBase Shape NoUniqueness]
mergeinit_ts

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

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

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

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

          (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Result, (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
            ( [SubExp] -> Result
subExpsRes [SubExp]
shapeargs Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ Result
loop_end_cond Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ [SubExp] -> Result
subExpsRes [SubExp]
ses,
              ( VName -> LoopForm
I.WhileLoop (VName -> LoopForm) -> VName -> LoopForm
forall a b. (a -> b) -> a -> b
$ Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName Param DeclType
loop_while,
                [Param DeclType]
[FParam SOACS]
shapepat,
                Param DeclType
loop_while Param DeclType -> [Param DeclType] -> [Param DeclType]
forall a. a -> [a] -> [a]
: [Param DeclType]
[FParam SOACS]
mergepat',
                SubExp
loop_initial_cond SubExp -> [SubExp] -> [SubExp]
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 = VName -> Info StructType -> SrcLoc -> PatBase Info VName StructType
forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
E.Id (IdentBase Info VName StructType -> VName
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
E.identName IdentBase Info VName StructType
name) (IdentBase Info VName StructType -> Info StructType
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 = IdentBase Info VName StructType -> Info StructType
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> f t
E.identType IdentBase Info VName StructType
src
      e :: Exp
e = Exp -> SliceBase Info VName -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> SliceBase f vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
E.Update (QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
E.Var (VName -> QualName VName
forall v. v -> QualName v
E.qualName (VName -> QualName VName) -> VName -> QualName VName
forall a b. (a -> b) -> a -> b
$ IdentBase Info VName StructType -> VName
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 (Exp -> InternaliseM [SubExp]) -> Exp -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
E.AppExp
      ([SizeBinder VName]
-> PatBase Info VName StructType -> Exp -> Exp -> SrcLoc -> 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)
      (AppRes -> Info AppRes
forall a. a -> Info a
Info (StructType -> [VName] -> AppRes
AppRes (Exp -> StructType
E.typeOf Exp
body) [VName]
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 [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"_scrutinee") Exp
e
  NonEmpty (Case (InternaliseM (Body SOACS)))
cs <- (CaseBase Info VName
 -> InternaliseM (Case (InternaliseM (Body SOACS))))
-> NonEmpty (CaseBase Info VName)
-> InternaliseM (NonEmpty (Case (InternaliseM (Body SOACS))))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> NonEmpty a -> m (NonEmpty b)
mapM ([SubExp]
-> CaseBase Info VName
-> InternaliseM (Case (InternaliseM (Body SOACS)))
onCase [SubExp]
ses) NonEmpty (CaseBase Info VName)
orig_cs
  case NonEmpty (Case (InternaliseM (Body SOACS)))
-> (Case (InternaliseM (Body SOACS)),
    Maybe (NonEmpty (Case (InternaliseM (Body SOACS)))))
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) ->
      (Result -> [SubExp])
-> InternaliseM Result -> InternaliseM [SubExp]
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((SubExpRes -> SubExp) -> Result -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp) (InternaliseM Result -> InternaliseM [SubExp])
-> InternaliseM Result -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Body (Rep InternaliseM) -> InternaliseM Result
Body SOACS -> InternaliseM Result
forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind (Body SOACS -> InternaliseM Result)
-> InternaliseM (Body SOACS) -> InternaliseM Result
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 (Exp SOACS -> InternaliseM [SubExp])
-> InternaliseM (Exp SOACS) -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [SubExp]
-> [Case (InternaliseM (Body (Rep InternaliseM)))]
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
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 (NonEmpty (Case (InternaliseM (Body SOACS)))
-> [Case (InternaliseM (Body SOACS))]
forall a. NonEmpty a -> [a]
NE.init NonEmpty (Case (InternaliseM (Body SOACS)))
cs) (Case (InternaliseM (Body (Rep InternaliseM)))
-> InternaliseM (Body (Rep InternaliseM))
forall body. Case body -> body
I.caseBody (Case (InternaliseM (Body (Rep InternaliseM)))
 -> InternaliseM (Body (Rep InternaliseM)))
-> Case (InternaliseM (Body (Rep InternaliseM)))
-> InternaliseM (Body (Rep InternaliseM))
forall a b. (a -> b) -> a -> b
$ NonEmpty (Case (InternaliseM (Body SOACS)))
-> Case (InternaliseM (Body SOACS))
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
      Case (InternaliseM (Body SOACS))
-> InternaliseM (Case (InternaliseM (Body SOACS)))
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Case (InternaliseM (Body SOACS))
 -> InternaliseM (Case (InternaliseM (Body SOACS))))
-> (InternaliseM (Body SOACS) -> Case (InternaliseM (Body SOACS)))
-> InternaliseM (Body SOACS)
-> InternaliseM (Case (InternaliseM (Body SOACS)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe PrimValue]
-> InternaliseM (Body SOACS) -> Case (InternaliseM (Body SOACS))
forall body. [Maybe PrimValue] -> body -> Case body
I.Case [Maybe PrimValue]
cmps (InternaliseM (Body SOACS)
 -> InternaliseM (Case (InternaliseM (Body SOACS))))
-> InternaliseM (Body SOACS)
-> InternaliseM (Case (InternaliseM (Body SOACS)))
forall a b. (a -> b) -> a -> b
$
        [SizeBinder VName]
-> PatBase Info VName StructType
-> [SubExp]
-> InternaliseM (Body SOACS)
-> InternaliseM (Body SOACS)
forall a.
[SizeBinder VName]
-> PatBase Info VName StructType
-> [SubExp]
-> InternaliseM a
-> InternaliseM a
internalisePat' [] PatBase Info VName StructType
p [SubExp]
pertinent (InternaliseM (Body SOACS) -> InternaliseM (Body SOACS))
-> InternaliseM (Body SOACS) -> InternaliseM (Body SOACS)
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
    (Exp SOACS -> InternaliseM [SubExp])
-> InternaliseM (Exp SOACS) -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
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
      (BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS)
-> (SubExp -> BasicOp) -> SubExp -> Exp SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
SubExp (SubExp -> Exp SOACS)
-> InternaliseM SubExp -> InternaliseM (Exp SOACS)
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 [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"_t") Exp
te)
      ([Char] -> Exp -> InternaliseM (Body SOACS)
internaliseBody ([Char]
desc [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"_f") Exp
fe)
internaliseAppExp [Char]
_ AppRes
_ e :: AppExp
e@E.BinOp {} =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseAppExp: Unexpected BinOp " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ AppExp -> [Char]
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 = Doc Any -> Text
forall a. Doc a -> Text
docText (Doc Any -> Text) -> Doc Any -> Text
forall a b. (a -> b) -> a -> b
$ Doc Any
"Reached hole of type: " Doc Any -> Doc Any -> Doc Any
forall a. Semigroup a => a -> a -> a
<> Doc Any -> Doc Any
forall ann. Doc ann -> Doc ann
align (StructType -> Doc Any
forall a ann. Pretty a => a -> Doc ann
forall ann. StructType -> Doc ann
pretty StructType
t)
      ts :: [TypeBase ExtShape Uniqueness]
ts = (Tree (TypeBase ExtShape Uniqueness)
 -> [TypeBase ExtShape Uniqueness])
-> [Tree (TypeBase ExtShape Uniqueness)]
-> [TypeBase ExtShape Uniqueness]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Tree (TypeBase ExtShape Uniqueness)
-> [TypeBase ExtShape Uniqueness]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList ([Tree (TypeBase ExtShape Uniqueness)]
 -> [TypeBase ExtShape Uniqueness])
-> [Tree (TypeBase ExtShape Uniqueness)]
-> [TypeBase ExtShape Uniqueness]
forall a b. (a -> b) -> a -> b
$ StructType -> [Tree (TypeBase ExtShape Uniqueness)]
internaliseType (StructType -> StructType
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" (Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
False) ([ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [Text -> ErrorMsgPart SubExp
forall a. Text -> ErrorMsgPart a
ErrorString Text
msg]) SrcLoc
loc
  case (TypeBase ExtShape Uniqueness -> Maybe DeclType)
-> [TypeBase ExtShape Uniqueness] -> Maybe [DeclType]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM TypeBase ExtShape Uniqueness -> Maybe DeclType
forall u. TypeBase ExtShape u -> Maybe (TypeBase Shape u)
hasStaticShape [TypeBase ExtShape Uniqueness]
ts of
    Maybe [DeclType]
Nothing ->
      [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"Hole at " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr SrcLoc
loc [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" has existential type:\n" [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [TypeBase ExtShape Uniqueness] -> [Char]
forall a. Show a => a -> [Char]
show [TypeBase ExtShape Uniqueness]
ts
    Just [DeclType]
ts' ->
      -- Make sure we always generate a binding, even for primitives.
      Certs -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a. Certs -> InternaliseM a -> InternaliseM a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ (DeclType -> InternaliseM SubExp)
-> [DeclType] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((VName -> SubExp) -> InternaliseM VName -> InternaliseM SubExp
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap VName -> SubExp
I.Var (InternaliseM VName -> InternaliseM SubExp)
-> (Exp SOACS -> InternaliseM VName)
-> Exp SOACS
-> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp (Rep InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
desc (Exp SOACS -> InternaliseM SubExp)
-> (DeclType -> InternaliseM (Exp SOACS))
-> DeclType
-> InternaliseM SubExp
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< TypeBase Shape NoUniqueness
-> InternaliseM (Exp (Rep InternaliseM))
TypeBase Shape NoUniqueness -> InternaliseM (Exp SOACS)
forall (m :: * -> *).
MonadBuilder m =>
TypeBase Shape NoUniqueness -> m (Exp (Rep m))
eBlank (TypeBase Shape NoUniqueness -> InternaliseM (Exp SOACS))
-> (DeclType -> TypeBase Shape NoUniqueness)
-> DeclType
-> InternaliseM (Exp SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DeclType -> TypeBase Shape NoUniqueness
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
_) =
  (SubExp -> [SubExp])
-> InternaliseM SubExp -> InternaliseM [SubExp]
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> [SubExp]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InternaliseM SubExp -> InternaliseM [SubExp])
-> (Exp SOACS -> InternaliseM SubExp)
-> Exp SOACS
-> InternaliseM [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc (Exp SOACS -> InternaliseM [SubExp])
-> Exp SOACS -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$
      [SubExp] -> TypeBase Shape NoUniqueness -> BasicOp
I.ArrayLit ((Word8 -> SubExp) -> [Word8] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map Word8 -> SubExp
forall v. IsValue v => v -> SubExp
constant [Word8]
vs) (TypeBase Shape NoUniqueness -> BasicOp)
-> TypeBase Shape NoUniqueness -> BasicOp
forall a b. (a -> b) -> a -> b
$
        PrimType -> TypeBase Shape NoUniqueness
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 -> [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
substs
    Maybe [SubExp]
Nothing -> [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
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
  [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
ses
internaliseExp [Char]
_ (E.TupLit [] SrcLoc
_) =
  [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
forall v. IsValue v => v -> SubExp
constant PrimValue
UnitValue]
internaliseExp [Char]
_ (E.RecordLit [] SrcLoc
_) =
  [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
forall v. IsValue v => v -> SubExp
constant PrimValue
UnitValue]
internaliseExp [Char]
desc (E.TupLit [Exp]
es SrcLoc
_) = [[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[SubExp]] -> [SubExp])
-> InternaliseM [[SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> InternaliseM [SubExp]) -> [Exp] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc) [Exp]
es
internaliseExp [Char]
desc (E.RecordLit [FieldBase Info VName]
orig_fields SrcLoc
_) =
  ((Name, [SubExp]) -> [SubExp]) -> [(Name, [SubExp])] -> [SubExp]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Name, [SubExp]) -> [SubExp]
forall a b. (a, b) -> b
snd ([(Name, [SubExp])] -> [SubExp])
-> ([Map Name [SubExp]] -> [(Name, [SubExp])])
-> [Map Name [SubExp]]
-> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name [SubExp] -> [(Name, [SubExp])]
forall a. Map Name a -> [(Name, a)]
sortFields (Map Name [SubExp] -> [(Name, [SubExp])])
-> ([Map Name [SubExp]] -> Map Name [SubExp])
-> [Map Name [SubExp]]
-> [(Name, [SubExp])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Map Name [SubExp]] -> Map Name [SubExp]
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
M.unions ([Map Name [SubExp]] -> [SubExp])
-> InternaliseM [Map Name [SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (FieldBase Info VName -> InternaliseM (Map Name [SubExp]))
-> [FieldBase Info VName] -> InternaliseM [Map Name [SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [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 (L Loc
_ Name
name) Exp
e SrcLoc
_) =
      Name -> [SubExp] -> Map Name [SubExp]
forall k a. k -> a -> Map k a
M.singleton Name
name ([SubExp] -> Map Name [SubExp])
-> InternaliseM [SubExp] -> InternaliseM (Map Name [SubExp])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
    internaliseField (E.RecordFieldImplicit (L Loc
_ VName
name) Info StructType
t SrcLoc
loc) =
      FieldBase Info VName -> InternaliseM (Map Name [SubExp])
internaliseField (FieldBase Info VName -> InternaliseM (Map Name [SubExp]))
-> FieldBase Info VName -> InternaliseM (Map Name [SubExp])
forall a b. (a -> b) -> a -> b
$
        L Name -> Exp -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
L Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
E.RecordFieldExplicit
          (Loc -> Name -> L Name
forall a. Loc -> a -> L a
L Loc
forall a. IsLocation a => a
noLoc (VName -> Name
baseName VName
name))
          (QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
E.Var (VName -> QualName VName
forall v. v -> QualName v
E.qualName VName
name) Info StructType
t SrcLoc
loc)
          SrcLoc
loc
internaliseExp [Char]
desc (E.ArrayVal [PrimValue]
vs PrimType
t SrcLoc
_) =
  (SubExp -> [SubExp])
-> InternaliseM SubExp -> InternaliseM [SubExp]
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> [SubExp]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InternaliseM SubExp -> InternaliseM [SubExp])
-> (BasicOp -> InternaliseM SubExp)
-> BasicOp
-> InternaliseM [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc (Exp SOACS -> InternaliseM SubExp)
-> (BasicOp -> Exp SOACS) -> BasicOp -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> InternaliseM [SubExp])
-> BasicOp -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    [PrimValue] -> PrimType -> BasicOp
I.ArrayVal ((PrimValue -> PrimValue) -> [PrimValue] -> [PrimValue]
forall a b. (a -> b) -> [a] -> [b]
map PrimValue -> PrimValue
internalisePrimValue [PrimValue]
vs) (PrimType -> PrimType
internalisePrimType PrimType
t)
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') <- (Exp -> Maybe ([Int], [Exp])) -> [Exp] -> Maybe [([Int], [Exp])]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Exp -> Maybe ([Int], [Exp])
isArrayLiteral [Exp]
es,
    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [Int] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
eshape,
    (([Int], [Exp]) -> Bool) -> [([Int], [Exp])] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (([Int]
eshape ==) ([Int] -> Bool)
-> (([Int], [Exp]) -> [Int]) -> ([Int], [Exp]) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Int], [Exp]) -> [Int]
forall a b. (a, b) -> a
fst) [([Int], [Exp])]
es',
    Just StructType
basetype <- Int -> StructType -> Maybe StructType
forall dim u. Int -> TypeBase dim u -> Maybe (TypeBase dim u)
E.peelArray ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
eshape) StructType
arr_t = do
      let flat_lit :: Exp
flat_lit = [Exp] -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
[ExpBase f vn] -> f StructType -> SrcLoc -> ExpBase f vn
E.ArrayLit ([Exp]
e' [Exp] -> [Exp] -> [Exp]
forall a. [a] -> [a] -> [a]
++ (([Int], [Exp]) -> [Exp]) -> [([Int], [Exp])] -> [Exp]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Int], [Exp]) -> [Exp]
forall a b. (a, b) -> b
snd [([Int], [Exp])]
es') (StructType -> Info StructType
forall a. a -> Info a
Info StructType
basetype) SrcLoc
loc
          new_shape :: [Int]
new_shape = [Exp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Exp]
es Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
eshape
      [VName]
flat_arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"flat_literal" Exp
flat_lit
      [VName] -> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
flat_arrs ((VName -> InternaliseM SubExp) -> InternaliseM [SubExp])
-> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
flat_arr -> do
        TypeBase Shape NoUniqueness
flat_arr_t <- VName -> InternaliseM (TypeBase Shape NoUniqueness)
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
                ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
I.Shape ([SubExp] -> Shape) -> [SubExp] -> Shape
forall a b. (a -> b) -> a -> b
$ (Int -> SubExp) -> [Int] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (IntType -> Integer -> SubExp
intConst IntType
Int64 (Integer -> SubExp) -> (Int -> Integer) -> Int -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Integer
forall a. Integral a => a -> Integer
toInteger) [Int]
new_shape)
                Int
1
                (Shape -> Shape) -> Shape -> Shape
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoUniqueness -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape TypeBase Shape NoUniqueness
flat_arr_t
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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' <- (Exp -> InternaliseM [SubExp]) -> [Exp] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"arr_elem") [Exp]
es
      let arr_t_ext :: [TypeBase ExtShape Uniqueness]
arr_t_ext = (Tree (TypeBase ExtShape Uniqueness)
 -> [TypeBase ExtShape Uniqueness])
-> [Tree (TypeBase ExtShape Uniqueness)]
-> [TypeBase ExtShape Uniqueness]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Tree (TypeBase ExtShape Uniqueness)
-> [TypeBase ExtShape Uniqueness]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList ([Tree (TypeBase ExtShape Uniqueness)]
 -> [TypeBase ExtShape Uniqueness])
-> [Tree (TypeBase ExtShape Uniqueness)]
-> [TypeBase ExtShape Uniqueness]
forall a b. (a -> b) -> a -> b
$ StructType -> [Tree (TypeBase ExtShape Uniqueness)]
internaliseType (StructType -> [Tree (TypeBase ExtShape Uniqueness)])
-> StructType -> [Tree (TypeBase ExtShape Uniqueness)]
forall a b. (a -> b) -> a -> b
$ StructType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
E.toStruct StructType
arr_t

      [TypeBase Shape NoUniqueness]
rowtypes <-
        case (TypeBase ExtShape Uniqueness
 -> Maybe (TypeBase Shape NoUniqueness))
-> [TypeBase ExtShape Uniqueness]
-> Maybe [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
-> Maybe (TypeBase Shape NoUniqueness)
-> Maybe (TypeBase Shape NoUniqueness)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall u. TypeBase Shape u -> TypeBase Shape u
rowType (Maybe (TypeBase Shape NoUniqueness)
 -> Maybe (TypeBase Shape NoUniqueness))
-> (TypeBase ExtShape Uniqueness
    -> Maybe (TypeBase Shape NoUniqueness))
-> TypeBase ExtShape Uniqueness
-> Maybe (TypeBase Shape NoUniqueness)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExtType -> Maybe (TypeBase Shape NoUniqueness)
forall u. TypeBase ExtShape u -> Maybe (TypeBase Shape u)
hasStaticShape (ExtType -> Maybe (TypeBase Shape NoUniqueness))
-> (TypeBase ExtShape Uniqueness -> ExtType)
-> TypeBase ExtShape Uniqueness
-> Maybe (TypeBase Shape NoUniqueness)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase ExtShape Uniqueness -> ExtType
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl) [TypeBase ExtShape Uniqueness]
arr_t_ext of
          Just [TypeBase Shape NoUniqueness]
ts -> [TypeBase Shape NoUniqueness]
-> InternaliseM [TypeBase Shape NoUniqueness]
forall a. a -> InternaliseM a
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
              [] -> [Char] -> InternaliseM [TypeBase Shape NoUniqueness]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [TypeBase Shape NoUniqueness])
-> [Char] -> InternaliseM [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp ArrayLit: existential type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ StructType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString StructType
arr_t
              [SubExp]
e' : [[SubExp]]
_ -> (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
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' <-
              (SubExp -> InternaliseM SubExp)
-> [SubExp] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [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
            Exp SOACS -> InternaliseM (Exp SOACS)
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp SOACS -> InternaliseM (Exp SOACS))
-> Exp SOACS -> InternaliseM (Exp SOACS)
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ [SubExp] -> TypeBase Shape NoUniqueness -> BasicOp
I.ArrayLit [SubExp]
ks' TypeBase Shape NoUniqueness
rt

      (Exp SOACS -> InternaliseM SubExp)
-> [Exp SOACS] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc)
        ([Exp SOACS] -> InternaliseM [SubExp])
-> InternaliseM [Exp SOACS] -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< if [[SubExp]] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[SubExp]]
es'
          then (TypeBase Shape NoUniqueness -> InternaliseM (Exp SOACS))
-> [TypeBase Shape NoUniqueness] -> InternaliseM [Exp SOACS]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([SubExp] -> TypeBase Shape NoUniqueness -> InternaliseM (Exp SOACS)
arraylit []) [TypeBase Shape NoUniqueness]
rowtypes
          else ([SubExp]
 -> TypeBase Shape NoUniqueness -> InternaliseM (Exp SOACS))
-> [[SubExp]]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM [Exp SOACS]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM [SubExp] -> TypeBase Shape NoUniqueness -> InternaliseM (Exp SOACS)
arraylit ([[SubExp]] -> [[SubExp]]
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' <- (Exp -> Maybe ([Int], [Exp])) -> [Exp] -> Maybe [([Int], [Exp])]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Exp -> Maybe ([Int], [Exp])
isArrayLiteral [Exp]
inner_es
      Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ (([Int], [Exp]) -> Bool) -> [([Int], [Exp])] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (([Int]
eshape ==) ([Int] -> Bool)
-> (([Int], [Exp]) -> [Int]) -> ([Int], [Exp]) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Int], [Exp]) -> [Int]
forall a b. (a, b) -> a
fst) [([Int], [Exp])]
inner_es'
      ([Int], [Exp]) -> Maybe ([Int], [Exp])
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Exp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Exp]
inner_es Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
eshape, [Exp]
e [Exp] -> [Exp] -> [Exp]
forall a. [a] -> [a] -> [a]
++ (([Int], [Exp]) -> [Exp]) -> [([Int], [Exp])] -> [Exp]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Int], [Exp]) -> [Exp]
forall a b. (a, b) -> b
snd [([Int], [Exp])]
inner_es')
    isArrayLiteral Exp
e =
      ([Int], [Exp]) -> Maybe ([Int], [Exp])
forall a. a -> Maybe a
Just ([], [Exp
e])
internaliseExp [Char]
desc (E.Ascript Exp
e TypeExp Exp VName
_ SrcLoc
_) =
  [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
internaliseExp [Char]
desc (E.Coerce Exp
e TypeExp Exp VName
_ (Info StructType
et) SrcLoc
loc) = do
  [SubExp]
ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
  [TypeBase ExtShape Uniqueness]
ts <- StructType
-> [TypeBase Shape NoUniqueness] -> [TypeBase ExtShape Uniqueness]
forall shape u.
StructType -> [TypeBase shape u] -> [TypeBase ExtShape Uniqueness]
internaliseCoerceType (StructType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
E.toStruct StructType
et) ([TypeBase Shape NoUniqueness] -> [TypeBase ExtShape Uniqueness])
-> InternaliseM [TypeBase Shape NoUniqueness]
-> InternaliseM [TypeBase ExtShape Uniqueness]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
ses
  [ErrorMsgPart SubExp]
dt' <- StructType -> InternaliseM [ErrorMsgPart SubExp]
forall u. TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError (StructType -> InternaliseM [ErrorMsgPart SubExp])
-> StructType -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ StructType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct StructType
et
  [(SubExp, TypeBase ExtShape Uniqueness)]
-> ((SubExp, TypeBase ExtShape Uniqueness) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([SubExp]
-> [TypeBase ExtShape Uniqueness]
-> [(SubExp, TypeBase ExtShape Uniqueness)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
ses [TypeBase ExtShape Uniqueness]
ts) (((SubExp, TypeBase ExtShape Uniqueness) -> InternaliseM SubExp)
 -> InternaliseM [SubExp])
-> ((SubExp, TypeBase ExtShape Uniqueness) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \(SubExp
e', TypeBase ExtShape Uniqueness
t') -> do
    [SubExp]
dims <- TypeBase Shape NoUniqueness -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
arrayDims (TypeBase Shape NoUniqueness -> [SubExp])
-> InternaliseM (TypeBase Shape NoUniqueness)
-> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType SubExp
e'
    let parts :: [ErrorMsgPart SubExp]
parts =
          [ErrorMsgPart SubExp
"Value of (desugared) shape ["]
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ ErrorMsgPart SubExp
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. a -> [a] -> [a]
intersperse ErrorMsgPart SubExp
"][" ((SubExp -> ErrorMsgPart SubExp)
-> [SubExp] -> [ErrorMsgPart SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64) [SubExp]
dims)
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"] cannot match shape of type `"]
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
dt'
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"`."]
    ErrorMsg SubExp
-> SrcLoc -> ExtType -> [Char] -> SubExp -> InternaliseM SubExp
ensureExtShape ([ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [ErrorMsgPart SubExp]
parts) SrcLoc
loc (TypeBase ExtShape Uniqueness -> ExtType
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 <- SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType SubExp
e'
  case TypeBase Shape NoUniqueness
et of
    I.Prim PrimType
pt ->
      [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (PrimType -> UnOp
I.Neg PrimType
pt) SubExp
e'
    TypeBase Shape NoUniqueness
_ -> [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise.internaliseExp: non-primitive 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 <- SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
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) ->
      [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 ->
      [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (PrimType -> UnOp
I.Neg PrimType
I.Bool) SubExp
e'
    TypeBase Shape NoUniqueness
_ ->
      [Char] -> InternaliseM [SubExp]
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]
src_dims, [SubExp]
ve_dims) <- case ([VName]
srcs, [SubExp]
ves) of
    (VName
src_v : [VName]
_, SubExp
ve_v : [SubExp]
_) ->
      (,)
        ([SubExp] -> [SubExp] -> ([SubExp], [SubExp]))
-> InternaliseM [SubExp]
-> InternaliseM ([SubExp] -> ([SubExp], [SubExp]))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeBase Shape NoUniqueness -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
I.arrayDims (TypeBase Shape NoUniqueness -> [SubExp])
-> InternaliseM (TypeBase Shape NoUniqueness)
-> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
src_v)
        InternaliseM ([SubExp] -> ([SubExp], [SubExp]))
-> InternaliseM [SubExp] -> InternaliseM ([SubExp], [SubExp])
forall a b.
InternaliseM (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (TypeBase Shape NoUniqueness -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
I.arrayDims (TypeBase Shape NoUniqueness -> [SubExp])
-> InternaliseM (TypeBase Shape NoUniqueness)
-> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType SubExp
ve_v)
    ([VName], [SubExp])
_ -> ([SubExp], [SubExp]) -> InternaliseM ([SubExp], [SubExp])
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], []) -- Will this happen?
  ([DimIndex SubExp]
idxs', Certs
cs) <- SrcLoc
-> [SubExp]
-> SliceBase Info VName
-> InternaliseM ([DimIndex SubExp], Certs)
internaliseSlice SrcLoc
loc [SubExp]
src_dims SliceBase Info VName
slice
  let src_dims' :: [SubExp]
src_dims' = Slice SubExp -> [SubExp]
forall d. Slice d -> [d]
sliceDims ([DimIndex SubExp] -> Slice SubExp
forall d. [DimIndex d] -> Slice d
Slice [DimIndex SubExp]
idxs')
      rank :: Int
rank = [SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
src_dims'
      errormsg :: ErrorMsg SubExp
errormsg =
        ErrorMsg SubExp
"Shape "
          ErrorMsg SubExp -> ErrorMsg SubExp -> ErrorMsg SubExp
forall a. Semigroup a => a -> a -> a
<> [SubExp] -> ErrorMsg SubExp
forall a. [a] -> ErrorMsg a
errorShape [SubExp]
src_dims'
          ErrorMsg SubExp -> ErrorMsg SubExp -> ErrorMsg SubExp
forall a. Semigroup a => a -> a -> a
<> ErrorMsg SubExp
" of slice does not match shape "
          ErrorMsg SubExp -> ErrorMsg SubExp -> ErrorMsg SubExp
forall a. Semigroup a => a -> a -> a
<> [SubExp] -> ErrorMsg SubExp
forall a. [a] -> ErrorMsg a
errorShape (Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
take Int
rank [SubExp]
ve_dims)
          ErrorMsg SubExp -> ErrorMsg SubExp -> ErrorMsg SubExp
forall a. Semigroup a => a -> a -> a
<> ErrorMsg SubExp
" of value."

  let comb :: VName -> SubExp -> InternaliseM VName
comb VName
sname SubExp
ve' = do
        TypeBase Shape NoUniqueness
sname_t <- VName -> InternaliseM (TypeBase Shape NoUniqueness)
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 TypeBase Shape NoUniqueness
-> [SubExp] -> TypeBase Shape NoUniqueness
forall oldshape u.
TypeBase oldshape u -> [SubExp] -> TypeBase Shape u
`setArrayDims` Slice SubExp -> [SubExp]
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
errormsg SrcLoc
loc TypeBase Shape NoUniqueness
rowtype [Char]
"lw_val_correct_shape" SubExp
ve'
        [Char]
-> VName
-> Slice SubExp
-> Exp (Rep InternaliseM)
-> InternaliseM VName
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> VName -> Slice SubExp -> Exp (Rep m) -> m VName
letInPlace [Char]
desc VName
sname Slice SubExp
full_slice (Exp (Rep InternaliseM) -> InternaliseM VName)
-> Exp (Rep InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
ve''
  Certs -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a. Certs -> InternaliseM a -> InternaliseM a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var ([VName] -> [SubExp])
-> InternaliseM [VName] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> SubExp -> InternaliseM VName)
-> [VName] -> [SubExp] -> InternaliseM [VName]
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
  StructType
-> [Name] -> [SubExp] -> [SubExp] -> InternaliseM [SubExp]
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 <- Name -> Map Name (TypeBase Exp als) -> Maybe (TypeBase Exp als)
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 =
                [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int)
-> ([(Name, TypeBase Exp als)] -> [Int])
-> [(Name, TypeBase Exp als)]
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Name, TypeBase Exp als) -> Int)
-> [(Name, TypeBase Exp als)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase Exp als -> Int
forall als. TypeBase Exp als -> Int
internalisedTypeSize (TypeBase Exp als -> Int)
-> ((Name, TypeBase Exp als) -> TypeBase Exp als)
-> (Name, TypeBase Exp als)
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, TypeBase Exp als) -> TypeBase Exp als
forall a b. (a, b) -> b
snd) ([(Name, TypeBase Exp als)] -> Int)
-> [(Name, TypeBase Exp als)] -> Int
forall a b. (a -> b) -> a -> b
$
                  ((Name, TypeBase Exp als) -> Bool)
-> [(Name, TypeBase Exp als)] -> [(Name, TypeBase Exp als)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
/= Name
f) (Name -> Bool)
-> ((Name, TypeBase Exp als) -> Name)
-> (Name, TypeBase Exp als)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, TypeBase Exp als) -> Name
forall a b. (a, b) -> a
fst) ([(Name, TypeBase Exp als)] -> [(Name, TypeBase Exp als)])
-> (Map Name (TypeBase Exp als) -> [(Name, TypeBase Exp als)])
-> Map Name (TypeBase Exp als)
-> [(Name, TypeBase Exp als)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name (TypeBase Exp als) -> [(Name, TypeBase Exp als)]
forall a. Map Name a -> [(Name, a)]
sortFields (Map Name (TypeBase Exp als) -> [(Name, TypeBase Exp als)])
-> Map Name (TypeBase Exp als) -> [(Name, TypeBase Exp als)]
forall a b. (a -> b) -> a -> b
$
                    Map Name (TypeBase Exp als)
m
              k :: Int
k = TypeBase Exp als -> Int
forall als. TypeBase Exp als -> Int
internalisedTypeSize TypeBase Exp als
t
              ([a]
bef, [a]
to_update, [a]
aft) = Int -> Int -> [a] -> ([a], [a], [a])
forall a. Int -> Int -> [a] -> ([a], [a], [a])
splitAt3 Int
i Int
k [a]
src'
          [a]
src'' <- TypeBase Exp als -> [Name] -> [a] -> [a] -> m [a]
replace TypeBase Exp als
t [Name]
fs [a]
ve' [a]
to_update
          [a] -> m [a]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([a] -> m [a]) -> [a] -> m [a]
forall a b. (a -> b) -> a -> b
$ [a]
bef [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
src'' [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
aft
    replace TypeBase Exp als
_ [Name]
_ [a]
ve' [a]
_ = [a] -> m [a]
forall a. a -> m 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' <- (InternaliseEnv -> InternaliseEnv)
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a.
(InternaliseEnv -> InternaliseEnv)
-> InternaliseM a -> InternaliseM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (Attr -> InternaliseEnv -> InternaliseEnv
f Attr
attr') (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
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 ([Char] -> Text) -> [Char] -> Text
forall a b. (a -> b) -> a -> b
$ SrcLoc -> [Char]
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" ->
      (SubExp -> InternaliseM SubExp)
-> [SubExp] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc (Exp SOACS -> InternaliseM SubExp)
-> (SubExp -> Exp SOACS) -> SubExp -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS)
-> (SubExp -> BasicOp) -> SubExp -> Exp SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpaqueOp -> SubExp -> BasicOp
Opaque OpaqueOp
OpaqueNil) [SubExp]
e'
    Attr
_ ->
      [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
e'
  where
    traceRes :: Text -> [SubExp] -> InternaliseM [SubExp]
traceRes Text
tag' =
      (SubExp -> InternaliseM SubExp)
-> [SubExp] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc (Exp SOACS -> InternaliseM SubExp)
-> (SubExp -> Exp SOACS) -> SubExp -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS)
-> (SubExp -> BasicOp) -> SubExp -> Exp SOACS
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' Attr -> Attr -> Bool
forall a. Eq a => a -> a -> Bool
== Attr
"unsafe",
        Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ InternaliseEnv -> Bool
envSafe InternaliseEnv
env =
          InternaliseEnv
env {envDoBoundsChecks = False}
      | Bool
otherwise =
          InternaliseEnv
env {envAttrs = envAttrs env <> oneAttr 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' ([ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [Text -> ErrorMsgPart SubExp
forall a. Text -> ErrorMsgPart a
ErrorString (Text -> ErrorMsgPart SubExp) -> Text -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ Text
"Assertion is false: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
check]) SrcLoc
loc
  -- Make sure there are some bindings to certify.
  Certs -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a. Certs -> InternaliseM a -> InternaliseM a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ (SubExp -> InternaliseM SubExp)
-> [SubExp] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> InternaliseM SubExp
forall {m :: * -> *}. MonadBuilder m => SubExp -> m SubExp
rebind ([SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e2
  where
    rebind :: SubExp -> m SubExp
rebind SubExp
v = do
      VName
v' <- [Char] -> m VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"assert_res"
      [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
v'] (Exp (Rep m) -> m ()) -> Exp (Rep m) -> m ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
v
      SubExp -> m SubExp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> m SubExp) -> SubExp -> m SubExp
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, [(Name, [Int])]
constr_map) <- Map Name [StructType]
-> InternaliseM ([TypeBase ExtShape Uniqueness], [(Name, [Int])])
internaliseSumType (Map Name [StructType]
 -> InternaliseM ([TypeBase ExtShape Uniqueness], [(Name, [Int])]))
-> Map Name [StructType]
-> InternaliseM ([TypeBase ExtShape Uniqueness], [(Name, [Int])])
forall a b. (a -> b) -> a -> b
$ ([StructType] -> [StructType])
-> Map Name [StructType] -> Map Name [StructType]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((StructType -> StructType) -> [StructType] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map StructType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
E.toStruct) Map Name [StructType]
fs
  [SubExp]
es' <- [[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[SubExp]] -> [SubExp])
-> InternaliseM [[SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> InternaliseM [SubExp]) -> [Exp] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"payload") [Exp]
es

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

  case Name -> [(Name, [Int])] -> Maybe (Int, [Int])
forall a b. Eq a => a -> [(a, b)] -> Maybe (Int, b)
lookupWithIndex Name
c [(Name, [Int])]
constr_map of
    Just (Int
i, [Int]
js) ->
      (IntType -> Integer -> SubExp
intConst IntType
Int8 (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
i) :) ([SubExp] -> [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int
-> [TypeBase Shape NoUniqueness]
-> [(Int, SubExp)]
-> InternaliseM [SubExp]
forall {f :: * -> *} {t}.
(Num t, MonadBuilder f, Eq t) =>
t -> [TypeBase Shape NoUniqueness] -> [(t, SubExp)] -> f [SubExp]
clauses Int
0 [TypeBase Shape NoUniqueness]
ts' ([Int] -> [SubExp] -> [(Int, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
js [SubExp]
es')
    Maybe (Int, [Int])
Nothing ->
      [Char] -> InternaliseM [SubExp]
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 t -> [(t, SubExp)] -> Maybe SubExp
forall a b. Eq a => a -> [(a, b)] -> Maybe b
`lookup` [(t, SubExp)]
js_to_es =
          (SubExp
e :) ([SubExp] -> [SubExp]) -> f [SubExp] -> f [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> t -> [TypeBase Shape NoUniqueness] -> [(t, SubExp)] -> f [SubExp]
clauses (t
j t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) [TypeBase Shape NoUniqueness]
ts [(t, SubExp)]
js_to_es
      | Bool
otherwise = do
          SubExp
blank <-
            -- Cannot use eBlank here for arrays, because when doing
            -- equality comparisons on sum types, we end up looking at
            -- the array elements. (#2081) This is a bit of an edge
            -- case, but arrays in sum types are known to be
            -- inefficient.
            [Char] -> Exp (Rep f) -> f SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"zero"
              (Exp (Rep f) -> f SubExp) -> f (Exp (Rep f)) -> f SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< case TypeBase Shape NoUniqueness
t of
                I.Array {} ->
                  Exp (Rep f) -> f (Exp (Rep f))
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp (Rep f) -> f (Exp (Rep f))) -> Exp (Rep f) -> f (Exp (Rep f))
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep f)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep f)) -> BasicOp -> Exp (Rep f)
forall a b. (a -> b) -> a -> b
$ Shape -> SubExp -> BasicOp
Replicate (TypeBase Shape NoUniqueness -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape TypeBase Shape NoUniqueness
t) (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimValue
blankPrimValue (PrimType -> PrimValue) -> PrimType -> PrimValue
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoUniqueness -> PrimType
forall shape u. TypeBase shape u -> PrimType
elemType TypeBase Shape NoUniqueness
t
                TypeBase Shape NoUniqueness
_ -> TypeBase Shape NoUniqueness -> f (Exp (Rep f))
forall (m :: * -> *).
MonadBuilder m =>
TypeBase Shape NoUniqueness -> m (Exp (Rep m))
eBlank TypeBase Shape NoUniqueness
t
          (SubExp
blank :) ([SubExp] -> [SubExp]) -> f [SubExp] -> f [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> t -> [TypeBase Shape NoUniqueness] -> [(t, SubExp)] -> f [SubExp]
clauses (t
j t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) [TypeBase Shape NoUniqueness]
ts [(t, SubExp)]
js_to_es
    clauses t
_ [] [(t, SubExp)]
_ =
      [SubExp] -> f [SubExp]
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
internaliseExp [Char]
_ (E.Constr Name
_ [Exp]
_ (Info StructType
t) SrcLoc
loc) =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: constructor with type " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ StructType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString StructType
t [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
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
_) =
  [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
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)) ->
      [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
I.IntValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> IntValue
forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
v]
    E.Scalar (E.Prim (E.Unsigned IntType
it)) ->
      [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
I.IntValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> IntValue
forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
v]
    E.Scalar (E.Prim (E.FloatType FloatType
ft)) ->
      [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
I.FloatValue (FloatValue -> PrimValue) -> FloatValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ FloatType -> Integer -> FloatValue
forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
ft Integer
v]
    StructType
_ -> [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: nonsensical type for integer literal: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ StructType -> [Char]
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)) ->
      [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
I.FloatValue (FloatValue -> PrimValue) -> FloatValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ FloatType -> Double -> FloatValue
forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
ft Double
v]
    StructType
_ -> [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: nonsensical type for float literal: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ StructType -> [Char]
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' = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> ([StructType] -> [Int]) -> [StructType] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (StructType -> Int) -> [StructType] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map StructType -> Int
forall als. TypeBase Exp als -> Int
internalisedTypeSize ([StructType] -> Int) -> [StructType] -> Int
forall a b. (a -> b) -> a -> b
$
        case Exp -> StructType
E.typeOf Exp
e of
          E.Scalar (Record Map Name StructType
fs) ->
            ((Name, StructType) -> StructType)
-> [(Name, StructType)] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map (Name, StructType) -> StructType
forall a b. (a, b) -> b
snd ([(Name, StructType)] -> [StructType])
-> [(Name, StructType)] -> [StructType]
forall a b. (a -> b) -> a -> b
$ ((Name, StructType) -> Bool)
-> [(Name, StructType)] -> [(Name, StructType)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
/= Name
k) (Name -> Bool)
-> ((Name, StructType) -> Name) -> (Name, StructType) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, StructType) -> Name
forall a b. (a, b) -> a
fst) ([(Name, StructType)] -> [(Name, StructType)])
-> [(Name, StructType)] -> [(Name, StructType)]
forall a b. (a -> b) -> a -> b
$ Map Name StructType -> [(Name, StructType)]
forall a. Map Name a -> [(Name, a)]
sortFields Map Name StructType
fs
          StructType
t -> [StructType
t]
  Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
take (StructType -> Int
forall als. TypeBase Exp als -> Int
internalisedTypeSize StructType
rt) ([SubExp] -> [SubExp])
-> ([SubExp] -> [SubExp]) -> [SubExp] -> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
drop Int
i' ([SubExp] -> [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
internaliseExp [Char]
_ e :: Exp
e@E.Lambda {} =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected lambda at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp [Char]
_ e :: Exp
e@E.OpSection {} =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected operator section at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp [Char]
_ e :: Exp
e@E.OpSectionLeft {} =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected left operator section at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp [Char]
_ e :: Exp
e@E.OpSectionRight {} =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected right operator section at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp [Char]
_ e :: Exp
e@E.ProjectSection {} =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected projection section at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp [Char]
_ e :: Exp
e@E.IndexSection {} =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected index section at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr (Exp -> SrcLoc
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 <- InternaliseM (Scope SOACS)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  case Maybe VName
argdim of
    Just VName
d | VName
d VName -> Scope SOACS -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Scope SOACS
exists -> [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
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
          [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
d] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
        ([SubExp], Maybe VName)
_ -> () -> InternaliseM ()
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
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 (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> IntValue
forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
x
internalisePatLit (E.PatLitInt Integer
x) (E.Scalar (E.Prim (E.Unsigned IntType
it))) =
  IntValue -> PrimValue
I.IntValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> IntValue
forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
x
internalisePatLit (E.PatLitFloat Double
x) (E.Scalar (E.Prim (E.FloatType FloatType
ft))) =
  FloatValue -> PrimValue
I.FloatValue (FloatValue -> PrimValue) -> FloatValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ FloatType -> Double -> FloatValue
forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
ft Double
x
internalisePatLit PatLit
l StructType
t =
  [Char] -> PrimValue
forall a. HasCallStack => [Char] -> a
error ([Char] -> PrimValue) -> [Char] -> PrimValue
forall a b. (a -> b) -> a -> b
$ [Char]
"Nonsensical pattern and type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (PatLit, StructType) -> [Char]
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]
_) <- PatBase Info VName StructType
-> [SubExp] -> InternaliseM ([Maybe PrimValue], [SubExp], [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
  ([Maybe PrimValue], [SubExp])
-> InternaliseM ([Maybe PrimValue], [SubExp])
forall a. a -> InternaliseM a
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) =
      ([Maybe PrimValue], [a], [a])
-> InternaliseM ([Maybe PrimValue], [a], [a])
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue) -> PrimValue -> Maybe PrimValue
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, [(Name, [Int])]
m) <- Map Name [StructType]
-> InternaliseM ([TypeBase ExtShape Uniqueness], [(Name, [Int])])
internaliseSumType (Map Name [StructType]
 -> InternaliseM ([TypeBase ExtShape Uniqueness], [(Name, [Int])]))
-> Map Name [StructType]
-> InternaliseM ([TypeBase ExtShape Uniqueness], [(Name, [Int])])
forall a b. (a -> b) -> a -> b
$ ([StructType] -> [StructType])
-> Map Name [StructType] -> Map Name [StructType]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((StructType -> StructType) -> [StructType] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map StructType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct) Map Name [StructType]
fs
      case Name -> [(Name, [Int])] -> Maybe (Int, [Int])
forall a b. Eq a => a -> [(a, b)] -> Maybe (Int, b)
lookupWithIndex Name
c [(Name, [Int])]
m of
        Just (Int
tag, [Int]
payload_is) -> do
          let ([a]
payload_ses, [a]
ses') = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt ([TypeBase ExtShape Uniqueness] -> Int
forall a. [a] -> Int
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 ([a] -> InternaliseM ([Maybe PrimValue], [a], [a]))
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
forall a b. (a -> b) -> a -> b
$ (Int -> a) -> [Int] -> [a]
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 Int -> [Int] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
`elemIndex` [Int]
payload_is of
                  Just Int
j -> [Maybe PrimValue]
cmps [Maybe PrimValue] -> Int -> Maybe PrimValue
forall a. HasCallStack => [a] -> Int -> a
!! Int
j
                  Maybe Int
Nothing -> Maybe PrimValue
forall a. Maybe a
Nothing
          ([Maybe PrimValue], [a], [a])
-> InternaliseM ([Maybe PrimValue], [a], [a])
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
            ( PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (IntValue -> PrimValue
I.IntValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> IntValue
forall int. Integral int => IntType -> int -> IntValue
intValue IntType
Int8 (Integer -> IntValue) -> Integer -> IntValue
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
tag)
                Maybe PrimValue -> [Maybe PrimValue] -> [Maybe PrimValue]
forall a. a -> [a] -> [a]
: (Int -> a -> Maybe PrimValue) -> [Int] -> [a] -> [Maybe PrimValue]
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 ->
          [Char] -> InternaliseM ([Maybe PrimValue], [a], [a])
forall a. HasCallStack => [Char] -> a
error [Char]
"generateCond: missing constructor"
    compares (E.PatConstr Name
_ (Info StructType
t) [PatBase Info vn StructType]
_ SrcLoc
_) [a]
_ =
      [Char] -> InternaliseM ([Maybe PrimValue], [a], [a])
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM ([Maybe PrimValue], [a], [a]))
-> [Char] -> InternaliseM ([Maybe PrimValue], [a], [a])
forall a b. (a -> b) -> a -> b
$ [Char]
"generateCond: PatConstr has nonsensical type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ StructType -> [Char]
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 (Info StructType -> SrcLoc -> PatBase Info vn StructType
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) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt (StructType -> Int
forall als. TypeBase Exp als -> Int
internalisedTypeSize (StructType -> Int) -> StructType -> Int
forall a b. (a -> b) -> a -> b
$ StructType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
E.toStruct StructType
t) [a]
ses
      ([Maybe PrimValue], [a], [a])
-> InternaliseM ([Maybe PrimValue], [a], [a])
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((a -> Maybe PrimValue) -> [a] -> [Maybe PrimValue]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe PrimValue -> a -> Maybe PrimValue
forall a b. a -> b -> a
const Maybe PrimValue
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 (Info StructType -> SrcLoc -> PatBase Info vn StructType
forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
E.Wildcard (StructType -> Info StructType
forall a. a -> Info a
Info (StructType -> Info StructType) -> StructType -> Info StructType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
E.Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
-> ScalarTypeBase Exp NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ Map Name StructType -> ScalarTypeBase Exp NoUniqueness
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
E.Record Map Name StructType
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 (Info StructType -> SrcLoc -> PatBase Info vn StructType
forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
E.Wildcard (StructType -> Info StructType
forall a. a -> Info a
Info (StructType -> Info StructType) -> StructType -> Info StructType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
E.Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
-> ScalarTypeBase Exp NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ Map Name StructType -> ScalarTypeBase Exp NoUniqueness
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
E.Record Map Name StructType
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 [(L Name, PatBase Info vn StructType)]
fs SrcLoc
_) [a]
ses =
      [PatBase Info vn StructType]
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
comparesMany (((Name, PatBase Info vn StructType) -> PatBase Info vn StructType)
-> [(Name, PatBase Info vn StructType)]
-> [PatBase Info vn StructType]
forall a b. (a -> b) -> [a] -> [b]
map (Name, PatBase Info vn StructType) -> PatBase Info vn StructType
forall a b. (a, b) -> b
snd ([(Name, PatBase Info vn StructType)]
 -> [PatBase Info vn StructType])
-> [(Name, PatBase Info vn StructType)]
-> [PatBase Info vn StructType]
forall a b. (a -> b) -> a -> b
$ Map Name (PatBase Info vn StructType)
-> [(Name, PatBase Info vn StructType)]
forall a. Map Name a -> [(Name, a)]
E.sortFields (Map Name (PatBase Info vn StructType)
 -> [(Name, PatBase Info vn StructType)])
-> Map Name (PatBase Info vn StructType)
-> [(Name, PatBase Info vn StructType)]
forall a b. (a -> b) -> a -> b
$ [(Name, PatBase Info vn StructType)]
-> Map Name (PatBase Info vn StructType)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, PatBase Info vn StructType)]
 -> Map Name (PatBase Info vn StructType))
-> [(Name, PatBase Info vn StructType)]
-> Map Name (PatBase Info vn StructType)
forall a b. (a -> b) -> a -> b
$ ((L Name, PatBase Info vn StructType)
 -> (Name, PatBase Info vn StructType))
-> [(L Name, PatBase Info vn StructType)]
-> [(Name, PatBase Info vn StructType)]
forall a b. (a -> b) -> [a] -> [b]
map ((L Name -> Name)
-> (L Name, PatBase Info vn StructType)
-> (Name, PatBase Info vn StructType)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first L Name -> Name
forall a. L a -> a
unLoc) [(L Name, PatBase Info vn StructType)]
fs) [a]
ses
    compares (E.PatAscription PatBase Info vn StructType
pat TypeExp (ExpBase Info vn) 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 [] =
      [Char] -> InternaliseM ([Maybe PrimValue], [a], [a])
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM ([Maybe PrimValue], [a], [a]))
-> [Char] -> InternaliseM ([Maybe PrimValue], [a], [a])
forall a b. (a -> b) -> a -> b
$ [Char]
"generateCond: No values left for pattern " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ PatBase Info vn StructType -> [Char]
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 = ([Maybe PrimValue], [a], [a])
-> InternaliseM ([Maybe PrimValue], [a], [a])
forall a. a -> InternaliseM a
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'
      ([Maybe PrimValue], [a], [a])
-> InternaliseM ([Maybe PrimValue], [a], [a])
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        ( [Maybe PrimValue]
cmps1 [Maybe PrimValue] -> [Maybe PrimValue] -> [Maybe PrimValue]
forall a. Semigroup a => a -> a -> a
<> [Maybe PrimValue]
cmps2,
          [a]
pertinent1 [a] -> [a] -> [a]
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
  [SizeBinder VName]
-> PatBase Info VName StructType
-> [SubExp]
-> InternaliseM a
-> InternaliseM a
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 PatBase Info VName StructType -> [IdentBase Info VName StructType]
forall (f :: * -> *) vn t. PatBase f vn t -> [IdentBase f vn t]
E.patIdents PatBase Info VName StructType
p of
      [IdentBase Info VName StructType
v] -> VName -> [Char]
baseString (VName -> [Char]) -> VName -> [Char]
forall a b. (a -> b) -> a -> b
$ IdentBase Info VName StructType -> VName
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 <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
ses
  PatBase Info VName ParamType
-> [TypeBase Shape NoUniqueness]
-> ([VName] -> InternaliseM a)
-> InternaliseM a
forall a.
PatBase Info VName ParamType
-> [TypeBase Shape NoUniqueness]
-> ([VName] -> InternaliseM a)
-> InternaliseM a
stmPat (Diet -> StructType -> ParamType
forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
E.Observe (StructType -> ParamType)
-> PatBase Info VName StructType -> PatBase Info VName ParamType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PatBase Info VName StructType
p) [TypeBase Shape NoUniqueness]
ses_ts (([VName] -> InternaliseM a) -> InternaliseM a)
-> ([VName] -> InternaliseM a) -> InternaliseM a
forall a b. (a -> b) -> a -> b
$ \[VName]
pat_names -> do
    AppRes -> [SubExp] -> InternaliseM ()
bindExtSizes (StructType -> [VName] -> AppRes
AppRes (PatBase Info VName StructType -> StructType
forall d u. Pat (TypeBase d u) -> TypeBase d u
E.patternType PatBase Info VName StructType
p) ((SizeBinder VName -> VName) -> [SizeBinder VName] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map SizeBinder VName -> VName
forall vn. SizeBinder vn -> vn
E.sizeName [SizeBinder VName]
sizes)) [SubExp]
ses
    [(VName, SubExp)]
-> ((VName, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([VName] -> [SubExp] -> [(VName, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
pat_names [SubExp]
ses) (((VName, SubExp) -> InternaliseM ()) -> InternaliseM ())
-> ((VName, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \(VName
v, SubExp
se) ->
      [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
v] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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) <- [(DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])]
-> ([DimIndex SubExp], [SubExp], [[ErrorMsgPart SubExp]])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 ([(DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])]
 -> ([DimIndex SubExp], [SubExp], [[ErrorMsgPart SubExp]]))
-> InternaliseM [(DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])]
-> InternaliseM
     ([DimIndex SubExp], [SubExp], [[ErrorMsgPart SubExp]])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp
 -> DimIndex
 -> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp]))
-> [SubExp]
-> SliceBase Info VName
-> InternaliseM [(DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM SubExp
-> DimIndex
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
internaliseDimIndex [SubExp]
dims SliceBase Info VName
idxs
  SubExp
ok <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"index_ok" (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [SubExp] -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
eAll [SubExp]
oks
  let msg :: ErrorMsg SubExp
msg =
        [ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg ([ErrorMsgPart SubExp] -> ErrorMsg SubExp)
-> [ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a b. (a -> b) -> a -> b
$
          [ErrorMsgPart SubExp
"Index ["]
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
-> [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
", "] [[ErrorMsgPart SubExp]]
parts
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"] out of bounds for array of shape ["]
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ ErrorMsgPart SubExp
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. a -> [a] -> [a]
intersperse ErrorMsgPart SubExp
"][" ((SubExp -> ErrorMsgPart SubExp)
-> [SubExp] -> [ErrorMsgPart SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64) ([SubExp] -> [ErrorMsgPart SubExp])
-> [SubExp] -> [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
take (SliceBase Info VName -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length SliceBase Info VName
idxs) [SubExp]
dims)
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
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
  ([DimIndex SubExp], Certs)
-> InternaliseM ([DimIndex SubExp], Certs)
forall a. a -> InternaliseM a
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 =
        BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
I.Int64) (Int64 -> SubExp
forall v. IsValue v => v -> SubExp
I.constant (Int64
0 :: I.Int64)) SubExp
i'
      upperBound :: Exp SOACS
upperBound =
        BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"bounds_check" (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinOp
-> InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eBinOp BinOp
I.LogAnd (Exp SOACS -> InternaliseM (Exp SOACS)
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp SOACS
lowerBound) (Exp SOACS -> InternaliseM (Exp SOACS)
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp SOACS
upperBound)
  (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
I.DimFix SubExp
i', SubExp
ok, [PrimType -> SubExp -> ErrorMsgPart SubExp
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 <-
      [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"w_minus_1" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
        BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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
    (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
      ( SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
I.DimSlice SubExp
w_minus_1 SubExp
w (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 (-Integer
1),
        Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
True,
        [ErrorMsgPart SubExp]
forall a. Monoid a => a
mempty
      )
    where
      one :: SubExp
one = Int64 -> SubExp
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' <- InternaliseM SubExp
-> (Exp -> InternaliseM SubExp) -> Maybe Exp -> InternaliseM SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (SubExp -> InternaliseM SubExp
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
one) (((SubExp, IntType) -> SubExp)
-> InternaliseM (SubExp, IntType) -> InternaliseM SubExp
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SubExp, IntType) -> SubExp
forall a b. (a, b) -> a
fst (InternaliseM (SubExp, IntType) -> InternaliseM SubExp)
-> (Exp -> InternaliseM (SubExp, IntType))
-> Exp
-> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp -> InternaliseM (SubExp, IntType)
internaliseSizeExp [Char]
"s") Maybe Exp
s
  SubExp
s_sign <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"s_sign" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (IntType -> UnOp
I.SSignum IntType
Int64) SubExp
s'
  SubExp
backwards <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"backwards" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"w_minus_1" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 =
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"i_def"
          (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
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
            (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
backwards)
            ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
w_minus_1])
            ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
zero])
      j_def :: InternaliseM SubExp
j_def =
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"j_def"
          (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
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
            (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
backwards)
            ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
negone])
            ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
w])
  SubExp
i' <- InternaliseM SubExp
-> (Exp -> InternaliseM SubExp) -> Maybe Exp -> InternaliseM SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe InternaliseM SubExp
i_def (((SubExp, IntType) -> SubExp)
-> InternaliseM (SubExp, IntType) -> InternaliseM SubExp
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SubExp, IntType) -> SubExp
forall a b. (a, b) -> a
fst (InternaliseM (SubExp, IntType) -> InternaliseM SubExp)
-> (Exp -> InternaliseM (SubExp, IntType))
-> Exp
-> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp -> InternaliseM (SubExp, IntType)
internaliseSizeExp [Char]
"i") Maybe Exp
i
  SubExp
j' <- InternaliseM SubExp
-> (Exp -> InternaliseM SubExp) -> Maybe Exp -> InternaliseM SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe InternaliseM SubExp
j_def (((SubExp, IntType) -> SubExp)
-> InternaliseM (SubExp, IntType) -> InternaliseM SubExp
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SubExp, IntType) -> SubExp
forall a b. (a, b) -> a
fst (InternaliseM (SubExp, IntType) -> InternaliseM SubExp)
-> (Exp -> InternaliseM (SubExp, IntType))
-> Exp
-> InternaliseM SubExp
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"j_m_i" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 =
        BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
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)
          ( BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
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
              (BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
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 (m (Exp (Rep m)) -> m (Exp (Rep m))
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"n" (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
forall {m :: * -> *}.
MonadBuilder m =>
m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
divRounding (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
toExp SubExp
j_m_i) (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
toExp SubExp
s')

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

  SubExp
m <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"m" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"m_t_s" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"i_p_m_t_s" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <-
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"zero_leq_i_p_m_t_s" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <-
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"i_p_m_t_s_leq_w" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <-
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"i_p_m_t_s_leq_w" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"zero_lte_i" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"i_lte_j" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <-
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"forwards_ok"
      (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [SubExp] -> InternaliseM (Exp (Rep InternaliseM))
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"negone_lte_j" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"j_lte_i" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <-
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"backwards_ok"
      (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [SubExp] -> InternaliseM (Exp (Rep InternaliseM))
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 <-
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"slice_ok"
      (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
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
        (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
backwards)
        ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
backwards_ok])
        ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
forwards_ok])

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

  SubExp
acceptable <-
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"slice_acceptable" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 {}) ->
          [ ErrorMsgPart SubExp
-> (Exp -> ErrorMsgPart SubExp) -> Maybe Exp -> ErrorMsgPart SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp
forall a b. a -> b -> a
const (ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp)
-> ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
i') Maybe Exp
i,
            ErrorMsgPart SubExp
":",
            ErrorMsgPart SubExp
-> (Exp -> ErrorMsgPart SubExp) -> Maybe Exp -> ErrorMsgPart SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp
forall a b. a -> b -> a
const (ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp)
-> ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
j') Maybe Exp
j,
            ErrorMsgPart SubExp
":",
            PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
s'
          ]
        (Maybe Exp
_, Just {}, Maybe Exp
_) ->
          [ ErrorMsgPart SubExp
-> (Exp -> ErrorMsgPart SubExp) -> Maybe Exp -> ErrorMsgPart SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp
forall a b. a -> b -> a
const (ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp)
-> ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
i') Maybe Exp
i,
            ErrorMsgPart SubExp
":",
            PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
j'
          ]
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
-> (Exp -> [ErrorMsgPart SubExp])
-> Maybe Exp
-> [ErrorMsgPart SubExp]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [ErrorMsgPart SubExp]
forall a. Monoid a => a
mempty ([ErrorMsgPart SubExp] -> Exp -> [ErrorMsgPart SubExp]
forall a b. a -> b -> a
const [ErrorMsgPart SubExp
":", PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
s']) Maybe Exp
s
        (Maybe Exp
_, Maybe Exp
Nothing, Maybe Exp
Nothing) ->
          [PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
i', ErrorMsgPart SubExp
":"]
  (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
I.DimSlice SubExp
i' SubExp
n SubExp
s', SubExp
acceptable, [ErrorMsgPart SubExp]
parts)
  where
    zero :: SubExp
zero = Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)
    negone :: SubExp
negone = Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (-Int64
1 :: Int64)
    one :: SubExp
one = Int64 -> SubExp
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 [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"_arr") Exp
arr
  [SubExp]
nes <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp ([Char]
what [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"_ne") Exp
ne
  [SubExp]
nes' <- [(SubExp, VName)]
-> ((SubExp, VName) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([SubExp] -> [VName] -> [(SubExp, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
nes [VName]
arrs) (((SubExp, VName) -> InternaliseM SubExp) -> InternaliseM [SubExp])
-> ((SubExp, VName) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \(SubExp
ne', VName
arr') -> do
    TypeBase Shape NoUniqueness
rowtype <- Int -> TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall u. Int -> TypeBase Shape u -> TypeBase Shape u
I.stripArray Int
1 (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
-> InternaliseM (TypeBase Shape NoUniqueness)
-> InternaliseM (TypeBase Shape NoUniqueness)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM (TypeBase Shape NoUniqueness)
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 [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"_ne_right_shape")
      SubExp
ne'
  [TypeBase Shape NoUniqueness]
nests <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
I.subExpType [SubExp]
nes'
  [TypeBase Shape NoUniqueness]
arrts <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
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 <- Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 ([TypeBase Shape NoUniqueness] -> SubExp)
-> InternaliseM [TypeBase Shape NoUniqueness]
-> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
  [Char] -> Exp SOACS -> InternaliseM [SubExp]
letValExp' [Char]
desc (Exp SOACS -> InternaliseM [SubExp])
-> (SOAC SOACS -> Exp SOACS) -> SOAC SOACS -> InternaliseM [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Op SOACS -> Exp SOACS
SOAC SOACS -> Exp SOACS
forall rep. Op rep -> Exp rep
I.Op (SOAC SOACS -> InternaliseM [SubExp])
-> InternaliseM (SOAC SOACS) -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SubExp
-> Lambda 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 <- [(SubExp, VName)]
-> ((SubExp, VName) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([SubExp] -> [VName] -> [(SubExp, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
ne' [VName]
hist') (((SubExp, VName) -> InternaliseM SubExp) -> InternaliseM [SubExp])
-> ((SubExp, VName) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \(SubExp
n, VName
h) -> do
    TypeBase Shape NoUniqueness
rowtype <- Int -> TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall u. Int -> TypeBase Shape u -> TypeBase Shape u
I.stripArray Int
1 (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
-> InternaliseM (TypeBase Shape NoUniqueness)
-> InternaliseM (TypeBase Shape NoUniqueness)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM (TypeBase Shape NoUniqueness)
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 <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
I.subExpType [SubExp]
ne_shp
  [TypeBase Shape NoUniqueness]
his_ts <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
-> InternaliseM (TypeBase Shape NoUniqueness)
-> InternaliseM (TypeBase Shape NoUniqueness)
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall u. Int -> TypeBase Shape u -> TypeBase Shape u
I.stripArray (Int
dim Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) (InternaliseM (TypeBase Shape NoUniqueness)
 -> InternaliseM (TypeBase Shape NoUniqueness))
-> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> VName
-> InternaliseM (TypeBase Shape NoUniqueness)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> InternaliseM (TypeBase Shape NoUniqueness)
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 <- Int
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
-> InternaliseM [Param (TypeBase Shape NoUniqueness)]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
dim ([Char]
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"bucket_p" (TypeBase Shape NoUniqueness
 -> InternaliseM (Param (TypeBase Shape NoUniqueness)))
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
  [Param (TypeBase Shape NoUniqueness)]
img_params <- (TypeBase Shape NoUniqueness
 -> InternaliseM (Param (TypeBase Shape NoUniqueness)))
-> [TypeBase Shape NoUniqueness]
-> InternaliseM [Param (TypeBase Shape NoUniqueness)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Char]
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"img_p" (TypeBase Shape NoUniqueness
 -> InternaliseM (Param (TypeBase Shape NoUniqueness)))
-> (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall u. TypeBase Shape u -> TypeBase Shape u
rowType) ([TypeBase Shape NoUniqueness]
 -> InternaliseM [Param (TypeBase Shape NoUniqueness)])
-> InternaliseM [TypeBase Shape NoUniqueness]
-> InternaliseM [Param (TypeBase Shape NoUniqueness)]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
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 [Param (TypeBase Shape NoUniqueness)]
-> [Param (TypeBase Shape NoUniqueness)]
-> [Param (TypeBase Shape NoUniqueness)]
forall a. [a] -> [a] -> [a]
++ [Param (TypeBase Shape NoUniqueness)]
img_params
      rettype :: [TypeBase Shape NoUniqueness]
rettype = Int -> TypeBase Shape NoUniqueness -> [TypeBase Shape NoUniqueness]
forall a. Int -> a -> [a]
replicate Int
dim (PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64) [TypeBase Shape NoUniqueness]
-> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
forall a. [a] -> [a] -> [a]
++ [TypeBase Shape NoUniqueness]
ne_ts
      body :: Body SOACS
body = Stms SOACS -> Result -> Body SOACS
forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody Stms SOACS
forall a. Monoid a => a
mempty (Result -> Body SOACS) -> Result -> Body SOACS
forall a b. (a -> b) -> a -> b
$ [VName] -> Result
varsRes ([VName] -> Result) -> [VName] -> Result
forall a b. (a -> b) -> a -> b
$ (Param (TypeBase Shape NoUniqueness) -> VName)
-> [Param (TypeBase Shape NoUniqueness)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (TypeBase Shape NoUniqueness) -> VName
forall dec. Param dec -> VName
I.paramName [Param (TypeBase Shape NoUniqueness)]
params
  Lambda SOACS
lam' <-
    [LParam (Rep InternaliseM)]
-> InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
mkLambda [Param (TypeBase Shape NoUniqueness)]
[LParam (Rep InternaliseM)]
params (InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM)))
-> InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM))
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"
        (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
img)
        [TypeBase Shape NoUniqueness]
rettype
        (Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Body (Rep InternaliseM) -> InternaliseM Result
forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind Body (Rep InternaliseM)
Body SOACS
body

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

  [Char] -> Exp SOACS -> InternaliseM [SubExp]
letValExp' [Char]
desc (Exp SOACS -> InternaliseM [SubExp])
-> (SOAC SOACS -> Exp SOACS) -> SOAC SOACS -> InternaliseM [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Op SOACS -> Exp SOACS
SOAC SOACS -> Exp SOACS
forall rep. Op rep -> Exp rep
I.Op (SOAC SOACS -> InternaliseM [SubExp])
-> SOAC SOACS -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    SubExp -> [VName] -> [HistOp SOACS] -> Lambda SOACS -> SOAC SOACS
forall rep.
SubExp -> [VName] -> [HistOp rep] -> Lambda rep -> SOAC rep
I.Hist SubExp
w_img ([VName]
buckets' [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
img') [Shape
-> SubExp -> [VName] -> [SubExp] -> Lambda SOACS -> HistOp SOACS
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 <- [Char] -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"acc_cert"
  [TypeBase Shape NoUniqueness]
dest_ts <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
dest'
  let dest_w :: SubExp
dest_w = Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [TypeBase Shape NoUniqueness]
dest_ts
      acc_t :: TypeBase Shape NoUniqueness
acc_t = VName
-> Shape
-> [TypeBase Shape NoUniqueness]
-> NoUniqueness
-> TypeBase Shape NoUniqueness
forall shape u.
VName
-> Shape -> [TypeBase Shape NoUniqueness] -> u -> TypeBase shape u
Acc VName
acc_cert_v ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
I.Shape [SubExp
dest_w]) ((TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
-> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall u. TypeBase Shape u -> TypeBase Shape u
rowType [TypeBase Shape NoUniqueness]
dest_ts) NoUniqueness
NoUniqueness
  Param (TypeBase Shape NoUniqueness)
acc_p <- [Char]
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"acc_p" TypeBase Shape NoUniqueness
acc_t
  Lambda SOACS
withacc_lam <- [LParam (Rep InternaliseM)]
-> InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
mkLambda [Attrs
-> VName
-> TypeBase Shape NoUniqueness
-> Param (TypeBase Shape NoUniqueness)
forall dec. Attrs -> VName -> dec -> Param dec
Param Attrs
forall a. Monoid a => a
mempty VName
acc_cert_v (PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.Unit), Param (TypeBase Shape NoUniqueness)
LParam (Rep InternaliseM)
acc_p] (InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM)))
-> InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM))
forall a b. (a -> b) -> a -> b
$ do
    [TypeBase Shape NoUniqueness]
bs_ts <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
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 ([TypeBase Shape NoUniqueness] -> InternaliseM (Lambda SOACS))
-> [TypeBase Shape NoUniqueness] -> InternaliseM (Lambda SOACS)
forall a b. (a -> b) -> a -> b
$ (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
-> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall u. TypeBase Shape u -> TypeBase Shape u
rowType ([TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness])
-> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> a -> b
$ Param (TypeBase Shape NoUniqueness) -> TypeBase Shape NoUniqueness
forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
paramType Param (TypeBase Shape NoUniqueness)
acc_p TypeBase Shape NoUniqueness
-> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
forall a. a -> [a] -> [a]
: [TypeBase Shape NoUniqueness]
bs_ts
    let w :: SubExp
w = Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [TypeBase Shape NoUniqueness]
bs_ts
    ([SubExp] -> Result)
-> InternaliseM [SubExp] -> InternaliseM Result
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [SubExp] -> Result
subExpsRes (InternaliseM [SubExp] -> InternaliseM Result)
-> (Exp SOACS -> InternaliseM [SubExp])
-> Exp SOACS
-> InternaliseM Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp SOACS -> InternaliseM [SubExp]
letValExp' [Char]
"acc_res" (Exp SOACS -> InternaliseM Result)
-> Exp SOACS -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
      Op SOACS -> Exp SOACS
forall rep. Op rep -> Exp rep
I.Op (Op SOACS -> Exp SOACS) -> Op SOACS -> Exp SOACS
forall a b. (a -> b) -> a -> b
$
        SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w (Param (TypeBase Shape NoUniqueness) -> VName
forall dec. Param dec -> VName
I.paramName Param (TypeBase Shape NoUniqueness)
acc_p VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: [VName]
bs') (Lambda SOACS -> ScremaForm SOACS
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 <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
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 ([TypeBase Shape NoUniqueness]
 -> InternaliseM
      ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness]))
-> [TypeBase Shape NoUniqueness]
-> InternaliseM
     ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness])
forall a b. (a -> b) -> a -> b
$ [TypeBase Shape NoUniqueness]
ne_ts [TypeBase Shape NoUniqueness]
-> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
forall a. [a] -> [a] -> [a]
++ [TypeBase Shape NoUniqueness]
ne_ts
        Param (TypeBase Shape NoUniqueness)
idxp <- [Char]
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"idx" (TypeBase Shape NoUniqueness
 -> InternaliseM (Param (TypeBase Shape NoUniqueness)))
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64
        let op_lam' :: Lambda SOACS
op_lam' = [LParam SOACS]
-> [TypeBase Shape NoUniqueness] -> Body SOACS -> Lambda SOACS
forall rep.
[LParam rep]
-> [TypeBase Shape NoUniqueness] -> Body rep -> Lambda rep
I.Lambda (Param (TypeBase Shape NoUniqueness)
idxp Param (TypeBase Shape NoUniqueness)
-> [Param (TypeBase Shape NoUniqueness)]
-> [Param (TypeBase Shape NoUniqueness)]
forall a. a -> [a] -> [a]
: [Param (TypeBase Shape NoUniqueness)]
lam_params) [TypeBase Shape NoUniqueness]
lam_rettype Body SOACS
lam_body
        Maybe (Lambda SOACS, [SubExp])
-> InternaliseM (Maybe (Lambda SOACS, [SubExp]))
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Lambda SOACS, [SubExp])
 -> InternaliseM (Maybe (Lambda SOACS, [SubExp])))
-> Maybe (Lambda SOACS, [SubExp])
-> InternaliseM (Maybe (Lambda SOACS, [SubExp]))
forall a b. (a -> b) -> a -> b
$ (Lambda SOACS, [SubExp]) -> Maybe (Lambda SOACS, [SubExp])
forall a. a -> Maybe a
Just (Lambda SOACS
op_lam', [SubExp]
ne')
      Maybe (Exp, Exp)
Nothing ->
        Maybe (Lambda SOACS, [SubExp])
-> InternaliseM (Maybe (Lambda SOACS, [SubExp]))
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Lambda SOACS, [SubExp])
forall a. Maybe a
Nothing

  SubExp
destw <- Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 ([TypeBase Shape NoUniqueness] -> SubExp)
-> InternaliseM [TypeBase Shape NoUniqueness]
-> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
dest'
  ([VName] -> [SubExp])
-> InternaliseM [VName] -> InternaliseM [SubExp]
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var) (InternaliseM [VName] -> InternaliseM [SubExp])
-> InternaliseM [VName] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM [VName]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [VName]
letTupExp [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM [VName])
-> Exp (Rep InternaliseM) -> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$
      [WithAccInput SOACS] -> Lambda SOACS -> Exp SOACS
forall rep. [WithAccInput rep] -> Lambda rep -> Exp rep
WithAcc [([SubExp] -> Shape
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] -> SubExp -> InternaliseM SubExp
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
se
    [SubExp]
_ -> [Char] -> InternaliseM 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) (SubExp -> (SubExp, IntType))
-> InternaliseM SubExp -> InternaliseM (SubExp, IntType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
e'
    StructType
_ -> [Char] -> InternaliseM (SubExp, IntType)
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 =
  (SubExp -> InternaliseM VName) -> [SubExp] -> InternaliseM [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> InternaliseM VName
asIdent ([SubExp] -> InternaliseM [VName])
-> InternaliseM [SubExp] -> InternaliseM [VName]
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) = VName -> InternaliseM VName
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v
    asIdent SubExp
se = [Char] -> Exp (Rep InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM VName)
-> Exp (Rep InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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
  (VName -> InternaliseM SubExp) -> [VName] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
s (Exp SOACS -> InternaliseM SubExp)
-> (BasicOp -> Exp SOACS) -> BasicOp -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> InternaliseM SubExp)
-> (VName -> InternaliseM BasicOp) -> VName -> InternaliseM SubExp
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 <-
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"zero" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"nonzero" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
UnOp (PrimType -> UnOp
I.Neg PrimType
I.Bool) SubExp
zero
  Certs
c <- [Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert [Char]
"nonzero_cert" SubExp
nonzero ErrorMsg SubExp
"division by zero" SrcLoc
loc
  Certs -> InternaliseM a -> InternaliseM a
forall a. Certs -> InternaliseM a -> InternaliseM a
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 <-
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"nonnegative" (Exp SOACS -> InternaliseM SubExp)
-> (BasicOp -> Exp SOACS) -> BasicOp -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> InternaliseM SubExp) -> BasicOp -> InternaliseM SubExp
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
  Certs -> InternaliseM a -> InternaliseM a
forall a. Certs -> InternaliseM a -> InternaliseM a
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
_ =
  SrcLoc
-> IntType
-> SubExp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    [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
_ =
  SrcLoc
-> IntType
-> SubExp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    [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
_ =
  SrcLoc
-> IntType
-> SubExp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonnegative SrcLoc
loc IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    [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
_ =
  SrcLoc
-> IntType
-> SubExp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    [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
_ =
  SrcLoc
-> IntType
-> SubExp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    [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
_ =
  SrcLoc
-> IntType
-> SubExp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    [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
_ =
  SrcLoc
-> IntType
-> SubExp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    [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
_ =
  SrcLoc
-> IntType
-> SubExp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    [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
_ =
  SrcLoc
-> IntType
-> SubExp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    [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 (PrimType -> CmpOp) -> PrimType -> CmpOp
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp ([Char]
desc [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"true") (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimType
internalisePrimType PrimType
t) SubExp
x SubExp
y
  (SubExp -> [SubExp])
-> InternaliseM SubExp -> InternaliseM [SubExp]
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> [SubExp]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InternaliseM SubExp -> InternaliseM [SubExp])
-> InternaliseM SubExp -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (PrimType -> UnOp
I.Neg PrimType
I.Bool) 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 =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    [Char]
"Invalid binary operator "
      [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ BinOp -> [Char]
forall a. Pretty a => a -> [Char]
prettyString BinOp
op
      [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" with operand types "
      [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ PrimType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString PrimType
t1
      [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
", "
      [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ PrimType -> [Char]
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 =
  [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 =
  [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 -> [Char] -> [Char]
[Function] -> [Char] -> [Char]
Function -> [Char]
(Int -> Function -> [Char] -> [Char])
-> (Function -> [Char])
-> ([Function] -> [Char] -> [Char])
-> Show Function
forall a.
(Int -> a -> [Char] -> [Char])
-> (a -> [Char]) -> ([a] -> [Char] -> [Char]) -> Show a
$cshowsPrec :: Int -> Function -> [Char] -> [Char]
showsPrec :: Int -> Function -> [Char] -> [Char]
$cshow :: Function -> [Char]
show :: Function -> [Char]
$cshowList :: [Function] -> [Char] -> [Char]
showList :: [Function] -> [Char] -> [Char]
Show)

findFuncall :: E.AppExp -> (Function, [(E.Exp, Maybe VName)])
findFuncall :: AppExp -> (Function, [(Exp, Maybe VName)])
findFuncall (E.Apply Exp
f NonEmpty (Info (Maybe VName), Exp)
args SrcLoc
_)
  | E.Var QualName VName
fname Info StructType
_ SrcLoc
_ <- Exp
f =
      (QualName VName -> Function
FunctionName QualName VName
fname, ((Info (Maybe VName), Exp) -> (Exp, Maybe VName))
-> [(Info (Maybe VName), Exp)] -> [(Exp, Maybe VName)]
forall a b. (a -> b) -> [a] -> [b]
map (Info (Maybe VName), Exp) -> (Exp, Maybe VName)
forall {b} {a}. (Info b, a) -> (a, b)
onArg ([(Info (Maybe VName), Exp)] -> [(Exp, Maybe VName)])
-> [(Info (Maybe VName), Exp)] -> [(Exp, Maybe VName)]
forall a b. (a -> b) -> a -> b
$ NonEmpty (Info (Maybe VName), Exp) -> [(Info (Maybe VName), Exp)]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Info (Maybe VName), Exp)
args)
  | E.Hole (Info StructType
_) SrcLoc
loc <- Exp
f =
      (SrcLoc -> Function
FunctionHole SrcLoc
loc, ((Info (Maybe VName), Exp) -> (Exp, Maybe VName))
-> [(Info (Maybe VName), Exp)] -> [(Exp, Maybe VName)]
forall a b. (a -> b) -> [a] -> [b]
map (Info (Maybe VName), Exp) -> (Exp, Maybe VName)
forall {b} {a}. (Info b, a) -> (a, b)
onArg ([(Info (Maybe VName), Exp)] -> [(Exp, Maybe VName)])
-> [(Info (Maybe VName), Exp)] -> [(Exp, Maybe VName)]
forall a b. (a -> b) -> a -> b
$ NonEmpty (Info (Maybe VName), Exp) -> [(Info (Maybe VName), Exp)]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Info (Maybe VName), Exp)
args)
  where
    onArg :: (Info b, a) -> (a, b)
onArg (Info b
argext, a
e) = (a
e, b
argext)
findFuncall AppExp
e =
  [Char] -> (Function, [(Exp, Maybe VName)])
forall a. HasCallStack => [Char] -> a
error ([Char] -> (Function, [(Exp, Maybe VName)]))
-> [Char] -> (Function, [(Exp, Maybe VName)])
forall a b. (a -> b) -> a -> b
$ [Char]
"Invalid function expression in application:\n" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ AppExp -> [Char]
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 (Scope SOACS -> [VName]
forall k a. Map k a -> [k]
M.keys Scope SOACS
stmsscope) ([ExtType] -> [ExtType])
-> ([TypeBase Shape NoUniqueness] -> [ExtType])
-> [TypeBase Shape NoUniqueness]
-> [ExtType]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TypeBase Shape NoUniqueness] -> [ExtType]
forall u. [TypeBase Shape u] -> [TypeBase ExtShape u]
staticShapes
    ([TypeBase Shape NoUniqueness] -> [ExtType])
-> InternaliseM [TypeBase Shape NoUniqueness]
-> InternaliseM [ExtType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExtendedScope SOACS InternaliseM [TypeBase Shape NoUniqueness]
-> Scope SOACS -> InternaliseM [TypeBase Shape NoUniqueness]
forall rep (m :: * -> *) a.
ExtendedScope rep m a -> Scope rep -> m a
extendedScope ((SubExpRes
 -> ExtendedScope SOACS InternaliseM (TypeBase Shape NoUniqueness))
-> Result
-> ExtendedScope SOACS InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse SubExpRes
-> ExtendedScope SOACS InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExpRes -> m (TypeBase Shape NoUniqueness)
subExpResType Result
res) Scope SOACS
stmsscope
  where
    stmsscope :: Scope SOACS
stmsscope = Stms SOACS -> Scope SOACS
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 Exp VName)
_ (Info (RetType [VName]
_ TypeBase Exp Uniqueness
rettype)) SrcLoc
_) [TypeBase Shape NoUniqueness]
rowtypes =
  [PatBase Info VName ParamType]
-> [TypeBase Shape NoUniqueness]
-> ([LParam SOACS]
    -> InternaliseM
         ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness]))
-> InternaliseM
     ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness])
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 (([LParam SOACS]
  -> InternaliseM
       ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness]))
 -> InternaliseM
      ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness]))
-> ([LParam SOACS]
    -> InternaliseM
         ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness]))
-> InternaliseM
     ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness])
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' <- TypeBase Exp Uniqueness
-> [ExtType] -> InternaliseM [TypeBase Shape NoUniqueness]
forall shape u.
TypeBase Exp Uniqueness
-> [TypeBase shape u] -> InternaliseM [TypeBase Shape NoUniqueness]
internaliseLambdaReturnType TypeBase Exp Uniqueness
rettype ([ExtType] -> InternaliseM [TypeBase Shape NoUniqueness])
-> InternaliseM [ExtType]
-> InternaliseM [TypeBase Shape NoUniqueness]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Body SOACS -> InternaliseM [ExtType]
bodyExtType Body SOACS
body'
    ([Param (TypeBase Shape NoUniqueness)], Body SOACS,
 [TypeBase Shape NoUniqueness])
-> InternaliseM
     ([Param (TypeBase Shape NoUniqueness)], Body SOACS,
      [TypeBase Shape NoUniqueness])
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Param (TypeBase Shape NoUniqueness)]
[LParam SOACS]
params', Body SOACS
body', [TypeBase Shape NoUniqueness]
rettype')
internaliseLambda Exp
e [TypeBase Shape NoUniqueness]
_ = [Char]
-> InternaliseM
     ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness])
forall a. HasCallStack => [Char] -> a
error ([Char]
 -> InternaliseM
      ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness]))
-> [Char]
-> InternaliseM
     ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness])
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseLambda: unexpected expression:\n" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Exp -> [Char]
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
  [LParam (Rep InternaliseM)]
-> InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
mkLambda [Param (TypeBase Shape NoUniqueness)]
[LParam (Rep InternaliseM)]
params (InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM)))
-> InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM))
forall a b. (a -> b) -> a -> b
$
    ErrorMsg SubExp
-> SrcLoc
-> [TypeBase Shape NoUniqueness]
-> Result
-> InternaliseM Result
ensureResultShape
      ([ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg [Text -> ErrorMsgPart SubExp
forall a. Text -> ErrorMsgPart a
ErrorString Text
"unexpected lambda result size"])
      (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
lam)
      [TypeBase Shape NoUniqueness]
rettype
      (Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Body (Rep InternaliseM) -> InternaliseM Result
forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind Body (Rep InternaliseM)
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
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ VName -> Int
baseTag (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qname) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag
  [Char] -> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
handle ([Char]
 -> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp]))
-> [Char]
-> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ VName -> [Char]
baseString (VName -> [Char]) -> VName -> [Char]
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
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 = ([(StructType, [SubExp])] -> InternaliseM [SubExp])
-> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([(StructType, [SubExp])] -> InternaliseM [SubExp])
 -> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp]))
-> ([(StructType, [SubExp])] -> InternaliseM [SubExp])
-> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[(StructType
_, [SubExp]
xe'), (StructType
_, [SubExp]
ye')] -> do
          [SubExp]
rs <- (SubExp -> SubExp -> InternaliseM SubExp)
-> [SubExp] -> [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM SubExp -> SubExp -> InternaliseM SubExp
doComparison [SubExp]
xe' [SubExp]
ye'
          SubExp -> InternaliseM [SubExp]
cmp_f (SubExp -> InternaliseM [SubExp])
-> InternaliseM SubExp -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"eq" (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [SubExp] -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
eAll [SubExp]
rs
      where
        isEqlOp :: [Char] -> Maybe (SubExp -> InternaliseM [SubExp])
isEqlOp [Char]
"!=" = (SubExp -> InternaliseM [SubExp])
-> Maybe (SubExp -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((SubExp -> InternaliseM [SubExp])
 -> Maybe (SubExp -> InternaliseM [SubExp]))
-> (SubExp -> InternaliseM [SubExp])
-> Maybe (SubExp -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \SubExp
eq ->
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (PrimType -> UnOp
I.Neg PrimType
I.Bool) SubExp
eq
        isEqlOp [Char]
"==" = (SubExp -> InternaliseM [SubExp])
-> Maybe (SubExp -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((SubExp -> InternaliseM [SubExp])
 -> Maybe (SubExp -> InternaliseM [SubExp]))
-> (SubExp -> InternaliseM [SubExp])
-> Maybe (SubExp -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \SubExp
eq ->
          [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp
eq]
        isEqlOp [Char]
_ = Maybe (SubExp -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

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

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

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

              [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"arrays_equal"
                (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
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 (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
shapes_match) InternaliseM (Body (Rep InternaliseM))
InternaliseM (Body SOACS)
compare_elems_body ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
False])
    handle [Char]
name
      | Just BinOp
bop <- (BinOp -> Bool) -> [BinOp] -> Maybe BinOp
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (([Char]
name ==) ([Char] -> Bool) -> (BinOp -> [Char]) -> BinOp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BinOp -> [Char]
forall a. Pretty a => a -> [Char]
prettyString) [BinOp
forall a. Bounded a => a
minBound .. BinOp
forall a. Bounded a => a
maxBound :: E.BinOp] =
          ([(StructType, [SubExp])] -> InternaliseM [SubExp])
-> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([(StructType, [SubExp])] -> InternaliseM [SubExp])
 -> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp]))
-> ([(StructType, [SubExp])] -> InternaliseM [SubExp])
-> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
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)
_ -> [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise.internaliseExp: non-primitive type in BinOp."
    handle [Char]
_ = Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
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
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ VName -> Int
baseTag (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qname) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag
  let handlers :: [[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])]
handlers =
        [ [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
forall {a}.
(Eq a, IsString a) =>
[Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleSign,
          [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
forall {f :: * -> *}.
Applicative f =>
[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM (f SubExp))
handleOps,
          [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
handleSOACs,
          [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
forall {a}.
(Eq a, IsString a) =>
[Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleAccs,
          [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
forall {a}.
(IsString a, Eq a) =>
[Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleAD,
          [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
handleRest
        ]
  [Maybe ([Char] -> InternaliseM [SubExp])]
-> Maybe ([Char] -> InternaliseM [SubExp])
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, MonadPlus m) =>
t (m a) -> m a
msum [[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
h [Exp]
args ([Char] -> Maybe ([Char] -> InternaliseM [SubExp]))
-> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ VName -> [Char]
baseString (VName -> [Char]) -> VName -> [Char]
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
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" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
I.Int8 Exp
x
    handleSign [Exp
x] a
"sign_i16" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
I.Int16 Exp
x
    handleSign [Exp
x] a
"sign_i32" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
I.Int32 Exp
x
    handleSign [Exp
x] a
"sign_i64" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
I.Int64 Exp
x
    handleSign [Exp
x] a
"unsign_i8" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
I.Int8 Exp
x
    handleSign [Exp
x] a
"unsign_i16" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
I.Int16 Exp
x
    handleSign [Exp
x] a
"unsign_i32" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
I.Int32 Exp
x
    handleSign [Exp
x] a
"unsign_i64" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
I.Int64 Exp
x
    handleSign [Exp]
_ a
_ = Maybe ([Char] -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

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

    handleSOACs :: [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
handleSOACs [Exp
lam, Exp
arr] [Char]
"map" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
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 <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
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 ([TypeBase Shape NoUniqueness] -> InternaliseM (Lambda SOACS))
-> [TypeBase Shape NoUniqueness] -> InternaliseM (Lambda SOACS)
forall a b. (a -> b) -> a -> b
$ (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
-> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall u. TypeBase Shape u -> TypeBase Shape u
rowType [TypeBase Shape NoUniqueness]
arr_ts
      let w :: SubExp
w = Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [TypeBase Shape NoUniqueness]
arr_ts
      [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Op (Rep InternaliseM) -> Exp (Rep InternaliseM)
forall rep. Op rep -> Exp rep
I.Op (Op (Rep InternaliseM) -> Exp (Rep InternaliseM))
-> Op (Rep InternaliseM) -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w [VName]
arr' (Lambda SOACS -> ScremaForm SOACS
forall rep. Lambda rep -> ScremaForm rep
I.mapSOAC Lambda SOACS
lam')
    handleSOACs [Exp
k, Exp
lam, Exp
arr] [Char]
"partition" = do
      Int
k' <- Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int32 -> Int) -> Maybe Int32 -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> Maybe Int32
forall {vn}. ExpBase Info vn -> Maybe Int32
fromInt32 Exp
k
      ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
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 ([SubExp] -> InternaliseM (Lambda SOACS))
-> [SubExp] -> InternaliseM (Lambda SOACS)
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var [VName]
arrs
        ([SubExp] -> [SubExp] -> [SubExp])
-> ([SubExp], [SubExp]) -> [SubExp]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
(++) (([SubExp], [SubExp]) -> [SubExp])
-> InternaliseM ([SubExp], [SubExp]) -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Lambda SOACS -> [VName] -> InternaliseM ([SubExp], [SubExp])
partitionWithSOACS (Int -> Int
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
_) = Int32 -> Maybe Int32
forall a. a -> Maybe a
Just Int32
k'
        fromInt32 (IntLit Integer
k' (Info (E.Scalar (E.Prim (E.Signed IntType
Int32)))) SrcLoc
_) = Int32 -> Maybe Int32
forall a. a -> Maybe a
Just (Int32 -> Maybe Int32) -> Int32 -> Maybe Int32
forall a b. (a -> b) -> a -> b
$ Integer -> Int32
forall a. Num a => Integer -> a
fromInteger Integer
k'
        fromInt32 ExpBase Info vn
_ = Maybe Int32
forall a. Maybe a
Nothing
    handleSOACs [Exp
lam, Exp
ne, Exp
arr] [Char]
"reduce" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
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" SubExp
-> Lambda SOACS -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)
forall {rep} {f :: * -> *}.
(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 =
          SubExp -> [VName] -> ScremaForm rep -> SOAC rep
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w [VName]
arrs
            (ScremaForm rep -> SOAC rep) -> f (ScremaForm rep) -> f (SOAC rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Reduce rep] -> f (ScremaForm rep)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Reduce rep] -> m (ScremaForm rep)
I.reduceSOAC [Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
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" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
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" SubExp
-> Lambda SOACS -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)
forall {rep} {f :: * -> *}.
(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 =
          SubExp -> [VName] -> ScremaForm rep -> SOAC rep
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w [VName]
arrs
            (ScremaForm rep -> SOAC rep) -> f (ScremaForm rep) -> f (SOAC rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Reduce rep] -> f (ScremaForm rep)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Reduce rep] -> m (ScremaForm rep)
I.reduceSOAC [Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
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" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
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" SubExp
-> Lambda SOACS -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)
forall {rep} {f :: * -> *}.
(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 =
          SubExp -> [VName] -> ScremaForm rep -> SOAC rep
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w [VName]
arrs (ScremaForm rep -> SOAC rep) -> f (ScremaForm rep) -> f (SOAC rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Scan rep] -> f (ScremaForm rep)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Scan rep] -> m (ScremaForm rep)
I.scanSOAC [Lambda rep -> [SubExp] -> Scan rep
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" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
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" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
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" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
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]
_ = Maybe ([Char] -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

    handleAccs :: [Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleAccs [Exp
dest, Exp
f, Exp
bs] a
"scatter_stream" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char]
-> Exp -> Maybe (Exp, Exp) -> Exp -> Exp -> InternaliseM [SubExp]
internaliseStreamAcc [Char]
desc Exp
dest Maybe (Exp, Exp)
forall a. Maybe a
Nothing Exp
f Exp
bs
    handleAccs [Exp
dest, Exp
op, Exp
ne, Exp
f, Exp
bs] a
"hist_stream" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char]
-> Exp -> Maybe (Exp, Exp) -> Exp -> Exp -> InternaliseM [SubExp]
internaliseStreamAcc [Char]
desc Exp
dest ((Exp, Exp) -> Maybe (Exp, Exp)
forall a. a -> Maybe a
Just (Exp
op, Exp
ne)) Exp
f Exp
bs
    handleAccs [Exp
acc, Exp
i, Exp
v] a
"acc_write" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      VName
acc' <- [VName] -> VName
forall a. HasCallStack => [a] -> a
head ([VName] -> VName) -> InternaliseM [VName] -> InternaliseM VName
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
      (SubExp -> [SubExp])
-> InternaliseM SubExp -> InternaliseM [SubExp]
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> [SubExp]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InternaliseM SubExp -> InternaliseM [SubExp])
-> InternaliseM SubExp -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ Safety -> VName -> [SubExp] -> [SubExp] -> BasicOp
UpdateAcc Safety
Safe VName
acc' [SubExp
i'] [SubExp]
vs
    handleAccs [Exp]
_ a
_ = Maybe ([Char] -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

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

    handleRest :: [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
handleRest [Exp
a, Exp
si, Exp
v] [Char]
"scatter" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
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" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
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" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
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" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
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 <- Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
I.arraysSize Int
0 ([TypeBase Shape NoUniqueness] -> SubExp)
-> InternaliseM [TypeBase Shape NoUniqueness]
-> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
      SubExp
dim_ok <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"dim_ok" (Exp SOACS -> InternaliseM SubExp)
-> (TPrimExp Bool VName -> InternaliseM (Exp SOACS))
-> TPrimExp Bool VName
-> InternaliseM SubExp
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< TPrimExp Bool VName -> InternaliseM (Exp (Rep InternaliseM))
TPrimExp Bool VName -> InternaliseM (Exp SOACS)
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
TPrimExp Bool VName -> m (Exp (Rep m))
toExp (TPrimExp Bool VName -> InternaliseM SubExp)
-> TPrimExp Bool VName -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          SubExp -> TPrimExp Int64 VName
pe64 SubExp
old_dim TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.==. SubExp -> TPrimExp Int64 VName
pe64 SubExp
n'
            TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* SubExp -> TPrimExp Int64 VName
pe64 SubExp
m'
              TPrimExp Bool VName -> TPrimExp Bool VName -> TPrimExp Bool VName
forall v.
Eq v =>
TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v
.&&. SubExp -> TPrimExp Int64 VName
pe64 SubExp
n'
              TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.>=. TPrimExp Int64 VName
0
              TPrimExp Bool VName -> TPrimExp Bool VName -> TPrimExp Bool VName
forall v.
Eq v =>
TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v
.&&. SubExp -> TPrimExp Int64 VName
pe64 SubExp
m'
              TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
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
          ( [ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg
              [ ErrorMsgPart SubExp
"Cannot unflatten array of shape [",
                PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
old_dim,
                ErrorMsgPart SubExp
"] to array of shape [",
                PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
n',
                ErrorMsgPart SubExp
"][",
                PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
m',
                ErrorMsgPart SubExp
"]"
              ]
          )
          SrcLoc
loc
      Certs -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a. Certs -> InternaliseM a -> InternaliseM a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
dim_ok_cert (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
        [VName] -> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
arrs ((VName -> InternaliseM SubExp) -> InternaliseM [SubExp])
-> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
arr' -> do
          TypeBase Shape NoUniqueness
arr_t <- VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
arr'
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc (Exp SOACS -> InternaliseM SubExp)
-> (BasicOp -> Exp SOACS) -> BasicOp -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> InternaliseM SubExp) -> BasicOp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
            ReshapeKind -> Shape -> VName -> BasicOp
I.Reshape
              ReshapeKind
I.ReshapeArbitrary
              (Shape -> Int -> Shape -> Shape
reshapeOuter ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
I.Shape [SubExp
n', SubExp
m']) Int
1 (Shape -> Shape) -> Shape -> Shape
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoUniqueness -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape TypeBase Shape NoUniqueness
arr_t)
              VName
arr'
    handleRest [Exp
arr] [Char]
"manifest" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      [VName]
arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"flatten_arr" Exp
arr
      [VName] -> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
arrs ((VName -> InternaliseM SubExp) -> InternaliseM [SubExp])
-> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
arr' -> do
        Int
r <- TypeBase Shape NoUniqueness -> Int
forall shape u. ArrayShape shape => TypeBase shape u -> Int
I.arrayRank (TypeBase Shape NoUniqueness -> Int)
-> InternaliseM (TypeBase Shape NoUniqueness) -> InternaliseM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
arr'
        if Int
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
          then SubExp -> InternaliseM SubExp
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> InternaliseM SubExp) -> SubExp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
arr'
          else [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ [Int] -> VName -> BasicOp
I.Manifest [Int
0 .. Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] VName
arr'
    handleRest [Exp
arr] [Char]
"flatten" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      [VName]
arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"flatten_arr" Exp
arr
      [VName] -> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
arrs ((VName -> InternaliseM SubExp) -> InternaliseM [SubExp])
-> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
arr' -> do
        TypeBase Shape NoUniqueness
arr_t <- VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
arr'
        let n :: SubExp
n = Int -> TypeBase Shape NoUniqueness -> SubExp
forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
0 TypeBase Shape NoUniqueness
arr_t
            m :: SubExp
m = Int -> TypeBase Shape NoUniqueness -> SubExp
forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
1 TypeBase Shape NoUniqueness
arr_t
        SubExp
k <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"flat_dim" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc (Exp SOACS -> InternaliseM SubExp)
-> (BasicOp -> Exp SOACS) -> BasicOp -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> InternaliseM SubExp) -> BasicOp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          ReshapeKind -> Shape -> VName -> BasicOp
I.Reshape
            ReshapeKind
I.ReshapeArbitrary
            (Shape -> Int -> Shape -> Shape
reshapeOuter ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
I.Shape [SubExp
k]) Int
2 (Shape -> Shape) -> Shape -> Shape
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoUniqueness -> Shape
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" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
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 <- Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 ([TypeBase Shape NoUniqueness] -> SubExp)
-> InternaliseM [TypeBase Shape NoUniqueness]
-> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
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 =
            [Char] -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"conc_tmp" (Exp (Rep m) -> m SubExp) -> Exp (Rep m) -> m SubExp
forall a b. (a -> b) -> a -> b
$
              BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
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 <-
        (SubExp -> SubExp -> InternaliseM SubExp)
-> SubExp -> [SubExp] -> InternaliseM SubExp
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM SubExp -> SubExp -> InternaliseM SubExp
forall {m :: * -> *}.
MonadBuilder m =>
SubExp -> SubExp -> m SubExp
sumdims SubExp
outer_size
          ([SubExp] -> InternaliseM SubExp)
-> InternaliseM [SubExp] -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ([VName] -> InternaliseM SubExp)
-> [[VName]] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (([TypeBase Shape NoUniqueness] -> SubExp)
-> InternaliseM [TypeBase Shape NoUniqueness]
-> InternaliseM SubExp
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0) (InternaliseM [TypeBase Shape NoUniqueness] -> InternaliseM SubExp)
-> ([VName] -> InternaliseM [TypeBase Shape NoUniqueness])
-> [VName]
-> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
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 =
            BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ Int -> NonEmpty VName -> SubExp -> BasicOp
I.Concat Int
0 (VName
xarr VName -> [VName] -> NonEmpty VName
forall a. a -> [a] -> NonEmpty a
:| [VName
yarr]) SubExp
ressize
      (Exp SOACS -> InternaliseM SubExp)
-> [Exp SOACS] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc) ([Exp SOACS] -> InternaliseM [SubExp])
-> [Exp SOACS] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ (VName -> VName -> Exp SOACS) -> [VName] -> [VName] -> [Exp SOACS]
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" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char]
-> Exp -> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
internaliseOperation [Char]
desc Exp
e ((VName -> InternaliseM BasicOp) -> InternaliseM [SubExp])
-> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
v -> do
        Int
r <- TypeBase Shape NoUniqueness -> Int
forall shape u. ArrayShape shape => TypeBase shape u -> Int
I.arrayRank (TypeBase Shape NoUniqueness -> Int)
-> InternaliseM (TypeBase Shape NoUniqueness) -> InternaliseM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
v
        BasicOp -> InternaliseM BasicOp
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (BasicOp -> InternaliseM BasicOp)
-> BasicOp -> InternaliseM BasicOp
forall a b. (a -> b) -> a -> b
$ [Int] -> VName -> BasicOp
I.Rearrange ([Int
1, Int
0] [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
2 .. Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]) VName
v
    handleRest [Exp
x, Exp
y] [Char]
"zip" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      (VName -> InternaliseM SubExp) -> [VName] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"zip_copy" (Exp SOACS -> InternaliseM SubExp)
-> (VName -> Exp SOACS) -> VName -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> (VName -> BasicOp) -> VName -> Exp SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shape -> SubExp -> BasicOp
Replicate Shape
forall a. Monoid a => a
mempty (SubExp -> BasicOp) -> (VName -> SubExp) -> VName -> BasicOp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
I.Var)
        ([VName] -> InternaliseM [SubExp])
-> InternaliseM [VName] -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ( [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
(++)
                ([VName] -> [VName] -> [VName])
-> InternaliseM [VName] -> InternaliseM ([VName] -> [VName])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars ([Char]
desc [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"_zip_x") Exp
x
                InternaliseM ([VName] -> [VName])
-> InternaliseM [VName] -> InternaliseM [VName]
forall a b.
InternaliseM (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars ([Char]
desc [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"_zip_y") Exp
y
            )
    handleRest [Exp
x] [Char]
"unzip" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      (VName -> InternaliseM SubExp) -> [VName] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc (Exp SOACS -> InternaliseM SubExp)
-> (VName -> Exp SOACS) -> VName -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> (VName -> BasicOp) -> VName -> Exp SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shape -> SubExp -> BasicOp
Replicate Shape
forall a. Monoid a => a
mempty (SubExp -> BasicOp) -> (VName -> SubExp) -> VName -> BasicOp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
I.Var)
        ([VName] -> InternaliseM [SubExp])
-> InternaliseM [VName] -> InternaliseM [SubExp]
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" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
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" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
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" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
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" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
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" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
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" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
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]
_ = Maybe ([Char] -> InternaliseM [SubExp])
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) ->
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc
            (Exp SOACS -> InternaliseM [SubExp])
-> InternaliseM (Exp SOACS) -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
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
              (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
e')
              ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
1])
              ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
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)) ->
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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)) ->
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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)) ->
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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
_ -> [Char] -> InternaliseM [SubExp]
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) ->
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc
            (Exp SOACS -> InternaliseM [SubExp])
-> InternaliseM (Exp SOACS) -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
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
              (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
e')
              ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
1])
              ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
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)) ->
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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)) ->
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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)) ->
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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
_ -> [Char] -> InternaliseM [SubExp]
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 <- Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
I.arraysSize Int
0 ([TypeBase Shape NoUniqueness] -> SubExp)
-> InternaliseM [TypeBase Shape NoUniqueness]
-> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
si'
      [TypeBase Shape NoUniqueness]
sv_ts <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
svs

      [VName]
svs' <- [(VName, TypeBase Shape NoUniqueness)]
-> ((VName, TypeBase Shape NoUniqueness) -> InternaliseM VName)
-> InternaliseM [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([VName]
-> [TypeBase Shape NoUniqueness]
-> [(VName, TypeBase Shape NoUniqueness)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
svs [TypeBase Shape NoUniqueness]
sv_ts) (((VName, TypeBase Shape NoUniqueness) -> InternaliseM VName)
 -> InternaliseM [VName])
-> ((VName, TypeBase Shape NoUniqueness) -> InternaliseM VName)
-> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$ \(VName
sv, TypeBase Shape NoUniqueness
sv_t) -> do
        let sv_shape :: Shape
sv_shape = TypeBase Shape NoUniqueness -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape TypeBase Shape NoUniqueness
sv_t
            sv_w :: SubExp
sv_w = Int -> TypeBase Shape NoUniqueness -> SubExp
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 <-
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"write_cmp" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
            BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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
        Certs -> InternaliseM VName -> InternaliseM VName
forall a. Certs -> InternaliseM a -> InternaliseM a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c (InternaliseM VName -> InternaliseM VName)
-> InternaliseM VName -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp (VName -> [Char]
baseString VName
sv [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"_write_sv") (Exp SOACS -> InternaliseM VName)
-> (BasicOp -> Exp SOACS) -> BasicOp -> InternaliseM VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> InternaliseM VName) -> BasicOp -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$
            ReshapeKind -> Shape -> VName -> BasicOp
I.Reshape ReshapeKind
I.ReshapeCoerce (Shape -> Int -> Shape -> Shape
reshapeOuter ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
I.Shape [SubExp
si_w]) Int
1 Shape
sv_shape) VName
sv

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

      [TypeBase Shape NoUniqueness]
sa_ts <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
sas
      let bodyTypes :: [TypeBase Shape NoUniqueness]
bodyTypes = [[TypeBase Shape NoUniqueness]] -> [TypeBase Shape NoUniqueness]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Int
-> [TypeBase Shape NoUniqueness] -> [[TypeBase Shape NoUniqueness]]
forall a. Int -> a -> [a]
replicate ([TypeBase Shape NoUniqueness] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase Shape NoUniqueness]
sv_ts) [TypeBase Shape NoUniqueness]
indexType) [TypeBase Shape NoUniqueness]
-> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
forall a. [a] -> [a] -> [a]
++ (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
-> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
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 [TypeBase Shape NoUniqueness]
-> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
forall a. Semigroup a => a -> a -> a
<> (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
-> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall u. TypeBase Shape u -> TypeBase Shape u
rowType [TypeBase Shape NoUniqueness]
sv_ts
          bodyNames :: [VName]
bodyNames = [VName]
indexName [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
valueNames
          bodyParams :: [Param (TypeBase Shape NoUniqueness)]
bodyParams = (VName
 -> TypeBase Shape NoUniqueness
 -> Param (TypeBase Shape NoUniqueness))
-> [VName]
-> [TypeBase Shape NoUniqueness]
-> [Param (TypeBase Shape NoUniqueness)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (Attrs
-> VName
-> TypeBase Shape NoUniqueness
-> Param (TypeBase Shape NoUniqueness)
forall dec. Attrs -> VName -> dec -> Param dec
I.Param Attrs
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 <- Scope SOACS
-> InternaliseM (Body SOACS) -> InternaliseM (Body SOACS)
forall a. Scope SOACS -> InternaliseM a -> InternaliseM a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope ([Param (TypeBase Shape NoUniqueness)] -> Scope SOACS
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams [Param (TypeBase Shape NoUniqueness)]
bodyParams) (InternaliseM (Body SOACS) -> InternaliseM (Body SOACS))
-> (InternaliseM Result -> InternaliseM (Body SOACS))
-> InternaliseM Result
-> InternaliseM (Body SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InternaliseM Result -> InternaliseM (Body (Rep InternaliseM))
InternaliseM Result -> InternaliseM (Body SOACS)
forall (m :: * -> *).
MonadBuilder m =>
m Result -> m (Body (Rep m))
buildBody_ (InternaliseM Result -> InternaliseM (Body SOACS))
-> InternaliseM Result -> InternaliseM (Body SOACS)
forall a b. (a -> b) -> a -> b
$ do
        let outs :: [VName]
outs = [[VName]] -> [VName]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Int -> [VName] -> [[VName]]
forall a. Int -> a -> [a]
replicate ([VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
valueNames) [VName]
indexName) [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
valueNames
        [SubExp]
results <- [VName] -> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
outs ((VName -> InternaliseM SubExp) -> InternaliseM [SubExp])
-> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
name ->
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"write_res" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
name
        ErrorMsg SubExp
-> SrcLoc
-> [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)]
[LParam SOACS]
bodyParams,
                lambdaReturnType :: [TypeBase Shape NoUniqueness]
I.lambdaReturnType = [TypeBase Shape NoUniqueness]
bodyTypes,
                lambdaBody :: Body SOACS
I.lambdaBody = Body SOACS
body
              }
          sivs :: [VName]
sivs = [VName]
si' [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
svs'

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

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 <- Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
I.arraysSize Int
0 ([TypeBase Shape NoUniqueness] -> SubExp)
-> InternaliseM [TypeBase Shape NoUniqueness]
-> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
  SubExp
offset_inbounds_down <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"offset_inbounds_down" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"offset_inbounds_up" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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' <-
    ((Exp, Exp) -> InternaliseM (SubExp, SubExp))
-> [(Exp, Exp)] -> InternaliseM [(SubExp, SubExp)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [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
          (SubExp, SubExp) -> InternaliseM (SubExp, SubExp)
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
n', SubExp
s')
      )
      [(Exp, Exp)]
slices
  (SubExp
min_bound, SubExp
max_bound) <-
    ((SubExp, SubExp)
 -> (SubExp, SubExp) -> InternaliseM (SubExp, SubExp))
-> (SubExp, SubExp)
-> [(SubExp, SubExp)]
-> InternaliseM (SubExp, SubExp)
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span_and_lower" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span_and_upper" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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' <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"minimum" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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' <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"maximum" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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

          (SubExp, SubExp) -> InternaliseM (SubExp, SubExp)
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
lower', SubExp
upper')
      )
      (SubExp
offset', SubExp
offset')
      [(SubExp, SubExp)]
slices'
  SubExp
min_in_bounds <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"min_in_bounds" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"max_in_bounds" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <-
    (SubExp -> SubExp -> InternaliseM SubExp)
-> SubExp -> [SubExp] -> InternaliseM SubExp
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
      (\SubExp
x SubExp
y -> [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"inBounds" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 ([ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg [Text -> ErrorMsgPart SubExp
forall a. Text -> ErrorMsgPart a
ErrorString (Text -> ErrorMsgPart SubExp) -> Text -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ Text
"Flat slice out of bounds: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> SubExp -> Text
forall a. Pretty a => a -> Text
prettyText SubExp
old_dim Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" and " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [(SubExp, SubExp)] -> Text
forall a. Pretty a => a -> Text
prettyText [(SubExp, SubExp)]
slices']) SrcLoc
loc
  let slice :: FlatSlice SubExp
slice = SubExp -> [FlatDimIndex SubExp] -> FlatSlice SubExp
forall d. d -> [FlatDimIndex d] -> FlatSlice d
I.FlatSlice SubExp
offset' ([FlatDimIndex SubExp] -> FlatSlice SubExp)
-> [FlatDimIndex SubExp] -> FlatSlice SubExp
forall a b. (a -> b) -> a -> b
$ ((SubExp, SubExp) -> FlatDimIndex SubExp)
-> [(SubExp, SubExp)] -> [FlatDimIndex SubExp]
forall a b. (a -> b) -> [a] -> [b]
map ((SubExp -> SubExp -> FlatDimIndex SubExp)
-> (SubExp, SubExp) -> FlatDimIndex SubExp
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry SubExp -> SubExp -> FlatDimIndex SubExp
forall d. d -> d -> FlatDimIndex d
FlatDimIndex) [(SubExp, SubExp)]
slices'
  Certs -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a. Certs -> InternaliseM a -> InternaliseM a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    [VName] -> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
arrs ((VName -> InternaliseM SubExp) -> InternaliseM [SubExp])
-> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
arr' ->
      [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
I.arraysSize Int
0 ([TypeBase Shape NoUniqueness] -> SubExp)
-> InternaliseM [TypeBase Shape NoUniqueness]
-> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs1
  SubExp
offset_inbounds_down <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"offset_inbounds_down" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"offset_inbounds_up" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs2
  [(SubExp, SubExp)]
slices' <-
    ((Exp, Int) -> InternaliseM (SubExp, SubExp))
-> [(Exp, Int)] -> InternaliseM [(SubExp, SubExp)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM
      ( \(Exp
s, Int
i) -> do
          SubExp
s' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"s" Exp
s
          let n :: SubExp
n = Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
i [TypeBase Shape NoUniqueness]
ts
          (SubExp, SubExp) -> InternaliseM (SubExp, SubExp)
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
n, SubExp
s')
      )
      ([(Exp, Int)] -> InternaliseM [(SubExp, SubExp)])
-> [(Exp, Int)] -> InternaliseM [(SubExp, SubExp)]
forall a b. (a -> b) -> a -> b
$ [Exp] -> [Int] -> [(Exp, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Exp]
slices [Int
0 ..]
  (SubExp
min_bound, SubExp
max_bound) <-
    ((SubExp, SubExp)
 -> (SubExp, SubExp) -> InternaliseM (SubExp, SubExp))
-> (SubExp, SubExp)
-> [(SubExp, SubExp)]
-> InternaliseM (SubExp, SubExp)
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span_and_lower" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span_and_upper" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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' <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"minimum" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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' <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"maximum" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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

          (SubExp, SubExp) -> InternaliseM (SubExp, SubExp)
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
lower', SubExp
upper')
      )
      (SubExp
offset', SubExp
offset')
      [(SubExp, SubExp)]
slices'
  SubExp
min_in_bounds <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"min_in_bounds" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"max_in_bounds" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <-
    (SubExp -> SubExp -> InternaliseM SubExp)
-> SubExp -> [SubExp] -> InternaliseM SubExp
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
      (\SubExp
x SubExp
y -> [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"inBounds" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 ([ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg [Text -> ErrorMsgPart SubExp
forall a. Text -> ErrorMsgPart a
ErrorString (Text -> ErrorMsgPart SubExp) -> Text -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ Text
"Flat slice out of bounds: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> SubExp -> Text
forall a. Pretty a => a -> Text
prettyText SubExp
old_dim Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" and " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [(SubExp, SubExp)] -> Text
forall a. Pretty a => a -> Text
prettyText [(SubExp, SubExp)]
slices']) SrcLoc
loc
  let slice :: FlatSlice SubExp
slice = SubExp -> [FlatDimIndex SubExp] -> FlatSlice SubExp
forall d. d -> [FlatDimIndex d] -> FlatSlice d
I.FlatSlice SubExp
offset' ([FlatDimIndex SubExp] -> FlatSlice SubExp)
-> [FlatDimIndex SubExp] -> FlatSlice SubExp
forall a b. (a -> b) -> a -> b
$ ((SubExp, SubExp) -> FlatDimIndex SubExp)
-> [(SubExp, SubExp)] -> [FlatDimIndex SubExp]
forall a b. (a -> b) -> [a] -> [b]
map ((SubExp -> SubExp -> FlatDimIndex SubExp)
-> (SubExp, SubExp) -> FlatDimIndex SubExp
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry SubExp -> SubExp -> FlatDimIndex SubExp
forall d. d -> d -> FlatDimIndex d
FlatDimIndex) [(SubExp, SubExp)]
slices'
  Certs -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a. Certs -> InternaliseM a -> InternaliseM a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    [(VName, VName)]
-> ((VName, VName) -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([VName] -> [VName] -> [(VName, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
arrs1 [VName]
arrs2) (((VName, VName) -> InternaliseM SubExp) -> InternaliseM [SubExp])
-> ((VName, VName) -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \(VName
arr1', VName
arr2') ->
      [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
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]
[FParam SOACS]
fun_params [TypeBase Shape NoUniqueness]
argts
  let diets :: [Diet]
diets =
        Int -> Diet -> [Diet]
forall a. Int -> a -> [a]
replicate ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
shapeargs) Diet
I.ObservePrim
          [Diet] -> [Diet] -> [Diet]
forall a. [a] -> [a] -> [a]
++ (DeclType -> Diet) -> [DeclType] -> [Diet]
forall a b. (a -> b) -> [a] -> [b]
map DeclType -> Diet
forall shape. TypeBase shape Uniqueness -> Diet
I.diet [DeclType]
value_paramts
  [SubExp]
args' <-
    ErrorMsg SubExp
-> SrcLoc
-> [VName]
-> [TypeBase Shape NoUniqueness]
-> [SubExp]
-> InternaliseM [SubExp]
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
      ((Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
fun_params)
      ((Param DeclType -> TypeBase Shape NoUniqueness)
-> [Param DeclType] -> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> TypeBase Shape NoUniqueness
forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
I.paramType [Param DeclType]
fun_params)
      ([SubExp]
shapeargs [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
args)
  [TypeBase Shape NoUniqueness]
argts' <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
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 ([(SubExp, TypeBase Shape NoUniqueness)]
 -> Maybe [(TypeBase ExtShape Uniqueness, RetAls)])
-> [(SubExp, TypeBase Shape NoUniqueness)]
-> Maybe [(TypeBase ExtShape Uniqueness, RetAls)]
forall a b. (a -> b) -> a -> b
$ [SubExp]
-> [TypeBase Shape NoUniqueness]
-> [(SubExp, TypeBase Shape NoUniqueness)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
args' [TypeBase Shape NoUniqueness]
argts' of
    Maybe [(TypeBase ExtShape Uniqueness, RetAls)]
Nothing ->
      [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> ([[Char]] -> [Char]) -> [[Char]] -> InternaliseM [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[Char]] -> [Char]
unlines ([[Char]] -> InternaliseM [SubExp])
-> [[Char]] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
        [ [Char]
"Cannot apply "
            [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> VName -> [Char]
forall a. Pretty a => a -> [Char]
prettyString VName
fname
            [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" to "
            [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> [Char]
forall a. Show a => a -> [Char]
show ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
args')
            [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" arguments",
          [Char]
" " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [SubExp] -> [Char]
forall a. Pretty a => a -> [Char]
prettyString [SubExp]
args',
          [Char]
"of types",
          [Char]
" " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [TypeBase Shape NoUniqueness] -> [Char]
forall a. Pretty a => a -> [Char]
prettyString [TypeBase Shape NoUniqueness]
argts',
          [Char]
"Function has " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> [Char]
forall a. Show a => a -> [Char]
show ([Param DeclType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Param DeclType]
fun_params) [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" parameters",
          [Char]
" " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Param DeclType] -> [Char]
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 <- (InternaliseEnv -> Attrs) -> InternaliseM Attrs
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks InternaliseEnv -> Attrs
envAttrs
      Attrs -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. MonadBuilder m => Attrs -> m a -> m a
attributing Attrs
attrs (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> (Exp SOACS -> InternaliseM [SubExp])
-> Exp SOACS
-> InternaliseM [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp SOACS -> InternaliseM [SubExp]
letValExp' [Char]
desc (Exp SOACS -> InternaliseM [SubExp])
-> Exp SOACS -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
        Name
-> [(SubExp, Diet)]
-> [(RetType SOACS, RetAls)]
-> (Safety, SrcLoc, [SrcLoc])
-> Exp SOACS
forall rep.
Name
-> [(SubExp, Diet)]
-> [(RetType rep, RetAls)]
-> (Safety, SrcLoc, [SrcLoc])
-> Exp rep
I.Apply (VName -> Name
internaliseFunName VName
fname) ([SubExp] -> [Diet] -> [(SubExp, Diet)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
args' [Diet]
diets) [(TypeBase ExtShape Uniqueness, RetAls)]
[(RetType SOACS, RetAls)]
ts (Safety
safety, SrcLoc
loc, [SrcLoc]
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 = (Tree (TypeBase ExtShape Uniqueness)
 -> [TypeBase ExtShape Uniqueness])
-> [Tree (TypeBase ExtShape Uniqueness)]
-> [TypeBase ExtShape Uniqueness]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Tree (TypeBase ExtShape Uniqueness)
-> [TypeBase ExtShape Uniqueness]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList ([Tree (TypeBase ExtShape Uniqueness)]
 -> [TypeBase ExtShape Uniqueness])
-> [Tree (TypeBase ExtShape Uniqueness)]
-> [TypeBase ExtShape Uniqueness]
forall a b. (a -> b) -> a -> b
$ StructType -> [Tree (TypeBase ExtShape Uniqueness)]
internaliseType (StructType -> [Tree (TypeBase ExtShape Uniqueness)])
-> StructType -> [Tree (TypeBase ExtShape Uniqueness)]
forall a b. (a -> b) -> a -> b
$ StructType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
E.toStruct StructType
ret
  [TypeBase Shape NoUniqueness]
ses_ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
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 =
        [Map VName SubExp] -> Map VName SubExp
forall a. Monoid a => [a] -> a
mconcat ([Map VName SubExp] -> Map VName SubExp)
-> [Map VName SubExp] -> Map VName SubExp
forall a b. (a -> b) -> a -> b
$ (Ext SubExp -> SubExp -> Map VName SubExp)
-> [Ext SubExp] -> [SubExp] -> [Map VName SubExp]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Ext SubExp -> SubExp -> Map VName SubExp
combine' (TypeBase ExtShape Uniqueness -> [Ext SubExp]
forall u. TypeBase ExtShape u -> [Ext SubExp]
arrayExtDims TypeBase ExtShape Uniqueness
t1) (TypeBase Shape NoUniqueness -> [SubExp]
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 VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
retext = VName -> SubExp -> Map VName SubExp
forall k a. k -> a -> Map k a
M.singleton VName
v SubExp
se
      combine' Ext SubExp
_ SubExp
_ = Map VName SubExp
forall a. Monoid a => a
mempty

  [(VName, SubExp)]
-> ((VName, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Map VName SubExp -> [(VName, SubExp)]
forall k a. Map k a -> [(k, a)]
M.toList (Map VName SubExp -> [(VName, SubExp)])
-> Map VName SubExp -> [(VName, SubExp)]
forall a b. (a -> b) -> a -> b
$ [Map VName SubExp] -> Map VName SubExp
forall a. Monoid a => [a] -> a
mconcat ([Map VName SubExp] -> Map VName SubExp)
-> [Map VName SubExp] -> Map VName SubExp
forall a b. (a -> b) -> a -> b
$ (TypeBase ExtShape Uniqueness
 -> TypeBase Shape NoUniqueness -> Map VName SubExp)
-> [TypeBase ExtShape Uniqueness]
-> [TypeBase Shape NoUniqueness]
-> [Map VName SubExp]
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) (((VName, SubExp) -> InternaliseM ()) -> InternaliseM ())
-> ((VName, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \(VName
v, SubExp
se) ->
    [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
v] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se

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

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

  [Param (TypeBase Shape NoUniqueness)]
add_lam_x_params <-
    Int
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
-> InternaliseM [Param (TypeBase Shape NoUniqueness)]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
k (InternaliseM (Param (TypeBase Shape NoUniqueness))
 -> InternaliseM [Param (TypeBase Shape NoUniqueness)])
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
-> InternaliseM [Param (TypeBase Shape NoUniqueness)]
forall a b. (a -> b) -> a -> b
$ [Char]
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"x" (PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
  [Param (TypeBase Shape NoUniqueness)]
add_lam_y_params <-
    Int
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
-> InternaliseM [Param (TypeBase Shape NoUniqueness)]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
k (InternaliseM (Param (TypeBase Shape NoUniqueness))
 -> InternaliseM [Param (TypeBase Shape NoUniqueness)])
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
-> InternaliseM [Param (TypeBase Shape NoUniqueness)]
forall a b. (a -> b) -> a -> b
$ [Char]
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"y" (PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
  Body SOACS
add_lam_body <- Builder SOACS Result -> InternaliseM (Body SOACS)
forall rep (m :: * -> *) somerep.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
 SameScope somerep rep) =>
Builder rep Result -> m (Body rep)
runBodyBuilder (Builder SOACS Result -> InternaliseM (Body SOACS))
-> Builder SOACS Result -> InternaliseM (Body SOACS)
forall a b. (a -> b) -> a -> b
$
    Scope SOACS -> Builder SOACS Result -> Builder SOACS Result
forall a.
Scope SOACS
-> BuilderT SOACS (State VNameSource) a
-> BuilderT SOACS (State VNameSource) a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope ([Param (TypeBase Shape NoUniqueness)] -> Scope SOACS
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams ([Param (TypeBase Shape NoUniqueness)] -> Scope SOACS)
-> [Param (TypeBase Shape NoUniqueness)] -> Scope SOACS
forall a b. (a -> b) -> a -> b
$ [Param (TypeBase Shape NoUniqueness)]
add_lam_x_params [Param (TypeBase Shape NoUniqueness)]
-> [Param (TypeBase Shape NoUniqueness)]
-> [Param (TypeBase Shape NoUniqueness)]
forall a. [a] -> [a] -> [a]
++ [Param (TypeBase Shape NoUniqueness)]
add_lam_y_params) (Builder SOACS Result -> Builder SOACS Result)
-> Builder SOACS Result -> Builder SOACS Result
forall a b. (a -> b) -> a -> b
$
      ([SubExp] -> Result)
-> BuilderT SOACS (State VNameSource) [SubExp]
-> Builder SOACS Result
forall a b.
(a -> b)
-> BuilderT SOACS (State VNameSource) a
-> BuilderT SOACS (State VNameSource) b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [SubExp] -> Result
subExpsRes (BuilderT SOACS (State VNameSource) [SubExp]
 -> Builder SOACS Result)
-> BuilderT SOACS (State VNameSource) [SubExp]
-> Builder SOACS Result
forall a b. (a -> b) -> a -> b
$
        [(Param (TypeBase Shape NoUniqueness),
  Param (TypeBase Shape NoUniqueness))]
-> ((Param (TypeBase Shape NoUniqueness),
     Param (TypeBase Shape NoUniqueness))
    -> BuilderT SOACS (State VNameSource) SubExp)
-> BuilderT SOACS (State VNameSource) [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([Param (TypeBase Shape NoUniqueness)]
-> [Param (TypeBase Shape NoUniqueness)]
-> [(Param (TypeBase Shape NoUniqueness),
     Param (TypeBase Shape NoUniqueness))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param (TypeBase Shape NoUniqueness)]
add_lam_x_params [Param (TypeBase Shape NoUniqueness)]
add_lam_y_params) (((Param (TypeBase Shape NoUniqueness),
   Param (TypeBase Shape NoUniqueness))
  -> BuilderT SOACS (State VNameSource) SubExp)
 -> BuilderT SOACS (State VNameSource) [SubExp])
-> ((Param (TypeBase Shape NoUniqueness),
     Param (TypeBase Shape NoUniqueness))
    -> BuilderT SOACS (State VNameSource) SubExp)
-> BuilderT SOACS (State VNameSource) [SubExp]
forall a b. (a -> b) -> a -> b
$ \(Param (TypeBase Shape NoUniqueness)
x, Param (TypeBase Shape NoUniqueness)
y) ->
          [Char]
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"z" (Exp (Rep (BuilderT SOACS (State VNameSource)))
 -> BuilderT SOACS (State VNameSource) SubExp)
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) SubExp
forall a b. (a -> b) -> a -> b
$
            BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource))))
-> BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
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 (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param (TypeBase Shape NoUniqueness) -> VName
forall dec. Param dec -> VName
I.paramName Param (TypeBase Shape NoUniqueness)
x)
                (VName -> SubExp
I.Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param (TypeBase Shape NoUniqueness) -> VName
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 [Param (TypeBase Shape NoUniqueness)]
-> [Param (TypeBase Shape NoUniqueness)]
-> [Param (TypeBase Shape NoUniqueness)]
forall a. [a] -> [a] -> [a]
++ [Param (TypeBase Shape NoUniqueness)]
add_lam_y_params,
            lambdaReturnType :: [TypeBase Shape NoUniqueness]
I.lambdaReturnType = Int -> TypeBase Shape NoUniqueness -> [TypeBase Shape NoUniqueness]
forall a. Int -> a -> [a]
replicate Int
k (TypeBase Shape NoUniqueness -> [TypeBase Shape NoUniqueness])
-> TypeBase Shape NoUniqueness -> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64
          }
      nes :: [SubExp]
nes = Int -> SubExp -> [SubExp]
forall a. Int -> a -> [a]
replicate ([VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
increments) (SubExp -> [SubExp]) -> SubExp -> [SubExp]
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0

  ScremaForm SOACS
scan <- [Scan SOACS] -> InternaliseM (ScremaForm SOACS)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Scan rep] -> m (ScremaForm rep)
I.scanSOAC [Lambda SOACS -> [SubExp] -> Scan SOACS
forall rep. Lambda rep -> [SubExp] -> Scan rep
I.Scan Lambda SOACS
add_lam [SubExp]
nes]
  [VName]
all_offsets <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM [VName]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [VName]
letTupExp [Char]
"offsets" (Exp (Rep InternaliseM) -> InternaliseM [VName])
-> Exp (Rep InternaliseM) -> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$ Op (Rep InternaliseM) -> Exp (Rep InternaliseM)
forall rep. Op rep -> Exp rep
I.Op (Op (Rep InternaliseM) -> Exp (Rep InternaliseM))
-> Op (Rep InternaliseM) -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
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 <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"last_index" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
Int64 Overflow
OverflowUndef) SubExp
w (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
1 :: Int64)
  let nonempty_body :: InternaliseM (Body SOACS)
nonempty_body = Builder SOACS Result -> InternaliseM (Body SOACS)
forall rep (m :: * -> *) somerep.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
 SameScope somerep rep) =>
Builder rep Result -> m (Body rep)
runBodyBuilder (Builder SOACS Result -> InternaliseM (Body SOACS))
-> Builder SOACS Result -> InternaliseM (Body SOACS)
forall a b. (a -> b) -> a -> b
$
        ([SubExp] -> Result)
-> BuilderT SOACS (State VNameSource) [SubExp]
-> Builder SOACS Result
forall a b.
(a -> b)
-> BuilderT SOACS (State VNameSource) a
-> BuilderT SOACS (State VNameSource) b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [SubExp] -> Result
subExpsRes (BuilderT SOACS (State VNameSource) [SubExp]
 -> Builder SOACS Result)
-> BuilderT SOACS (State VNameSource) [SubExp]
-> Builder SOACS Result
forall a b. (a -> b) -> a -> b
$
          [VName]
-> (VName -> BuilderT SOACS (State VNameSource) SubExp)
-> BuilderT SOACS (State VNameSource) [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
all_offsets ((VName -> BuilderT SOACS (State VNameSource) SubExp)
 -> BuilderT SOACS (State VNameSource) [SubExp])
-> (VName -> BuilderT SOACS (State VNameSource) SubExp)
-> BuilderT SOACS (State VNameSource) [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
offset_array ->
            [Char]
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"last_offset" (Exp (Rep (BuilderT SOACS (State VNameSource)))
 -> BuilderT SOACS (State VNameSource) SubExp)
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource))))
-> BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp
I.Index VName
offset_array (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ [DimIndex SubExp] -> Slice SubExp
forall d. [DimIndex d] -> Slice d
Slice [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
I.DimFix SubExp
last_index]
      empty_body :: InternaliseM (Body (Rep InternaliseM))
empty_body = [SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM ([SubExp] -> InternaliseM (Body (Rep InternaliseM)))
-> [SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall a b. (a -> b) -> a -> b
$ Int -> SubExp -> [SubExp]
forall a. Int -> a -> [a]
replicate Int
k (SubExp -> [SubExp]) -> SubExp -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)
  SubExp
is_empty <- [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"is_empty" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
CmpEq PrimType
int64) SubExp
w (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)
  [VName]
sizes <-
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM [VName]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [VName]
letTupExp [Char]
"partition_size" (Exp SOACS -> InternaliseM [VName])
-> InternaliseM (Exp SOACS) -> InternaliseM [VName]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
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 (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
is_empty) InternaliseM (Body (Rep InternaliseM))
empty_body InternaliseM (Body (Rep InternaliseM))
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 <- [TypeBase Shape NoUniqueness]
-> (TypeBase Shape NoUniqueness -> InternaliseM VName)
-> InternaliseM [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [TypeBase Shape NoUniqueness]
arr_ts ((TypeBase Shape NoUniqueness -> InternaliseM VName)
 -> InternaliseM [VName])
-> (TypeBase Shape NoUniqueness -> InternaliseM VName)
-> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$ \TypeBase Shape NoUniqueness
arr_t ->
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
"partition_dest" (Exp (Rep InternaliseM) -> InternaliseM VName)
-> Exp (Rep InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$
      BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
        PrimType -> [SubExp] -> BasicOp
Scratch (TypeBase Shape NoUniqueness -> PrimType
forall shape u. TypeBase shape u -> PrimType
I.elemType TypeBase Shape NoUniqueness
arr_t) (SubExp
w SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
drop Int
1 (TypeBase Shape NoUniqueness -> [SubExp]
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 <- [Char]
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"c" (PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
    [Param (TypeBase Shape NoUniqueness)]
offset_params <- Int
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
-> InternaliseM [Param (TypeBase Shape NoUniqueness)]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
k (InternaliseM (Param (TypeBase Shape NoUniqueness))
 -> InternaliseM [Param (TypeBase Shape NoUniqueness)])
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
-> InternaliseM [Param (TypeBase Shape NoUniqueness)]
forall a b. (a -> b) -> a -> b
$ [Char]
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"offset" (PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
    [Param (TypeBase Shape NoUniqueness)]
value_params <- (TypeBase Shape NoUniqueness
 -> InternaliseM (Param (TypeBase Shape NoUniqueness)))
-> [TypeBase Shape NoUniqueness]
-> InternaliseM [Param (TypeBase Shape NoUniqueness)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Char]
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"v" (TypeBase Shape NoUniqueness
 -> InternaliseM (Param (TypeBase Shape NoUniqueness)))
-> (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall u. TypeBase Shape u -> TypeBase Shape u
I.rowType) [TypeBase Shape NoUniqueness]
arr_ts
    (SubExp
offset, Stms SOACS
offset_stms) <-
      InternaliseM SubExp
-> InternaliseM (SubExp, Stms (Rep InternaliseM))
forall a.
InternaliseM a -> InternaliseM (a, Stms (Rep InternaliseM))
forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms (InternaliseM SubExp
 -> InternaliseM (SubExp, Stms (Rep InternaliseM)))
-> InternaliseM SubExp
-> InternaliseM (SubExp, Stms (Rep InternaliseM))
forall a b. (a -> b) -> a -> b
$
        [SubExp] -> SubExp -> Int -> [LParam SOACS] -> InternaliseM SubExp
mkOffsetLambdaBody
          ((VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var [VName]
sizes)
          (VName -> SubExp
I.Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param (TypeBase Shape NoUniqueness) -> VName
forall dec. Param dec -> VName
I.paramName Param (TypeBase Shape NoUniqueness)
c_param)
          Int
0
          [Param (TypeBase Shape NoUniqueness)]
[LParam SOACS]
offset_params
    Lambda SOACS -> InternaliseM (Lambda SOACS)
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
      I.Lambda
        { lambdaParams :: [LParam SOACS]
I.lambdaParams = Param (TypeBase Shape NoUniqueness)
c_param Param (TypeBase Shape NoUniqueness)
-> [Param (TypeBase Shape NoUniqueness)]
-> [Param (TypeBase Shape NoUniqueness)]
forall a. a -> [a] -> [a]
: [Param (TypeBase Shape NoUniqueness)]
offset_params [Param (TypeBase Shape NoUniqueness)]
-> [Param (TypeBase Shape NoUniqueness)]
-> [Param (TypeBase Shape NoUniqueness)]
forall a. [a] -> [a] -> [a]
++ [Param (TypeBase Shape NoUniqueness)]
value_params,
          lambdaReturnType :: [TypeBase Shape NoUniqueness]
I.lambdaReturnType =
            Int -> TypeBase Shape NoUniqueness -> [TypeBase Shape NoUniqueness]
forall a. Int -> a -> [a]
replicate ([TypeBase Shape NoUniqueness] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase Shape NoUniqueness]
arr_ts) (PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
              [TypeBase Shape NoUniqueness]
-> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
forall a. [a] -> [a] -> [a]
++ (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
-> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall u. TypeBase Shape u -> TypeBase Shape u
I.rowType [TypeBase Shape NoUniqueness]
arr_ts,
          lambdaBody :: Body SOACS
I.lambdaBody =
            Stms SOACS -> Result -> Body SOACS
forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody Stms SOACS
offset_stms (Result -> Body SOACS) -> Result -> Body SOACS
forall a b. (a -> b) -> a -> b
$
              Int -> SubExpRes -> Result
forall a. Int -> a -> [a]
replicate ([TypeBase Shape NoUniqueness] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase Shape NoUniqueness]
arr_ts) (SubExp -> SubExpRes
subExpRes SubExp
offset)
                Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ [VName] -> Result
I.varsRes ((Param (TypeBase Shape NoUniqueness) -> VName)
-> [Param (TypeBase Shape NoUniqueness)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (TypeBase Shape NoUniqueness) -> VName
forall dec. Param dec -> VName
I.paramName [Param (TypeBase Shape NoUniqueness)]
value_params)
        }
  let spec :: [(Shape, Int, VName)]
spec = [Shape] -> [Int] -> [VName] -> [(Shape, Int, VName)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 (Shape -> [Shape]
forall a. a -> [a]
repeat (Shape -> [Shape]) -> Shape -> [Shape]
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Shape
forall d. [d] -> ShapeBase d
I.Shape [SubExp
w]) (Int -> [Int]
forall a. a -> [a]
repeat Int
1) [VName]
blanks
  [VName]
results <-
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM [VName]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [VName]
letTupExp [Char]
"partition_res" (Exp SOACS -> InternaliseM [VName])
-> (SOAC SOACS -> Exp SOACS) -> SOAC SOACS -> InternaliseM [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Op SOACS -> Exp SOACS
SOAC SOACS -> Exp SOACS
forall rep. Op rep -> Exp rep
I.Op (SOAC SOACS -> InternaliseM [VName])
-> SOAC SOACS -> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$
      SubExp
-> [VName] -> [(Shape, Int, VName)] -> Lambda SOACS -> SOAC SOACS
forall rep.
SubExp
-> [VName] -> [(Shape, Int, VName)] -> Lambda rep -> SOAC rep
I.Scatter SubExp
w (VName
classes VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: [VName]
all_offsets [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
arrs) [(Shape, Int, VName)]
spec Lambda SOACS
write_lam
  SubExp
sizes' <-
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"partition_sizes" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
        [SubExp] -> TypeBase Shape NoUniqueness -> BasicOp
I.ArrayLit ((VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var [VName]
sizes) (TypeBase Shape NoUniqueness -> BasicOp)
-> TypeBase Shape NoUniqueness -> BasicOp
forall a b. (a -> b) -> a -> b
$
          PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64
  ([SubExp], [SubExp]) -> InternaliseM ([SubExp], [SubExp])
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((VName -> SubExp) -> [VName] -> [SubExp]
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
_ [] =
      SubExp -> InternaliseM SubExp
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> InternaliseM SubExp) -> SubExp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ Int64 -> SubExp
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 <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"is_this_one" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
CmpEq PrimType
int64) SubExp
c (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$
              IntType -> Integer -> SubExp
intConst IntType
Int64 (Integer -> SubExp) -> Integer -> SubExp
forall a b. (a -> b) -> a -> b
$
                Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
i
      SubExp
next_one <- [SubExp] -> SubExp -> Int -> [LParam SOACS] -> InternaliseM SubExp
mkOffsetLambdaBody [SubExp]
sizes SubExp
c (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [LParam SOACS]
ps
      SubExp
this_one <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"this_offset"
          (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinOp
-> SubExp -> [SubExp] -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
foldBinOp
            (IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
OverflowUndef)
            (Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (-Int64
1 :: Int64))
            (VName -> SubExp
I.Var (Param (TypeBase Shape NoUniqueness) -> VName
forall dec. Param dec -> VName
I.paramName Param (TypeBase Shape NoUniqueness)
LParam SOACS
p) SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
take Int
i [SubExp]
sizes)
      [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"total_res"
        (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
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
          (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
is_this_one)
          ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
this_one])
          ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
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
  [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [ErrorMsgPart SubExp
"[", PrimType -> SubExp -> 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)) = [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Text -> ErrorMsgPart SubExp
forall a. Text -> ErrorMsgPart a
ErrorString (Text -> ErrorMsgPart SubExp) -> Text -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> Text
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' <- [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp])
-> InternaliseM [[ErrorMsgPart SubExp]]
-> InternaliseM [ErrorMsgPart SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeArg Exp -> InternaliseM [ErrorMsgPart SubExp])
-> [TypeArg Exp] -> InternaliseM [[ErrorMsgPart SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM TypeArg Exp -> InternaliseM [ErrorMsgPart SubExp]
onArg [TypeArg Exp]
args
  [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ ErrorMsgPart SubExp
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. a -> [a] -> [a]
intersperse ErrorMsgPart SubExp
" " ([ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ Text -> ErrorMsgPart SubExp
forall a. Text -> ErrorMsgPart a
ErrorString (QualName VName -> Text
forall a. Pretty a => a -> Text
prettyText QualName VName
v) ErrorMsgPart SubExp
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
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) = StructType -> InternaliseM [ErrorMsgPart SubExp]
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 <- Map Name (TypeBase Exp u) -> Maybe [TypeBase Exp u]
forall a. Map Name a -> Maybe [a]
E.areTupleFields Map Name (TypeBase Exp u)
fs = do
      [[ErrorMsgPart SubExp]]
ts' <- (TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp])
-> [TypeBase Exp u] -> InternaliseM [[ErrorMsgPart SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
forall u. TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError [TypeBase Exp u]
ts
      [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp
"("] [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
-> [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
", "] [[ErrorMsgPart SubExp]]
ts' [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
")"]
  | Bool
otherwise = do
      [[ErrorMsgPart SubExp]]
fs' <- ((Name, TypeBase Exp u) -> InternaliseM [ErrorMsgPart SubExp])
-> [(Name, TypeBase Exp u)] -> InternaliseM [[ErrorMsgPart SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Name, TypeBase Exp u) -> InternaliseM [ErrorMsgPart SubExp]
forall {a} {u}.
Pretty a =>
(a, TypeBase Exp u) -> InternaliseM [ErrorMsgPart SubExp]
onField ([(Name, TypeBase Exp u)] -> InternaliseM [[ErrorMsgPart SubExp]])
-> [(Name, TypeBase Exp u)] -> InternaliseM [[ErrorMsgPart SubExp]]
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase Exp u) -> [(Name, TypeBase Exp u)]
forall k a. Map k a -> [(k, a)]
M.toList Map Name (TypeBase Exp u)
fs
      [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp
"{"] [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
-> [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
", "] [[ErrorMsgPart SubExp]]
fs' [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"}"]
  where
    onField :: (a, TypeBase Exp u) -> InternaliseM [ErrorMsgPart SubExp]
onField (a
k, TypeBase Exp u
te) =
      (Text -> ErrorMsgPart SubExp
forall a. Text -> ErrorMsgPart a
ErrorString (a -> Text
forall a. Pretty a => a -> Text
prettyText a
k Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
": ") :) ([ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp])
-> InternaliseM [ErrorMsgPart SubExp]
-> InternaliseM [ErrorMsgPart SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
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' <- [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
forall a. Monoid a => [a] -> a
mconcat ([[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp])
-> InternaliseM [[ErrorMsgPart SubExp]]
-> InternaliseM [ErrorMsgPart SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> InternaliseM [ErrorMsgPart SubExp])
-> [Exp] -> InternaliseM [[ErrorMsgPart SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Exp -> InternaliseM [ErrorMsgPart SubExp]
sizeExpForError (Shape Exp -> [Exp]
forall dim. Shape dim -> [dim]
E.shapeDims Shape Exp
shape)
  [ErrorMsgPart SubExp]
et' <- StructType -> InternaliseM [ErrorMsgPart SubExp]
forall u. TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError (StructType -> InternaliseM [ErrorMsgPart SubExp])
-> StructType -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar ScalarTypeBase Exp NoUniqueness
et
  [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp]
shape' [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
et'
typeExpForError (E.Scalar (E.Sum Map Name [TypeBase Exp u]
cs)) = do
  [[ErrorMsgPart SubExp]]
cs' <- ((Name, [TypeBase Exp u]) -> InternaliseM [ErrorMsgPart SubExp])
-> [(Name, [TypeBase Exp u])]
-> InternaliseM [[ErrorMsgPart SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Name, [TypeBase Exp u]) -> InternaliseM [ErrorMsgPart SubExp]
forall {a} {u}.
Pretty a =>
(a, [TypeBase Exp u]) -> InternaliseM [ErrorMsgPart SubExp]
onConstructor ([(Name, [TypeBase Exp u])]
 -> InternaliseM [[ErrorMsgPart SubExp]])
-> [(Name, [TypeBase Exp u])]
-> InternaliseM [[ErrorMsgPart SubExp]]
forall a b. (a -> b) -> a -> b
$ Map Name [TypeBase Exp u] -> [(Name, [TypeBase Exp u])]
forall k a. Map k a -> [(k, a)]
M.toList Map Name [TypeBase Exp u]
cs
  [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp]
-> [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
" | "] [[ErrorMsgPart SubExp]]
cs'
  where
    onConstructor :: (a, [TypeBase Exp u]) -> InternaliseM [ErrorMsgPart SubExp]
onConstructor (a
c, [TypeBase Exp u]
ts) = do
      [[ErrorMsgPart SubExp]]
ts' <- (TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp])
-> [TypeBase Exp u] -> InternaliseM [[ErrorMsgPart SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
forall u. TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError [TypeBase Exp u]
ts
      [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ Text -> ErrorMsgPart SubExp
forall a. Text -> ErrorMsgPart a
ErrorString (Text
"#" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> a -> Text
forall a. Pretty a => a -> Text
prettyText a
c Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" ") ErrorMsgPart SubExp
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. a -> [a] -> [a]
: [ErrorMsgPart SubExp]
-> [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
" "] [[ErrorMsgPart SubExp]]
ts'
typeExpForError (E.Scalar Arrow {}) = [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a. a -> InternaliseM a
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 = [ErrorMsgPart a] -> ErrorMsg a
forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg ([ErrorMsgPart a] -> ErrorMsg a)
-> ([ErrorMsgPart a] -> [ErrorMsgPart a])
-> [ErrorMsgPart a]
-> ErrorMsg a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ErrorMsgPart a] -> [ErrorMsgPart a]
forall {a}. [ErrorMsgPart a] -> [ErrorMsgPart a]
compact
  where
    compact :: [ErrorMsgPart a] -> [ErrorMsgPart a]
compact [] = []
    compact (ErrorString Text
x : ErrorString Text
y : [ErrorMsgPart a]
parts) =
      [ErrorMsgPart a] -> [ErrorMsgPart a]
compact (Text -> ErrorMsgPart a
forall a. Text -> ErrorMsgPart a
ErrorString (Text
x Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
y) ErrorMsgPart a -> [ErrorMsgPart a] -> [ErrorMsgPart a]
forall a. a -> [a] -> [a]
: [ErrorMsgPart a]
parts)
    compact (ErrorMsgPart a
x : [ErrorMsgPart a]
y) = ErrorMsgPart a
x ErrorMsgPart a -> [ErrorMsgPart a] -> [ErrorMsgPart a]
forall a. a -> [a] -> [a]
: [ErrorMsgPart a] -> [ErrorMsgPart a]
compact [ErrorMsgPart a]
y

errorShape :: [a] -> ErrorMsg a
errorShape :: forall a. [a] -> ErrorMsg a
errorShape [a]
dims =
  ErrorMsg a
"["
    ErrorMsg a -> ErrorMsg a -> ErrorMsg a
forall a. Semigroup a => a -> a -> a
<> [ErrorMsg a] -> ErrorMsg a
forall a. Monoid a => [a] -> a
mconcat (ErrorMsg a -> [ErrorMsg a] -> [ErrorMsg a]
forall a. a -> [a] -> [a]
intersperse ErrorMsg a
"][" ([ErrorMsg a] -> [ErrorMsg a]) -> [ErrorMsg a] -> [ErrorMsg a]
forall a b. (a -> b) -> a -> b
$ (a -> ErrorMsg a) -> [a] -> [ErrorMsg a]
forall a b. (a -> b) -> [a] -> [b]
map ([ErrorMsgPart a] -> ErrorMsg a
forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg ([ErrorMsgPart a] -> ErrorMsg a)
-> (a -> [ErrorMsgPart a]) -> a -> ErrorMsg a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ErrorMsgPart a -> [ErrorMsgPart a]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ErrorMsgPart a -> [ErrorMsgPart a])
-> (a -> ErrorMsgPart a) -> a -> [ErrorMsgPart a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimType -> a -> ErrorMsgPart a
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64) [a]
dims)
    ErrorMsg a -> ErrorMsg a -> ErrorMsg a
forall a. Semigroup a => a -> a -> a
<> ErrorMsg a
"]"