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

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

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

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

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

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

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

internaliseFunName :: VName -> Name
internaliseFunName :: VName -> Name
internaliseFunName = String -> Name
nameFromString (String -> Name) -> (VName -> String) -> VName -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> String
forall a. Pretty a => a -> String
pretty

internaliseValBind :: E.ValBind -> InternaliseM ()
internaliseValBind :: ValBind -> InternaliseM ()
internaliseValBind fb :: ValBind
fb@(E.ValBind Maybe (Info EntryPoint)
entry VName
fname Maybe (TypeExp VName)
retdecl (Info (StructType
rettype, [VName]
_)) [TypeParamBase VName]
tparams [PatBase Info VName]
params ExpBase Info VName
body Maybe DocComment
_ [AttrInfo]
attrs SrcLoc
loc) = do
  InternaliseM () -> InternaliseM ()
forall a. InternaliseM a -> InternaliseM a
localConstsScope (InternaliseM () -> InternaliseM ())
-> InternaliseM () -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
    [TypeParamBase VName]
-> [PatBase Info VName]
-> ([FParam] -> [[FParam]] -> InternaliseM ())
-> InternaliseM ()
forall a.
[TypeParamBase VName]
-> [PatBase Info VName]
-> ([FParam] -> [[FParam]] -> InternaliseM a)
-> InternaliseM a
bindingFParams [TypeParamBase VName]
tparams [PatBase Info VName]
params (([FParam] -> [[FParam]] -> InternaliseM ()) -> InternaliseM ())
-> ([FParam] -> [[FParam]] -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \[FParam]
shapeparams [[FParam]]
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]
shapeparams

      ErrorMsg SubExp
msg <- case Maybe (TypeExp VName)
retdecl of
        Just TypeExp VName
dt ->
          [ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg
            ([ErrorMsgPart SubExp] -> ErrorMsg SubExp)
-> ([ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp]
-> ErrorMsg SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ErrorMsgPart SubExp
"Function return value does not match shape of type " ErrorMsgPart SubExp
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. a -> [a] -> [a]
:)
            ([ErrorMsgPart SubExp] -> ErrorMsg SubExp)
-> InternaliseM [ErrorMsgPart SubExp]
-> InternaliseM (ErrorMsg SubExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
dt
        Maybe (TypeExp VName)
Nothing -> ErrorMsg SubExp -> InternaliseM (ErrorMsg SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (ErrorMsg SubExp -> InternaliseM (ErrorMsg SubExp))
-> ErrorMsg SubExp -> InternaliseM (ErrorMsg SubExp)
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [ErrorMsgPart SubExp
"Function return value does not match shape of declared return type."]

      (BodyT SOACS
body', [DeclExtType]
rettype') <- InternaliseM (Result, [DeclExtType])
-> InternaliseM (Body (Rep InternaliseM), [DeclExtType])
forall (m :: * -> *) a.
MonadBuilder m =>
m (Result, a) -> m (Body (Rep m), a)
buildBody (InternaliseM (Result, [DeclExtType])
 -> InternaliseM (Body (Rep InternaliseM), [DeclExtType]))
-> InternaliseM (Result, [DeclExtType])
-> InternaliseM (Body (Rep InternaliseM), [DeclExtType])
forall a b. (a -> b) -> a -> b
$ do
        [SubExp]
body_res <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp (VName -> String
baseString VName
fname String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"_res") ExpBase Info VName
body
        [DeclExtType]
rettype' <-
          ([DeclExtType] -> [DeclExtType])
-> InternaliseM [DeclExtType] -> InternaliseM [DeclExtType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [DeclExtType] -> [DeclExtType]
forall u. [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts (InternaliseM [DeclExtType] -> InternaliseM [DeclExtType])
-> ([TypeBase Shape NoUniqueness] -> InternaliseM [DeclExtType])
-> [TypeBase Shape NoUniqueness]
-> InternaliseM [DeclExtType]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StructType
-> [TypeBase Shape NoUniqueness] -> InternaliseM [DeclExtType]
forall shape u.
StructType -> [TypeBase shape u] -> InternaliseM [DeclExtType]
internaliseReturnType StructType
rettype ([TypeBase Shape NoUniqueness] -> InternaliseM [DeclExtType])
-> InternaliseM [TypeBase Shape NoUniqueness]
-> InternaliseM [DeclExtType]
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)
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
body_res
        Result
body_res' <-
          ErrorMsg SubExp
-> SrcLoc -> [ExtType] -> Result -> InternaliseM Result
ensureResultExtShape ErrorMsg SubExp
msg SrcLoc
loc ((DeclExtType -> ExtType) -> [DeclExtType] -> [ExtType]
forall a b. (a -> b) -> [a] -> [b]
map DeclExtType -> ExtType
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl [DeclExtType]
rettype') (Result -> InternaliseM Result) -> Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Result
subExpsRes [SubExp]
body_res
        (Result, [DeclExtType]) -> InternaliseM (Result, [DeclExtType])
forall (f :: * -> *) a. Applicative f => a -> f a
pure
          ( Result
body_res',
            Int -> DeclExtType -> [DeclExtType]
forall a. Int -> a -> [a]
replicate (Set Int -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([DeclExtType] -> Set Int
forall u. [TypeBase ExtShape u] -> Set Int
shapeContext [DeclExtType]
rettype')) (PrimType -> DeclExtType
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64) [DeclExtType] -> [DeclExtType] -> [DeclExtType]
forall a. [a] -> [a] -> [a]
++ [DeclExtType]
rettype'
          )

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

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

      if [[Param DeclType]] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[Param DeclType]]
[[FParam]]
params'
        then VName -> FunDef SOACS -> InternaliseM ()
bindConstant VName
fname FunDef SOACS
fd
        else
          VName -> FunDef SOACS -> FunInfo -> InternaliseM ()
bindFunction
            VName
fname
            FunDef SOACS
fd
            ( [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
$ [[Param DeclType]] -> [Param DeclType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Param DeclType]]
[[FParam]]
params',
              [Param DeclType]
[FParam]
all_params,
              [DeclExtType]
-> [Param DeclType]
-> [(SubExp, TypeBase Shape NoUniqueness)]
-> Maybe [DeclExtType]
forall rt dec.
(IsRetType rt, Typed dec) =>
[rt]
-> [Param dec]
-> [(SubExp, TypeBase Shape NoUniqueness)]
-> Maybe [rt]
applyRetType [DeclExtType]
rettype' [Param DeclType]
all_params
            )

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

generateEntryPoint :: E.EntryPoint -> E.ValBind -> InternaliseM ()
generateEntryPoint :: EntryPoint -> ValBind -> InternaliseM ()
generateEntryPoint (E.EntryPoint [EntryType]
e_paramts EntryType
e_rettype) ValBind
vb = InternaliseM () -> InternaliseM ()
forall a. InternaliseM a -> InternaliseM a
localConstsScope (InternaliseM () -> InternaliseM ())
-> InternaliseM () -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ do
  let (E.ValBind Maybe (Info EntryPoint)
_ VName
ofname Maybe (TypeExp VName)
_ (Info (StructType
rettype, [VName]
_)) [TypeParamBase VName]
tparams [PatBase Info VName]
params ExpBase Info VName
_ Maybe DocComment
_ [AttrInfo]
attrs SrcLoc
loc) = ValBind
vb
  [TypeParamBase VName]
-> [PatBase Info VName]
-> ([FParam] -> [[FParam]] -> InternaliseM ())
-> InternaliseM ()
forall a.
[TypeParamBase VName]
-> [PatBase Info VName]
-> ([FParam] -> [[FParam]] -> InternaliseM a)
-> InternaliseM a
bindingFParams [TypeParamBase VName]
tparams [PatBase Info VName]
params (([FParam] -> [[FParam]] -> InternaliseM ()) -> InternaliseM ())
-> ([FParam] -> [[FParam]] -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \[FParam]
shapeparams [[FParam]]
params' -> do
    [[DeclExtType]]
entry_rettype <- StructType -> InternaliseM [[DeclExtType]]
internaliseEntryReturnType StructType
rettype
    let entry' :: EntryPoint
entry' = Name
-> [(EntryType, [FParam])]
-> (EntryType, [[DeclExtType]])
-> EntryPoint
entryPoint (VName -> Name
baseName VName
ofname) ([EntryType]
-> [[Param DeclType]] -> [(EntryType, [Param DeclType])]
forall a b. [a] -> [b] -> [(a, b)]
zip [EntryType]
e_paramts [[Param DeclType]]
[[FParam]]
params') (EntryType
e_rettype, [[DeclExtType]]
entry_rettype)
        args :: [SubExp]
args = (Param DeclType -> SubExp) -> [Param DeclType] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
I.Var (VName -> SubExp)
-> (Param DeclType -> VName) -> Param DeclType -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName) ([Param DeclType] -> [SubExp]) -> [Param DeclType] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ [[Param DeclType]] -> [Param DeclType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Param DeclType]]
[[FParam]]
params'

    (BodyT SOACS
entry_body, [DeclExtType]
ctx_ts) <- InternaliseM (Result, [DeclExtType])
-> InternaliseM (Body (Rep InternaliseM), [DeclExtType])
forall (m :: * -> *) a.
MonadBuilder m =>
m (Result, a) -> m (Body (Rep m), a)
buildBody (InternaliseM (Result, [DeclExtType])
 -> InternaliseM (Body (Rep InternaliseM), [DeclExtType]))
-> InternaliseM (Result, [DeclExtType])
-> InternaliseM (Body (Rep InternaliseM), [DeclExtType])
forall a b. (a -> b) -> a -> b
$ do
      -- Special case the (rare) situation where the entry point is
      -- not a function.
      Maybe [SubExp]
maybe_const <- VName -> InternaliseM (Maybe [SubExp])
lookupConst VName
ofname
      [SubExp]
vals <- case Maybe [SubExp]
maybe_const of
        Just [SubExp]
ses ->
          [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [SubExp]
ses
        Maybe [SubExp]
Nothing ->
          ([SubExp], [ExtType]) -> [SubExp]
forall a b. (a, b) -> a
fst (([SubExp], [ExtType]) -> [SubExp])
-> InternaliseM ([SubExp], [ExtType]) -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String
-> QualName VName
-> [SubExp]
-> SrcLoc
-> InternaliseM ([SubExp], [ExtType])
funcall String
"entry_result" (VName -> QualName VName
forall v. v -> QualName v
E.qualName VName
ofname) [SubExp]
args SrcLoc
loc
      [SubExp]
ctx <-
        [DeclExtType] -> [[SubExp]] -> [SubExp]
forall u a. [TypeBase ExtShape u] -> [[a]] -> [a]
extractShapeContext ([[DeclExtType]] -> [DeclExtType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[DeclExtType]]
entry_rettype)
          ([[SubExp]] -> [SubExp])
-> InternaliseM [[SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> InternaliseM [SubExp])
-> [SubExp] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((TypeBase Shape NoUniqueness -> [SubExp])
-> InternaliseM (TypeBase Shape NoUniqueness)
-> InternaliseM [SubExp]
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, [DeclExtType]) -> InternaliseM (Result, [DeclExtType])
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 -> DeclExtType) -> [SubExp] -> [DeclExtType]
forall a b. (a -> b) -> [a] -> [b]
map (DeclExtType -> SubExp -> DeclExtType
forall a b. a -> b -> a
const (PrimType -> DeclExtType
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)) [SubExp]
ctx)

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

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

    entryPointType :: (EntryType, [TypeBase shape Uniqueness]) -> [EntryPointType]
entryPointType (EntryType
t, [TypeBase shape Uniqueness]
ts)
      | E.Scalar (E.Prim E.Unsigned {}) <- EntryType -> StructType
E.entryType EntryType
t =
        [Uniqueness -> EntryPointType
I.TypeUnsigned Uniqueness
u]
      | E.Array ()
_ Uniqueness
_ (E.Prim E.Unsigned {}) ShapeDecl (DimDecl VName)
_ <- EntryType -> StructType
E.entryType EntryType
t =
        [Uniqueness -> EntryPointType
I.TypeUnsigned Uniqueness
u]
      | E.Scalar E.Prim {} <- EntryType -> StructType
E.entryType EntryType
t =
        [Uniqueness -> EntryPointType
I.TypeDirect Uniqueness
u]
      | E.Array ()
_ Uniqueness
_ E.Prim {} ShapeDecl (DimDecl VName)
_ <- EntryType -> StructType
E.entryType EntryType
t =
        [Uniqueness -> EntryPointType
I.TypeDirect Uniqueness
u]
      | Bool
otherwise =
        [Uniqueness -> String -> Int -> EntryPointType
I.TypeOpaque Uniqueness
u String
desc (Int -> EntryPointType) -> Int -> EntryPointType
forall a b. (a -> b) -> a -> b
$ [TypeBase shape Uniqueness] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase shape Uniqueness]
ts]
      where
        u :: Uniqueness
u = (Uniqueness -> Uniqueness -> Uniqueness)
-> Uniqueness -> [Uniqueness] -> Uniqueness
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Uniqueness -> Uniqueness -> Uniqueness
forall a. Ord a => a -> a -> a
max Uniqueness
Nonunique ([Uniqueness] -> Uniqueness) -> [Uniqueness] -> Uniqueness
forall a b. (a -> b) -> a -> b
$ (TypeBase shape Uniqueness -> Uniqueness)
-> [TypeBase shape Uniqueness] -> [Uniqueness]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase shape Uniqueness -> Uniqueness
forall shape. TypeBase shape Uniqueness -> Uniqueness
I.uniqueness [TypeBase shape Uniqueness]
ts
        desc :: String
desc = String
-> (TypeExp VName -> String) -> Maybe (TypeExp VName) -> String
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (TypeBase () () -> String
forall a. Pretty a => a -> String
prettyOneLine TypeBase () ()
t') TypeExp VName -> String
forall vn. (Eq vn, IsName vn) => TypeExp vn -> String
typeExpOpaqueName (Maybe (TypeExp VName) -> String)
-> Maybe (TypeExp VName) -> String
forall a b. (a -> b) -> a -> b
$ EntryType -> Maybe (TypeExp VName)
E.entryAscribed EntryType
t
        t' :: TypeBase () ()
t' = StructType -> TypeBase () ()
forall vn as. TypeBase (DimDecl vn) as -> TypeBase () as
noSizes (EntryType -> StructType
E.entryType EntryType
t) TypeBase () () -> Uniqueness -> TypeBase () ()
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`E.setUniqueness` Uniqueness
Nonunique
    typeExpOpaqueName :: TypeExp vn -> String
typeExpOpaqueName (TEApply TypeExp vn
te TypeArgExpDim {} SrcLoc
_) =
      TypeExp vn -> String
typeExpOpaqueName TypeExp vn
te
    typeExpOpaqueName (TEArray TypeExp vn
te DimExp vn
_ SrcLoc
_) =
      let (Int
d, TypeExp vn
te') = TypeExp vn -> (Int, TypeExp vn)
forall vn. TypeExp vn -> (Int, TypeExp vn)
withoutDims TypeExp vn
te
       in String
"arr_" String -> String -> String
forall a. [a] -> [a] -> [a]
++ TypeExp vn -> String
typeExpOpaqueName TypeExp vn
te'
            String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_"
            String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
d)
            String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"d"
    typeExpOpaqueName TypeExp vn
te = TypeExp vn -> String
forall a. Pretty a => a -> String
prettyOneLine TypeExp vn
te

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

internaliseBody :: String -> E.Exp -> InternaliseM Body
internaliseBody :: String -> ExpBase Info VName -> InternaliseM (BodyT SOACS)
internaliseBody String
desc ExpBase Info VName
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
<$> String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp (String
desc String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"_res") ExpBase Info VName
e

bodyFromStms ::
  InternaliseM (Result, a) ->
  InternaliseM (Body, a)
bodyFromStms :: InternaliseM (Result, a) -> InternaliseM (BodyT 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 (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms InternaliseM (Result, a)
m
  (,a
a) (BodyT SOACS -> (BodyT SOACS, a))
-> InternaliseM (BodyT SOACS) -> InternaliseM (BodyT 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

internaliseAppExp :: String -> E.AppExp -> InternaliseM [I.SubExp]
internaliseAppExp :: String -> AppExp -> InternaliseM [SubExp]
internaliseAppExp String
desc (E.Index ExpBase Info VName
e SliceBase Info VName
idxs SrcLoc
loc) = do
  [VName]
vs <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"indexed" ExpBase Info VName
e
  [SubExp]
dims <- case [VName]
vs of
    [] -> [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [] -- 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 (ExpT 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
        ExpT SOACS -> InternaliseM (ExpT SOACS)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpT SOACS -> InternaliseM (ExpT SOACS))
-> ExpT SOACS -> InternaliseM (ExpT SOACS)
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT 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 (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
$ String -> [Exp (Rep InternaliseM)] -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
String -> [Exp (Rep m)] -> m [SubExp]
letSubExps String
desc ([ExpT SOACS] -> InternaliseM [SubExp])
-> InternaliseM [ExpT SOACS] -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (VName -> InternaliseM (ExpT SOACS))
-> [VName] -> InternaliseM [ExpT SOACS]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM (ExpT SOACS)
index [VName]
vs
internaliseAppExp String
desc (E.Range ExpBase Info VName
start Maybe (ExpBase Info VName)
maybe_second Inclusiveness (ExpBase Info VName)
end SrcLoc
loc) = do
  SubExp
start' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"range_start" ExpBase Info VName
start
  SubExp
end' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"range_end" (ExpBase Info VName -> InternaliseM SubExp)
-> ExpBase Info VName -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ case Inclusiveness (ExpBase Info VName)
end of
    DownToExclusive ExpBase Info VName
e -> ExpBase Info VName
e
    ToInclusive ExpBase Info VName
e -> ExpBase Info VName
e
    UpToExclusive ExpBase Info VName
e -> ExpBase Info VName
e
  Maybe SubExp
maybe_second' <-
    (ExpBase Info VName -> InternaliseM SubExp)
-> Maybe (ExpBase Info VName) -> InternaliseM (Maybe SubExp)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"range_second") Maybe (ExpBase Info VName)
maybe_second

  -- Construct an error message in case the range is invalid.
  let conv :: SubExp -> InternaliseM SubExp
conv = case ExpBase Info VName -> PatType
E.typeOf ExpBase Info VName
start of
        E.Scalar (E.Prim (E.Unsigned IntType
_)) -> IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntZ IntType
Int64
        PatType
_ -> 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)
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 (ExpBase Info VName)
end of
                   DownToExclusive {} -> [ErrorMsgPart SubExp
"..>"]
                   ToInclusive {} -> [ErrorMsgPart SubExp
"..."]
                   UpToExclusive {} -> [ErrorMsgPart SubExp
"..<"]
               )
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
end'_i64, ErrorMsgPart SubExp
" is invalid."]

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

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

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

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

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

  (SubExp
distance, SubExp
step_wrong_dir, SubExp
bounds_invalid) <- case Inclusiveness (ExpBase Info VName)
end of
    DownToExclusive {} -> do
      SubExp
step_wrong_dir <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"step_wrong_dir" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
one
      SubExp
distance <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"distance" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
it Overflow
I.OverflowWrap) SubExp
start' SubExp
end'
      SubExp
distance_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 (m :: * -> *) a. Monad m => a -> m a
return (SubExp
distance_i64, SubExp
step_wrong_dir, SubExp
bounds_invalid_downwards)
    UpToExclusive {} -> do
      SubExp
step_wrong_dir <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"step_wrong_dir" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
negone
      SubExp
distance <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"distance" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
it Overflow
I.OverflowWrap) SubExp
end' SubExp
start'
      SubExp
distance_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 (m :: * -> *) a. Monad m => a -> m a
return (SubExp
distance_i64, SubExp
step_wrong_dir, SubExp
bounds_invalid_upwards)
    ToInclusive {} -> do
      SubExp
downwards <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"downwards" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
negone
      SubExp
distance_downwards_exclusive <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"distance_downwards_exclusive" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
it Overflow
I.OverflowWrap) SubExp
start' SubExp
end'
      SubExp
distance_upwards_exclusive <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"distance_upwards_exclusive" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
it Overflow
I.OverflowWrap) SubExp
end' SubExp
start'

      SubExp
bounds_invalid <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"bounds_invalid" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          SubExp
-> BodyT SOACS
-> BodyT SOACS
-> IfDec (BranchType SOACS)
-> ExpT SOACS
forall rep.
SubExp
-> BodyT rep -> BodyT rep -> IfDec (BranchType rep) -> ExpT rep
I.If
            SubExp
downwards
            ([SubExp] -> BodyT SOACS
forall rep. Buildable rep => [SubExp] -> Body rep
resultBody [SubExp
bounds_invalid_downwards])
            ([SubExp] -> BodyT SOACS
forall rep. Buildable rep => [SubExp] -> Body rep
resultBody [SubExp
bounds_invalid_upwards])
            (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ [TypeBase Shape NoUniqueness] -> IfDec ExtType
ifCommon [PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.Bool]
      SubExp
distance_exclusive <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"distance_exclusive" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          SubExp
-> BodyT SOACS
-> BodyT SOACS
-> IfDec (BranchType SOACS)
-> ExpT SOACS
forall rep.
SubExp
-> BodyT rep -> BodyT rep -> IfDec (BranchType rep) -> ExpT rep
I.If
            SubExp
downwards
            ([SubExp] -> BodyT SOACS
forall rep. Buildable rep => [SubExp] -> Body rep
resultBody [SubExp
distance_downwards_exclusive])
            ([SubExp] -> BodyT SOACS
forall rep. Buildable rep => [SubExp] -> Body rep
resultBody [SubExp
distance_upwards_exclusive])
            (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ [TypeBase Shape NoUniqueness] -> IfDec ExtType
ifCommon [PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim (PrimType -> TypeBase Shape NoUniqueness)
-> PrimType -> TypeBase Shape NoUniqueness
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it]
      SubExp
distance_exclusive_i64 <- IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
distance_exclusive
      SubExp
distance <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"distance" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
            BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp
              (IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
I.OverflowWrap)
              SubExp
distance_exclusive_i64
              (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1)
      (SubExp, SubExp, SubExp) -> InternaliseM (SubExp, SubExp, SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
distance, Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
False, SubExp
bounds_invalid)

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

  SubExp
invalid <-
    String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"range_invalid" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogOr SubExp
step_invalid SubExp
bounds_invalid
  SubExp
valid <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"valid" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
I.Not SubExp
invalid
  Certs
cs <- String -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert String
"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 <-
    String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"pos_step" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Mul IntType
Int64 Overflow
I.OverflowWrap) SubExp
step_i64 SubExp
step_sign_i64

  SubExp
num_elems <-
    Certs -> InternaliseM SubExp -> InternaliseM SubExp
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
$
      String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"num_elems" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
        BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Safety -> BinOp
SDivUp IntType
Int64 Safety
I.Unsafe) SubExp
distance SubExp
pos_step

  SubExp
se <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
desc (BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> SubExp -> SubExp -> IntType -> BasicOp
I.Iota SubExp
num_elems SubExp
start' SubExp
step IntType
it)
  [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [SubExp
se]
internaliseAppExp String
desc (E.Coerce ExpBase Info VName
e (TypeDecl TypeExp VName
dt (Info StructType
et)) SrcLoc
loc) = do
  [SubExp]
ses <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
e
  [DeclExtType]
ts <- StructType
-> [TypeBase Shape NoUniqueness] -> InternaliseM [DeclExtType]
forall shape u.
StructType -> [TypeBase shape u] -> InternaliseM [DeclExtType]
internaliseReturnType StructType
et ([TypeBase Shape NoUniqueness] -> InternaliseM [DeclExtType])
-> InternaliseM [TypeBase Shape NoUniqueness]
-> InternaliseM [DeclExtType]
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)
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
ses
  [ErrorMsgPart SubExp]
dt' <- TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
dt
  [(SubExp, DeclExtType)]
-> ((SubExp, DeclExtType) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([SubExp] -> [DeclExtType] -> [(SubExp, DeclExtType)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
ses [DeclExtType]
ts) (((SubExp, DeclExtType) -> InternaliseM SubExp)
 -> InternaliseM [SubExp])
-> ((SubExp, DeclExtType) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \(SubExp
e', DeclExtType
t') -> do
    [SubExp]
dims <- 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 (core language) shape ("]
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ ErrorMsgPart SubExp
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. a -> [a] -> [a]
intersperse ErrorMsgPart SubExp
", " ((SubExp -> ErrorMsgPart SubExp)
-> [SubExp] -> [ErrorMsgPart SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (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 -> String -> SubExp -> InternaliseM SubExp
ensureExtShape ([ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [ErrorMsgPart SubExp]
parts) SrcLoc
loc (DeclExtType -> ExtType
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl DeclExtType
t') String
desc SubExp
e'
internaliseAppExp String
desc e :: AppExp
e@E.Apply {} = do
  (QualName VName
qfname, [(ExpBase Info VName, Maybe VName)]
args) <- AppExp
-> InternaliseM
     (QualName VName, [(ExpBase Info VName, Maybe VName)])
findFuncall AppExp
e

  -- Argument evaluation is outermost-in so that any existential sizes
  -- created by function applications can be brought into scope.
  let fname :: Name
fname = String -> Name
nameFromString (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ Name -> String
forall a. Pretty a => a -> String
pretty (Name -> String) -> Name -> String
forall a b. (a -> b) -> a -> b
$ VName -> Name
baseName (VName -> Name) -> VName -> Name
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname
      loc :: SrcLoc
loc = AppExp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf AppExp
e
      arg_desc :: String
arg_desc = Name -> String
nameToString Name
fname String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_arg"

  -- Some functions are magical (overloaded) and we handle that here.
  case () of
    -- Overloaded functions never take array arguments (except
    -- equality, but those cannot be existential), so we can safely
    -- ignore the existential dimensions.
    ()
      | Just String -> InternaliseM [SubExp]
internalise <- QualName VName
-> [ExpBase Info VName]
-> SrcLoc
-> Maybe (String -> InternaliseM [SubExp])
isOverloadedFunction QualName VName
qfname (((ExpBase Info VName, Maybe VName) -> ExpBase Info VName)
-> [(ExpBase Info VName, Maybe VName)] -> [ExpBase Info VName]
forall a b. (a -> b) -> [a] -> [b]
map (ExpBase Info VName, Maybe VName) -> ExpBase Info VName
forall a b. (a, b) -> a
fst [(ExpBase Info VName, Maybe VName)]
args) SrcLoc
loc ->
        String -> InternaliseM [SubExp]
internalise String
desc
      | 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
<$> ((ExpBase Info VName, Maybe VName) -> InternaliseM [SubExp])
-> [(ExpBase Info VName, Maybe VName)] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String
-> (ExpBase Info VName, Maybe VName) -> InternaliseM [SubExp]
internaliseArg String
arg_desc) ([(ExpBase Info VName, Maybe VName)]
-> [(ExpBase Info VName, Maybe VName)]
forall a. [a] -> [a]
reverse [(ExpBase Info VName, Maybe VName)]
args)
        let args'' :: [(SubExp, Diet)]
args'' = ([SubExp] -> [(SubExp, Diet)]) -> [[SubExp]] -> [(SubExp, Diet)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap [SubExp] -> [(SubExp, Diet)]
forall a. [a] -> [(a, Diet)]
tag [[SubExp]]
args'
        String -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [SubExp]
letTupExp' String
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Name
-> [(SubExp, Diet)]
-> [RetType SOACS]
-> (Safety, SrcLoc, [SrcLoc])
-> ExpT SOACS
forall rep.
Name
-> [(SubExp, Diet)]
-> [RetType rep]
-> (Safety, SrcLoc, [SrcLoc])
-> ExpT rep
I.Apply Name
fname [(SubExp, Diet)]
args'' [PrimType -> DeclExtType
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
rettype] (Safety
Safe, SrcLoc
loc, [])
      | Bool
otherwise -> do
        [SubExp]
args' <- [[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[SubExp]] -> [SubExp])
-> ([[SubExp]] -> [[SubExp]]) -> [[SubExp]] -> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[SubExp]] -> [[SubExp]]
forall a. [a] -> [a]
reverse ([[SubExp]] -> [SubExp])
-> InternaliseM [[SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((ExpBase Info VName, Maybe VName) -> InternaliseM [SubExp])
-> [(ExpBase Info VName, Maybe VName)] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String
-> (ExpBase Info VName, Maybe VName) -> InternaliseM [SubExp]
internaliseArg String
arg_desc) ([(ExpBase Info VName, Maybe VName)]
-> [(ExpBase Info VName, Maybe VName)]
forall a. [a] -> [a]
reverse [(ExpBase Info VName, Maybe VName)]
args)
        ([SubExp], [ExtType]) -> [SubExp]
forall a b. (a, b) -> a
fst (([SubExp], [ExtType]) -> [SubExp])
-> InternaliseM ([SubExp], [ExtType]) -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String
-> QualName VName
-> [SubExp]
-> SrcLoc
-> InternaliseM ([SubExp], [ExtType])
funcall String
desc QualName VName
qfname [SubExp]
args' SrcLoc
loc
internaliseAppExp String
desc (E.LetPat [SizeBinder VName]
sizes PatBase Info VName
pat ExpBase Info VName
e ExpBase Info VName
body SrcLoc
_) =
  String
-> [SizeBinder VName]
-> PatBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM [SubExp])
-> InternaliseM [SubExp]
forall a.
String
-> [SizeBinder VName]
-> PatBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat String
desc [SizeBinder VName]
sizes PatBase Info VName
pat ExpBase Info VName
e ExpBase Info VName
body (String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc)
internaliseAppExp String
_ (E.LetFun VName
ofname ([TypeParamBase VName], [PatBase Info VName],
 Maybe (TypeExp VName), Info StructType, ExpBase Info VName)
_ ExpBase Info VName
_ SrcLoc
_) =
  String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error (String -> InternaliseM [SubExp])
-> String -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String
"Unexpected LetFun " String -> String -> String
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Pretty a => a -> String
pretty VName
ofname
internaliseAppExp String
desc (E.DoLoop [VName]
sparams PatBase Info VName
mergepat ExpBase Info VName
mergeexp LoopFormBase Info VName
form ExpBase Info VName
loopbody SrcLoc
loc) = do
  [SubExp]
ses <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"loop_init" ExpBase Info VName
mergeexp
  ((BodyT SOACS
loopbody', (LoopForm SOACS
form', [Param DeclType]
shapepat, [Param DeclType]
mergepat', [SubExp]
mergeinit')), Stms SOACS
initstms) <-
    InternaliseM
  (BodyT SOACS,
   (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     ((BodyT SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])),
      Stms (Rep InternaliseM))
forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms (InternaliseM
   (BodyT SOACS,
    (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
 -> InternaliseM
      ((BodyT SOACS,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])),
       Stms (Rep InternaliseM)))
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     ((BodyT SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])),
      Stms (Rep InternaliseM))
forall a b. (a -> b) -> a -> b
$ [SubExp]
-> LoopFormBase Info VName
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [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)
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
mergeinit'

  [SubExp]
ctxinit <- [VName]
-> [FParam]
-> [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]
mergepat' [TypeBase Shape NoUniqueness]
mergeinit_ts'

  let ctxmerge :: [(Param DeclType, SubExp)]
ctxmerge = [Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
shapepat [SubExp]
ctxinit
      valmerge :: [(Param DeclType, SubExp)]
valmerge = [Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
mergepat' [SubExp]
mergeinit'
      dropCond :: [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

  -- Ensure that the result of the loop matches the shapes of the
  -- merge parameters.  XXX: Ideally they should already match (by
  -- the source language type rules), but some of our
  -- transformations (esp. defunctionalisation) strips out some size
  -- information.  For a type-correct source program, these reshapes
  -- should simplify away.
  let merge :: [(Param DeclType, SubExp)]
merge = [(Param DeclType, SubExp)]
ctxmerge [(Param DeclType, SubExp)]
-> [(Param DeclType, SubExp)] -> [(Param DeclType, SubExp)]
forall a. [a] -> [a] -> [a]
++ [(Param DeclType, SubExp)]
valmerge
      merge_ts :: [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
  BodyT SOACS
loopbody'' <-
    Scope SOACS
-> InternaliseM (BodyT SOACS) -> InternaliseM (BodyT SOACS)
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] -> Scope SOACS)
-> [Param DeclType] -> Scope SOACS
forall a b. (a -> b) -> a -> b
$ ((Param DeclType, SubExp) -> Param DeclType)
-> [(Param DeclType, SubExp)] -> [Param DeclType]
forall a b. (a -> b) -> [a] -> [b]
map (Param DeclType, SubExp) -> Param DeclType
forall a b. (a, b) -> a
fst [(Param DeclType, SubExp)]
merge) (InternaliseM (BodyT SOACS) -> InternaliseM (BodyT SOACS))
-> (InternaliseM Result -> InternaliseM (BodyT SOACS))
-> InternaliseM Result
-> InternaliseM (BodyT SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LoopForm SOACS
-> InternaliseM (BodyT SOACS) -> InternaliseM (BodyT SOACS)
forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf LoopForm SOACS
form' (InternaliseM (BodyT SOACS) -> InternaliseM (BodyT SOACS))
-> (InternaliseM Result -> InternaliseM (BodyT SOACS))
-> InternaliseM Result
-> InternaliseM (BodyT SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InternaliseM Result -> InternaliseM (BodyT SOACS)
forall (m :: * -> *).
MonadBuilder m =>
m Result -> m (Body (Rep m))
buildBody_ (InternaliseM Result -> InternaliseM (BodyT SOACS))
-> InternaliseM Result -> InternaliseM (BodyT SOACS)
forall a b. (a -> b) -> a -> b
$
      ([SubExp] -> Result)
-> InternaliseM [SubExp] -> InternaliseM Result
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)]
ctxmerge)
          [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)
BodyT 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
      (String -> Exp (Rep InternaliseM) -> InternaliseM [VName]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [VName]
letTupExp String
desc ([(FParam, SubExp)] -> LoopForm SOACS -> BodyT SOACS -> ExpT SOACS
forall rep.
[(FParam rep, SubExp)] -> LoopForm rep -> BodyT rep -> ExpT rep
I.DoLoop ([(Param DeclType, SubExp)]
ctxmerge [(Param DeclType, SubExp)]
-> [(Param DeclType, SubExp)] -> [(Param DeclType, SubExp)]
forall a. Semigroup a => a -> a -> a
<> [(Param DeclType, SubExp)]
valmerge) LoopForm SOACS
form' BodyT 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

    forLoop :: [Param DeclType]
-> [Param DeclType]
-> [SubExp]
-> LoopForm SOACS
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forLoop [Param DeclType]
mergepat' [Param DeclType]
shapepat [SubExp]
mergeinit LoopForm SOACS
form' =
      InternaliseM
  (Result,
   (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a. InternaliseM (Result, a) -> InternaliseM (BodyT SOACS, a)
bodyFromStms (InternaliseM
   (Result,
    (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
 -> InternaliseM
      (BodyT SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> (InternaliseM
      (Result,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
    -> InternaliseM
         (Result,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LoopForm SOACS
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf LoopForm SOACS
form' (InternaliseM
   (Result,
    (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
 -> InternaliseM
      (BodyT SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ do
        [SubExp]
ses <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"loopres" ExpBase Info VName
loopbody
        [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)
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
ses
        [SubExp]
shapeargs <- [VName]
-> [FParam]
-> [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]
mergepat' [TypeBase Shape NoUniqueness]
sets
        (Result,
 (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
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 SOACS
form',
              [Param DeclType]
shapepat,
              [Param DeclType]
mergepat',
              [SubExp]
mergeinit
            )
          )

    handleForm :: [SubExp]
-> LoopFormBase Info VName
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
handleForm [SubExp]
mergeinit (E.ForIn PatBase Info VName
x ExpBase Info VName
arr) = do
      [VName]
arr' <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"for_in_arr" ExpBase Info VName
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)
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 <- String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"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)
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
-> [TypeBase Shape NoUniqueness]
-> ([FParam]
    -> [FParam]
    -> InternaliseM
         (BodyT SOACS,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a.
[TypeParamBase VName]
-> PatBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([FParam] -> [FParam] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' PatBase Info VName
mergepat [TypeBase Shape NoUniqueness]
ts (([FParam]
  -> [FParam]
  -> InternaliseM
       (BodyT SOACS,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
 -> InternaliseM
      (BodyT SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> ([FParam]
    -> [FParam]
    -> InternaliseM
         (BodyT SOACS,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ \[FParam]
shapepat [FParam]
mergepat' ->
        [PatBase Info VName]
-> [TypeBase Shape NoUniqueness]
-> ([LParam]
    -> InternaliseM
         (BodyT SOACS,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a.
[PatBase Info VName]
-> [TypeBase Shape NoUniqueness]
-> ([LParam] -> InternaliseM a)
-> InternaliseM a
bindingLambdaParams [PatBase Info VName
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 shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
rowType [TypeBase Shape NoUniqueness]
arr_ts) (([LParam]
  -> InternaliseM
       (BodyT SOACS,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
 -> InternaliseM
      (BodyT SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> ([LParam]
    -> InternaliseM
         (BodyT SOACS,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ \[LParam]
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]
x_params [VName]
arr'
          [Param DeclType]
-> [Param DeclType]
-> [SubExp]
-> LoopForm SOACS
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forLoop [Param DeclType]
[FParam]
mergepat' [Param DeclType]
[FParam]
shapepat [SubExp]
mergeinit (LoopForm SOACS
 -> InternaliseM
      (BodyT SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> LoopForm SOACS
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$
            VName -> IntType -> SubExp -> [(LParam, VName)] -> LoopForm SOACS
forall rep.
VName -> IntType -> SubExp -> [(LParam rep, VName)] -> LoopForm rep
I.ForLoop VName
i IntType
Int64 SubExp
w [(Param (TypeBase Shape NoUniqueness), VName)]
[(LParam, VName)]
loopvars
    handleForm [SubExp]
mergeinit (E.For IdentBase Info VName
i ExpBase Info VName
num_iterations) = do
      SubExp
num_iterations' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"upper_bound" ExpBase Info VName
num_iterations
      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 (f :: * -> *) a. Applicative f => a -> f a
pure IntType
it
        TypeBase Shape NoUniqueness
_ -> String -> InternaliseM IntType
forall a. HasCallStack => String -> a
error String
"internaliseExp DoLoop: 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)
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
-> [TypeBase Shape NoUniqueness]
-> ([FParam]
    -> [FParam]
    -> InternaliseM
         (BodyT SOACS,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a.
[TypeParamBase VName]
-> PatBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([FParam] -> [FParam] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' PatBase Info VName
mergepat [TypeBase Shape NoUniqueness]
ts (([FParam]
  -> [FParam]
  -> InternaliseM
       (BodyT SOACS,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
 -> InternaliseM
      (BodyT SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> ([FParam]
    -> [FParam]
    -> InternaliseM
         (BodyT SOACS,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$
        \[FParam]
shapepat [FParam]
mergepat' ->
          [Param DeclType]
-> [Param DeclType]
-> [SubExp]
-> LoopForm SOACS
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forLoop [Param DeclType]
[FParam]
mergepat' [Param DeclType]
[FParam]
shapepat [SubExp]
mergeinit (LoopForm SOACS
 -> InternaliseM
      (BodyT SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> LoopForm SOACS
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$
            VName -> IntType -> SubExp -> [(LParam, VName)] -> LoopForm SOACS
forall rep.
VName -> IntType -> SubExp -> [(LParam rep, VName)] -> LoopForm rep
I.ForLoop (IdentBase Info VName -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
E.identName IdentBase Info VName
i) IntType
it SubExp
num_iterations' []
    handleForm [SubExp]
mergeinit (E.While ExpBase Info VName
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)
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
-> [TypeBase Shape NoUniqueness]
-> ([FParam]
    -> [FParam]
    -> InternaliseM
         (BodyT SOACS,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a.
[TypeParamBase VName]
-> PatBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([FParam] -> [FParam] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' PatBase Info VName
mergepat [TypeBase Shape NoUniqueness]
ts (([FParam]
  -> [FParam]
  -> InternaliseM
       (BodyT SOACS,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
 -> InternaliseM
      (BodyT SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> ([FParam]
    -> [FParam]
    -> InternaliseM
         (BodyT SOACS,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ \[FParam]
shapepat [FParam]
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)
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]
-> [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]
shapepat) [FParam]
mergepat' [TypeBase Shape NoUniqueness]
mergeinit_ts

        (SubExp
loop_initial_cond, Stms SOACS
init_loop_cond_stms) <- InternaliseM SubExp
-> InternaliseM (SubExp, 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]
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
paramName Param DeclType
p] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
          [(Param DeclType, SubExp)]
-> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
[FParam]
mergepat' [SubExp]
mergeinit) (((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ())
-> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \(Param DeclType
p, SubExp
se) ->
            Bool -> InternaliseM () -> InternaliseM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (SubExp
se SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== VName -> SubExp
I.Var (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p)) (InternaliseM () -> InternaliseM ())
-> InternaliseM () -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
              [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
                BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
                  case SubExp
se of
                    I.Var VName
v
                      | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ 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 ->
                        ShapeChange SubExp -> VName -> BasicOp
Reshape ((SubExp -> DimChange SubExp) -> [SubExp] -> ShapeChange SubExp
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimCoercion ([SubExp] -> ShapeChange SubExp) -> [SubExp] -> ShapeChange SubExp
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoUniqueness -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
arrayDims (TypeBase Shape NoUniqueness -> [SubExp])
-> TypeBase Shape NoUniqueness -> [SubExp]
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
          String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"loop_cond" ExpBase Info VName
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 SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a. InternaliseM (Result, a) -> InternaliseM (BodyT SOACS, a)
bodyFromStms (InternaliseM
   (Result,
    (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
 -> InternaliseM
      (BodyT SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ do
          [SubExp]
ses <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"loopres" ExpBase Info VName
loopbody
          [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)
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
ses
          Param DeclType
loop_while <- String -> DeclType -> InternaliseM (Param DeclType)
forall (m :: * -> *) dec.
MonadFreshNames m =>
String -> dec -> m (Param dec)
newParam String
"loop_while" (DeclType -> InternaliseM (Param DeclType))
-> DeclType -> InternaliseM (Param DeclType)
forall a b. (a -> b) -> a -> b
$ PrimType -> DeclType
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.Bool
          [SubExp]
shapeargs <- [VName]
-> [FParam]
-> [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]
shapepat) [FParam]
mergepat' [TypeBase Shape NoUniqueness]
sets

          -- Careful not to clobber anything.
          BodyT SOACS
loop_end_cond_body <- BodyT SOACS -> InternaliseM (BodyT SOACS)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Body rep -> m (Body rep)
renameBody (BodyT SOACS -> InternaliseM (BodyT SOACS))
-> (InternaliseM Result -> InternaliseM (BodyT SOACS))
-> InternaliseM Result
-> InternaliseM (BodyT SOACS)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< InternaliseM Result -> InternaliseM (BodyT SOACS)
forall (m :: * -> *).
MonadBuilder m =>
m Result -> m (Body (Rep m))
buildBody_ (InternaliseM Result -> InternaliseM (BodyT SOACS))
-> InternaliseM Result -> InternaliseM (BodyT 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]
shapepat [SubExp]
shapeargs) (((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ())
-> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \(Param DeclType
p, SubExp
se) ->
              Bool -> InternaliseM () -> InternaliseM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (SubExp
se SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== VName -> SubExp
I.Var (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p)) (InternaliseM () -> InternaliseM ())
-> InternaliseM () -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
                [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
            [(Param DeclType, SubExp)]
-> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
[FParam]
mergepat' [SubExp]
ses) (((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ())
-> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \(Param DeclType
p, SubExp
se) ->
              Bool -> InternaliseM () -> InternaliseM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (SubExp
se SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== VName -> SubExp
I.Var (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p)) (InternaliseM () -> InternaliseM ())
-> InternaliseM () -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
                [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
                  BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
                    case SubExp
se of
                      I.Var VName
v
                        | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ 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 ->
                          ShapeChange SubExp -> VName -> BasicOp
Reshape ((SubExp -> DimChange SubExp) -> [SubExp] -> ShapeChange SubExp
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimCoercion ([SubExp] -> ShapeChange SubExp) -> [SubExp] -> ShapeChange SubExp
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoUniqueness -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
arrayDims (TypeBase Shape NoUniqueness -> [SubExp])
-> TypeBase Shape NoUniqueness -> [SubExp]
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
<$> String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"loop_cond" ExpBase Info VName
cond
          Result
loop_end_cond <- Body (Rep InternaliseM) -> InternaliseM Result
forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind Body (Rep InternaliseM)
BodyT SOACS
loop_end_cond_body

          (Result,
 (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
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 SOACS
forall rep. VName -> LoopForm rep
I.WhileLoop (VName -> LoopForm SOACS) -> VName -> LoopForm SOACS
forall a b. (a -> b) -> a -> b
$ Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName Param DeclType
loop_while,
                [Param DeclType]
[FParam]
shapepat,
                Param DeclType
loop_while Param DeclType -> [Param DeclType] -> [Param DeclType]
forall a. a -> [a] -> [a]
: [Param DeclType]
[FParam]
mergepat',
                SubExp
loop_initial_cond SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: [SubExp]
mergeinit
              )
            )
internaliseAppExp String
desc (E.LetWith IdentBase Info VName
name IdentBase Info VName
src SliceBase Info VName
idxs ExpBase Info VName
ve ExpBase Info VName
body SrcLoc
loc) = do
  let pat :: PatBase Info VName
pat = VName -> Info PatType -> SrcLoc -> PatBase Info VName
forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
E.Id (IdentBase Info VName -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
E.identName IdentBase Info VName
name) (IdentBase Info VName -> Info PatType
forall (f :: * -> *) vn. IdentBase f vn -> f PatType
E.identType IdentBase Info VName
name) SrcLoc
loc
      src_t :: Info PatType
src_t = PatType -> PatType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
E.fromStruct (PatType -> PatType) -> Info PatType -> Info PatType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IdentBase Info VName -> Info PatType
forall (f :: * -> *) vn. IdentBase f vn -> f PatType
E.identType IdentBase Info VName
src
      e :: ExpBase Info VName
e = ExpBase Info VName
-> SliceBase Info VName
-> ExpBase Info VName
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> SliceBase f vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
E.Update (QualName VName -> Info PatType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
QualName vn -> f PatType -> 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 -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
E.identName IdentBase Info VName
src) Info PatType
src_t SrcLoc
loc) SliceBase Info VName
idxs ExpBase Info VName
ve SrcLoc
loc
  String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc (ExpBase Info VName -> InternaliseM [SubExp])
-> ExpBase Info VName -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    AppExp -> Info AppRes -> ExpBase Info VName
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
E.AppExp
      ([SizeBinder VName]
-> PatBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> SrcLoc
-> AppExp
forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
E.LetPat [] PatBase Info VName
pat ExpBase Info VName
e ExpBase Info VName
body SrcLoc
loc)
      (AppRes -> Info AppRes
forall a. a -> Info a
Info (PatType -> [VName] -> AppRes
AppRes (ExpBase Info VName -> PatType
E.typeOf ExpBase Info VName
body) [VName]
forall a. Monoid a => a
mempty))
internaliseAppExp String
desc (E.Match ExpBase Info VName
e NonEmpty (CaseBase Info VName)
cs SrcLoc
_) = do
  [SubExp]
ses <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp (String
desc String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_scrutinee") ExpBase Info VName
e
  case NonEmpty (CaseBase Info VName)
-> (CaseBase Info VName, Maybe (NonEmpty (CaseBase Info VName)))
forall a. NonEmpty a -> (a, Maybe (NonEmpty a))
NE.uncons NonEmpty (CaseBase Info VName)
cs of
    (CasePat PatBase Info VName
pCase ExpBase Info VName
eCase SrcLoc
_, Maybe (NonEmpty (CaseBase Info VName))
Nothing) -> do
      (SubExp
_, [SubExp]
pertinent) <- PatBase Info VName -> [SubExp] -> InternaliseM (SubExp, [SubExp])
generateCond PatBase Info VName
pCase [SubExp]
ses
      [SizeBinder VName]
-> PatBase Info VName
-> [SubExp]
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM [SubExp])
-> InternaliseM [SubExp]
forall a.
[SizeBinder VName]
-> PatBase Info VName
-> [SubExp]
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat' [] PatBase Info VName
pCase [SubExp]
pertinent ExpBase Info VName
eCase (String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc)
    (CaseBase Info VName
c, Just NonEmpty (CaseBase Info VName)
cs') -> do
      let CasePat PatBase Info VName
pLast ExpBase Info VName
eLast SrcLoc
_ = NonEmpty (CaseBase Info VName) -> CaseBase Info VName
forall a. NonEmpty a -> a
NE.last NonEmpty (CaseBase Info VName)
cs'
      BodyT SOACS
bFalse <- do
        (SubExp
_, [SubExp]
pertinent) <- PatBase Info VName -> [SubExp] -> InternaliseM (SubExp, [SubExp])
generateCond PatBase Info VName
pLast [SubExp]
ses
        BodyT SOACS
eLast' <- [SizeBinder VName]
-> PatBase Info VName
-> [SubExp]
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM (BodyT SOACS))
-> InternaliseM (BodyT SOACS)
forall a.
[SizeBinder VName]
-> PatBase Info VName
-> [SubExp]
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat' [] PatBase Info VName
pLast [SubExp]
pertinent ExpBase Info VName
eLast (String -> ExpBase Info VName -> InternaliseM (BodyT SOACS)
internaliseBody String
desc)
        (BodyT SOACS -> CaseBase Info VName -> InternaliseM (BodyT SOACS))
-> BodyT SOACS
-> [CaseBase Info VName]
-> InternaliseM (BodyT SOACS)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\BodyT SOACS
bf CaseBase Info VName
c' -> [InternaliseM (Exp (Rep InternaliseM))]
-> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[m (Exp (Rep m))] -> m (Body (Rep m))
eBody ([InternaliseM (Exp (Rep InternaliseM))]
 -> InternaliseM (Body (Rep InternaliseM)))
-> [InternaliseM (Exp (Rep InternaliseM))]
-> InternaliseM (Body (Rep InternaliseM))
forall a b. (a -> b) -> a -> b
$ InternaliseM (ExpT SOACS) -> [InternaliseM (ExpT SOACS)]
forall (m :: * -> *) a. Monad m => a -> m a
return (InternaliseM (ExpT SOACS) -> [InternaliseM (ExpT SOACS)])
-> InternaliseM (ExpT SOACS) -> [InternaliseM (ExpT SOACS)]
forall a b. (a -> b) -> a -> b
$ [SubExp]
-> CaseBase Info VName -> BodyT SOACS -> InternaliseM (ExpT SOACS)
generateCaseIf [SubExp]
ses CaseBase Info VName
c' BodyT SOACS
bf) BodyT SOACS
eLast' ([CaseBase Info VName] -> InternaliseM (BodyT SOACS))
-> [CaseBase Info VName] -> InternaliseM (BodyT SOACS)
forall a b. (a -> b) -> a -> b
$
          [CaseBase Info VName] -> [CaseBase Info VName]
forall a. [a] -> [a]
reverse ([CaseBase Info VName] -> [CaseBase Info VName])
-> [CaseBase Info VName] -> [CaseBase Info VName]
forall a b. (a -> b) -> a -> b
$ NonEmpty (CaseBase Info VName) -> [CaseBase Info VName]
forall a. NonEmpty a -> [a]
NE.init NonEmpty (CaseBase Info VName)
cs'
      String -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [SubExp]
letTupExp' String
desc (ExpT SOACS -> InternaliseM [SubExp])
-> InternaliseM (ExpT SOACS) -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [SubExp]
-> CaseBase Info VName -> BodyT SOACS -> InternaliseM (ExpT SOACS)
generateCaseIf [SubExp]
ses CaseBase Info VName
c BodyT SOACS
bFalse
internaliseAppExp String
desc (E.If ExpBase Info VName
ce ExpBase Info VName
te ExpBase Info VName
fe SrcLoc
_) =
  String -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [SubExp]
letTupExp' String
desc
    (ExpT SOACS -> InternaliseM [SubExp])
-> InternaliseM (ExpT SOACS) -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (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 -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT SOACS)
-> (SubExp -> BasicOp) -> SubExp -> ExpT SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
SubExp (SubExp -> ExpT SOACS)
-> InternaliseM SubExp -> InternaliseM (ExpT SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"cond" ExpBase Info VName
ce)
      (String -> ExpBase Info VName -> InternaliseM (BodyT SOACS)
internaliseBody (String
desc String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"_t") ExpBase Info VName
te)
      (String -> ExpBase Info VName -> InternaliseM (BodyT SOACS)
internaliseBody (String
desc String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"_f") ExpBase Info VName
fe)
internaliseAppExp String
_ e :: AppExp
e@E.BinOp {} =
  String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error (String -> InternaliseM [SubExp])
-> String -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String
"internaliseAppExp: Unexpected BinOp " String -> String -> String
forall a. [a] -> [a] -> [a]
++ AppExp -> String
forall a. Pretty a => a -> String
pretty AppExp
e

internaliseExp :: String -> E.Exp -> InternaliseM [I.SubExp]
internaliseExp :: String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc (E.Parens ExpBase Info VName
e SrcLoc
_) =
  String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
e
internaliseExp String
desc (E.QualParens (QualName VName, SrcLoc)
_ ExpBase Info VName
e SrcLoc
_) =
  String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
e
internaliseExp String
desc (E.StringLit [Word8]
vs SrcLoc
_) =
  (SubExp -> [SubExp])
-> InternaliseM SubExp -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> [SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InternaliseM SubExp -> InternaliseM [SubExp])
-> (ExpT SOACS -> InternaliseM SubExp)
-> ExpT SOACS
-> InternaliseM [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
desc (ExpT SOACS -> InternaliseM [SubExp])
-> ExpT SOACS -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT 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 String
_ (E.Var (E.QualName [VName]
_ VName
name) Info PatType
_ SrcLoc
_) = do
  Maybe [SubExp]
subst <- VName -> InternaliseM (Maybe [SubExp])
lookupSubst VName
name
  case Maybe [SubExp]
subst of
    Just [SubExp]
substs -> [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [SubExp]
substs
    Maybe [SubExp]
Nothing -> [SubExp] -> InternaliseM [SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [VName -> SubExp
I.Var VName
name]
internaliseExp String
desc (E.AppExp AppExp
e (Info AppRes
appres)) = do
  [SubExp]
ses <- String -> AppExp -> InternaliseM [SubExp]
internaliseAppExp String
desc AppExp
e
  AppRes -> [SubExp] -> InternaliseM ()
bindExtSizes AppRes
appres [SubExp]
ses
  [SubExp] -> InternaliseM [SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
ses

-- XXX: we map empty records and tuples to units, because otherwise
-- arrays of unit will lose their sizes.
internaliseExp String
_ (E.TupLit [] SrcLoc
_) =
  [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [PrimValue -> SubExp
forall v. IsValue v => v -> SubExp
constant PrimValue
UnitValue]
internaliseExp String
_ (E.RecordLit [] SrcLoc
_) =
  [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [PrimValue -> SubExp
forall v. IsValue v => v -> SubExp
constant PrimValue
UnitValue]
internaliseExp String
desc (E.TupLit [ExpBase Info VName]
es SrcLoc
_) = [[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[SubExp]] -> [SubExp])
-> InternaliseM [[SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ExpBase Info VName -> InternaliseM [SubExp])
-> [ExpBase Info VName] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc) [ExpBase Info VName]
es
internaliseExp String
desc (E.RecordLit [FieldBase Info VName]
orig_fields SrcLoc
_) =
  ((Name, [SubExp]) -> [SubExp]) -> [(Name, [SubExp])] -> [SubExp]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Name, [SubExp]) -> [SubExp]
forall a b. (a, b) -> b
snd ([(Name, [SubExp])] -> [SubExp])
-> ([Map Name [SubExp]] -> [(Name, [SubExp])])
-> [Map Name [SubExp]]
-> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name [SubExp] -> [(Name, [SubExp])]
forall a. Map Name a -> [(Name, a)]
sortFields (Map Name [SubExp] -> [(Name, [SubExp])])
-> ([Map Name [SubExp]] -> Map Name [SubExp])
-> [Map Name [SubExp]]
-> [(Name, [SubExp])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Map Name [SubExp]] -> Map Name [SubExp]
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
M.unions ([Map Name [SubExp]] -> [SubExp])
-> InternaliseM [Map Name [SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (FieldBase Info VName -> InternaliseM (Map Name [SubExp]))
-> [FieldBase Info VName] -> InternaliseM [Map Name [SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM FieldBase Info VName -> InternaliseM (Map Name [SubExp])
internaliseField [FieldBase Info VName]
orig_fields
  where
    internaliseField :: FieldBase Info VName -> InternaliseM (Map Name [SubExp])
internaliseField (E.RecordFieldExplicit Name
name ExpBase Info VName
e SrcLoc
_) =
      Name -> [SubExp] -> Map Name [SubExp]
forall k a. k -> a -> Map k a
M.singleton Name
name ([SubExp] -> Map Name [SubExp])
-> InternaliseM [SubExp] -> InternaliseM (Map Name [SubExp])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
e
    internaliseField (E.RecordFieldImplicit VName
name Info PatType
t SrcLoc
loc) =
      FieldBase Info VName -> InternaliseM (Map Name [SubExp])
internaliseField (FieldBase Info VName -> InternaliseM (Map Name [SubExp]))
-> FieldBase Info VName -> InternaliseM (Map Name [SubExp])
forall a b. (a -> b) -> a -> b
$
        Name -> ExpBase Info VName -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
E.RecordFieldExplicit
          (VName -> Name
baseName VName
name)
          (QualName VName -> Info PatType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
E.Var (VName -> QualName VName
forall v. v -> QualName v
E.qualName VName
name) Info PatType
t SrcLoc
loc)
          SrcLoc
loc
internaliseExp String
desc (E.ArrayLit [ExpBase Info VName]
es (Info PatType
arr_t) SrcLoc
loc)
  -- If this is a multidimensional array literal of primitives, we
  -- treat it specially by flattening it out followed by a reshape.
  -- This cuts down on the amount of statements that are produced, and
  -- thus allows us to efficiently handle huge array literals - a
  -- corner case, but an important one.
  | Just (([Int]
eshape, [ExpBase Info VName]
e') : [([Int], [ExpBase Info VName])]
es') <- (ExpBase Info VName -> Maybe ([Int], [ExpBase Info VName]))
-> [ExpBase Info VName] -> Maybe [([Int], [ExpBase Info VName])]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ExpBase Info VName -> Maybe ([Int], [ExpBase Info VName])
isArrayLiteral [ExpBase Info VName]
es,
    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [Int] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
eshape,
    (([Int], [ExpBase Info VName]) -> Bool)
-> [([Int], [ExpBase Info VName])] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (([Int]
eshape [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
==) ([Int] -> Bool)
-> (([Int], [ExpBase Info VName]) -> [Int])
-> ([Int], [ExpBase Info VName])
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Int], [ExpBase Info VName]) -> [Int]
forall a b. (a, b) -> a
fst) [([Int], [ExpBase Info VName])]
es',
    Just PatType
basetype <- Int -> PatType -> Maybe PatType
forall dim as. Int -> TypeBase dim as -> Maybe (TypeBase dim as)
E.peelArray ([Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
eshape) PatType
arr_t = do
    let flat_lit :: ExpBase Info VName
flat_lit = [ExpBase Info VName]
-> Info PatType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
[ExpBase f vn] -> f PatType -> SrcLoc -> ExpBase f vn
E.ArrayLit ([ExpBase Info VName]
e' [ExpBase Info VName]
-> [ExpBase Info VName] -> [ExpBase Info VName]
forall a. [a] -> [a] -> [a]
++ (([Int], [ExpBase Info VName]) -> [ExpBase Info VName])
-> [([Int], [ExpBase Info VName])] -> [ExpBase Info VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Int], [ExpBase Info VName]) -> [ExpBase Info VName]
forall a b. (a, b) -> b
snd [([Int], [ExpBase Info VName])]
es') (PatType -> Info PatType
forall a. a -> Info a
Info PatType
basetype) SrcLoc
loc
        new_shape :: [Int]
new_shape = [ExpBase Info VName] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ExpBase Info VName]
es Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
eshape
    [VName]
flat_arrs <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"flat_literal" ExpBase Info VName
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' :: ShapeChange SubExp
new_shape' =
            ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp
reshapeOuter
              ((Int -> DimChange SubExp) -> [Int] -> ShapeChange SubExp
forall a b. (a -> b) -> [a] -> [b]
map (SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimNew (SubExp -> DimChange SubExp)
-> (Int -> SubExp) -> Int -> DimChange SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntType -> Integer -> SubExp
intConst IntType
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 -> ShapeChange SubExp) -> Shape -> ShapeChange SubExp
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
      String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
desc (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
I.Reshape ShapeChange SubExp
new_shape' VName
flat_arr
  | Bool
otherwise = do
    [[SubExp]]
es' <- (ExpBase Info VName -> InternaliseM [SubExp])
-> [ExpBase Info VName] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"arr_elem") [ExpBase Info VName]
es
    [DeclExtType]
arr_t_ext <- StructType -> InternaliseM [DeclExtType]
internaliseType (StructType -> InternaliseM [DeclExtType])
-> StructType -> InternaliseM [DeclExtType]
forall a b. (a -> b) -> a -> b
$ PatType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatType
arr_t

    [TypeBase Shape NoUniqueness]
rowtypes <-
      case (DeclExtType -> Maybe (TypeBase Shape NoUniqueness))
-> [DeclExtType] -> Maybe [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
-> Maybe (TypeBase Shape NoUniqueness)
-> Maybe (TypeBase Shape NoUniqueness)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
rowType (Maybe (TypeBase Shape NoUniqueness)
 -> Maybe (TypeBase Shape NoUniqueness))
-> (DeclExtType -> Maybe (TypeBase Shape NoUniqueness))
-> DeclExtType
-> 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))
-> (DeclExtType -> ExtType)
-> DeclExtType
-> Maybe (TypeBase Shape NoUniqueness)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DeclExtType -> ExtType
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl) [DeclExtType]
arr_t_ext of
        Just [TypeBase Shape NoUniqueness]
ts -> [TypeBase Shape NoUniqueness]
-> InternaliseM [TypeBase Shape NoUniqueness]
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
            [] -> String -> InternaliseM [TypeBase Shape NoUniqueness]
forall a. HasCallStack => String -> a
error (String -> InternaliseM [TypeBase Shape NoUniqueness])
-> String -> InternaliseM [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp ArrayLit: existential type: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PatType -> String
forall a. Pretty a => a -> String
pretty PatType
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)
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 (ExpT 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)
mapM
              ( ErrorMsg SubExp
-> SrcLoc
-> TypeBase Shape NoUniqueness
-> String
-> SubExp
-> InternaliseM SubExp
ensureShape
                  ErrorMsg SubExp
"shape of element differs from shape of first element"
                  SrcLoc
loc
                  TypeBase Shape NoUniqueness
rt
                  String
"elem_reshaped"
              )
              [SubExp]
ks
          ExpT SOACS -> InternaliseM (ExpT SOACS)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpT SOACS -> InternaliseM (ExpT SOACS))
-> ExpT SOACS -> InternaliseM (ExpT SOACS)
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ [SubExp] -> TypeBase Shape NoUniqueness -> BasicOp
I.ArrayLit [SubExp]
ks' TypeBase Shape NoUniqueness
rt

    String -> [Exp (Rep InternaliseM)] -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
String -> [Exp (Rep m)] -> m [SubExp]
letSubExps String
desc
      ([ExpT SOACS] -> InternaliseM [SubExp])
-> InternaliseM [ExpT SOACS] -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< if [[SubExp]] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[SubExp]]
es'
        then (TypeBase Shape NoUniqueness -> InternaliseM (ExpT SOACS))
-> [TypeBase Shape NoUniqueness] -> InternaliseM [ExpT SOACS]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([SubExp]
-> TypeBase Shape NoUniqueness -> InternaliseM (ExpT SOACS)
arraylit []) [TypeBase Shape NoUniqueness]
rowtypes
        else ([SubExp]
 -> TypeBase Shape NoUniqueness -> InternaliseM (ExpT SOACS))
-> [[SubExp]]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM [ExpT SOACS]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM [SubExp]
-> TypeBase Shape NoUniqueness -> InternaliseM (ExpT SOACS)
arraylit ([[SubExp]] -> [[SubExp]]
forall a. [[a]] -> [[a]]
transpose [[SubExp]]
es') [TypeBase Shape NoUniqueness]
rowtypes
  where
    isArrayLiteral :: E.Exp -> Maybe ([Int], [E.Exp])
    isArrayLiteral :: ExpBase Info VName -> Maybe ([Int], [ExpBase Info VName])
isArrayLiteral (E.ArrayLit [ExpBase Info VName]
inner_es Info PatType
_ SrcLoc
_) = do
      ([Int]
eshape, [ExpBase Info VName]
e) : [([Int], [ExpBase Info VName])]
inner_es' <- (ExpBase Info VName -> Maybe ([Int], [ExpBase Info VName]))
-> [ExpBase Info VName] -> Maybe [([Int], [ExpBase Info VName])]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ExpBase Info VName -> Maybe ([Int], [ExpBase Info VName])
isArrayLiteral [ExpBase Info VName]
inner_es
      Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ (([Int], [ExpBase Info VName]) -> Bool)
-> [([Int], [ExpBase Info VName])] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (([Int]
eshape [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
==) ([Int] -> Bool)
-> (([Int], [ExpBase Info VName]) -> [Int])
-> ([Int], [ExpBase Info VName])
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Int], [ExpBase Info VName]) -> [Int]
forall a b. (a, b) -> a
fst) [([Int], [ExpBase Info VName])]
inner_es'
      ([Int], [ExpBase Info VName])
-> Maybe ([Int], [ExpBase Info VName])
forall (m :: * -> *) a. Monad m => a -> m a
return ([ExpBase Info VName] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ExpBase Info VName]
inner_es Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
eshape, [ExpBase Info VName]
e [ExpBase Info VName]
-> [ExpBase Info VName] -> [ExpBase Info VName]
forall a. [a] -> [a] -> [a]
++ (([Int], [ExpBase Info VName]) -> [ExpBase Info VName])
-> [([Int], [ExpBase Info VName])] -> [ExpBase Info VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Int], [ExpBase Info VName]) -> [ExpBase Info VName]
forall a b. (a, b) -> b
snd [([Int], [ExpBase Info VName])]
inner_es')
    isArrayLiteral ExpBase Info VName
e =
      ([Int], [ExpBase Info VName])
-> Maybe ([Int], [ExpBase Info VName])
forall a. a -> Maybe a
Just ([], [ExpBase Info VName
e])
internaliseExp String
desc (E.Ascript ExpBase Info VName
e TypeDeclBase Info VName
_ SrcLoc
_) =
  String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
e
internaliseExp String
desc (E.Negate ExpBase Info VName
e SrcLoc
_) = do
  SubExp
e' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"negate_arg" ExpBase Info VName
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) ->
      String -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [SubExp]
letTupExp' String
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
t Overflow
I.OverflowWrap) (IntType -> Integer -> SubExp
I.intConst IntType
t Integer
0) SubExp
e'
    I.Prim (I.FloatType FloatType
t) ->
      String -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [SubExp]
letTupExp' String
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (FloatType -> BinOp
I.FSub FloatType
t) (FloatType -> Double -> SubExp
I.floatConst FloatType
t Double
0) SubExp
e'
    TypeBase Shape NoUniqueness
_ -> String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error String
"Futhark.Internalise.internaliseExp: non-numeric type in Negate"
internaliseExp String
desc (E.Not ExpBase Info VName
e SrcLoc
_) = do
  SubExp
e' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"not_arg" ExpBase Info VName
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) ->
      String -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [SubExp]
letTupExp' String
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (IntType -> UnOp
I.Complement IntType
t) SubExp
e'
    I.Prim PrimType
I.Bool ->
      String -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [SubExp]
letTupExp' String
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
I.Not SubExp
e'
    TypeBase Shape NoUniqueness
_ ->
      String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error String
"Futhark.Internalise.internaliseExp: non-int/bool type in Not"
internaliseExp String
desc (E.Update ExpBase Info VName
src SliceBase Info VName
slice ExpBase Info VName
ve SrcLoc
loc) = do
  [SubExp]
ves <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"lw_val" ExpBase Info VName
ve
  [VName]
srcs <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"src" ExpBase Info VName
src
  [SubExp]
dims <- case [VName]
srcs of
    [] -> [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [] -- 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
slice

  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
-> String
-> SubExp
-> InternaliseM SubExp
ensureShape
            ErrorMsg SubExp
"shape of value does not match shape of source array"
            SrcLoc
loc
            TypeBase Shape NoUniqueness
rowtype
            String
"lw_val_correct_shape"
            SubExp
ve'
        String
-> VName
-> Slice SubExp
-> Exp (Rep InternaliseM)
-> InternaliseM VName
forall (m :: * -> *).
MonadBuilder m =>
String -> VName -> Slice SubExp -> Exp (Rep m) -> m VName
letInPlace String
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 -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
ve''
  Certs -> InternaliseM [SubExp] -> InternaliseM [SubExp]
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 String
desc (E.RecordUpdate ExpBase Info VName
src [Name]
fields ExpBase Info VName
ve Info PatType
_ SrcLoc
_) = do
  [SubExp]
src' <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
src
  [SubExp]
ve' <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
ve
  StructType
-> [Name] -> [SubExp] -> [SubExp] -> InternaliseM [SubExp]
forall als a.
TypeBase (DimDecl VName) als
-> [Name] -> [a] -> [a] -> InternaliseM [a]
replace (ExpBase Info VName -> PatType
E.typeOf ExpBase Info VName
src PatType -> () -> StructType
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` ()) [Name]
fields [SubExp]
ve' [SubExp]
src'
  where
    replace :: TypeBase (DimDecl VName) als
-> [Name] -> [a] -> [a] -> InternaliseM [a]
replace (E.Scalar (E.Record Map Name (TypeBase (DimDecl VName) als)
m)) (Name
f : [Name]
fs) [a]
ve' [a]
src'
      | Just TypeBase (DimDecl VName) als
t <- Name
-> Map Name (TypeBase (DimDecl VName) als)
-> Maybe (TypeBase (DimDecl VName) als)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f Map Name (TypeBase (DimDecl VName) als)
m = do
        Int
i <-
          ([Int] -> Int) -> InternaliseM [Int] -> InternaliseM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (InternaliseM [Int] -> InternaliseM Int)
-> InternaliseM [Int] -> InternaliseM Int
forall a b. (a -> b) -> a -> b
$
            ((Name, TypeBase (DimDecl VName) als) -> InternaliseM Int)
-> [(Name, TypeBase (DimDecl VName) als)] -> InternaliseM [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (TypeBase (DimDecl VName) als -> InternaliseM Int
forall als. TypeBase (DimDecl VName) als -> InternaliseM Int
internalisedTypeSize (TypeBase (DimDecl VName) als -> InternaliseM Int)
-> ((Name, TypeBase (DimDecl VName) als)
    -> TypeBase (DimDecl VName) als)
-> (Name, TypeBase (DimDecl VName) als)
-> InternaliseM Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, TypeBase (DimDecl VName) als)
-> TypeBase (DimDecl VName) als
forall a b. (a, b) -> b
snd) ([(Name, TypeBase (DimDecl VName) als)] -> InternaliseM [Int])
-> [(Name, TypeBase (DimDecl VName) als)] -> InternaliseM [Int]
forall a b. (a -> b) -> a -> b
$
              ((Name, TypeBase (DimDecl VName) als) -> Bool)
-> [(Name, TypeBase (DimDecl VName) als)]
-> [(Name, TypeBase (DimDecl VName) als)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
/= Name
f) (Name -> Bool)
-> ((Name, TypeBase (DimDecl VName) als) -> Name)
-> (Name, TypeBase (DimDecl VName) als)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, TypeBase (DimDecl VName) als) -> Name
forall a b. (a, b) -> a
fst) ([(Name, TypeBase (DimDecl VName) als)]
 -> [(Name, TypeBase (DimDecl VName) als)])
-> [(Name, TypeBase (DimDecl VName) als)]
-> [(Name, TypeBase (DimDecl VName) als)]
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase (DimDecl VName) als)
-> [(Name, TypeBase (DimDecl VName) als)]
forall a. Map Name a -> [(Name, a)]
sortFields Map Name (TypeBase (DimDecl VName) als)
m
        Int
k <- TypeBase (DimDecl VName) als -> InternaliseM Int
forall als. TypeBase (DimDecl VName) als -> InternaliseM Int
internalisedTypeSize TypeBase (DimDecl VName) als
t
        let ([a]
bef, [a]
to_update, [a]
aft) = Int -> Int -> [a] -> ([a], [a], [a])
forall a. Int -> Int -> [a] -> ([a], [a], [a])
splitAt3 Int
i Int
k [a]
src'
        [a]
src'' <- TypeBase (DimDecl VName) als
-> [Name] -> [a] -> [a] -> InternaliseM [a]
replace TypeBase (DimDecl VName) als
t [Name]
fs [a]
ve' [a]
to_update
        [a] -> InternaliseM [a]
forall (m :: * -> *) a. Monad m => a -> m a
return ([a] -> InternaliseM [a]) -> [a] -> InternaliseM [a]
forall a b. (a -> b) -> a -> b
$ [a]
bef [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
src'' [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
aft
    replace TypeBase (DimDecl VName) als
_ [Name]
_ [a]
ve' [a]
_ = [a] -> InternaliseM [a]
forall (m :: * -> *) a. Monad m => a -> m a
return [a]
ve'
internaliseExp String
desc (E.Attr AttrInfo
attr ExpBase Info VName
e SrcLoc
loc) = do
  [SubExp]
e' <- (InternaliseEnv -> InternaliseEnv)
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local InternaliseEnv -> InternaliseEnv
f (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
e
  case Attr
attr' of
    Attr
"trace" ->
      String -> [SubExp] -> InternaliseM [SubExp]
traceRes (SrcLoc -> String
forall a. Located a => a -> String
locStr SrcLoc
loc) [SubExp]
e'
    I.AttrComp Name
"trace" [I.AttrAtom Name
tag] ->
      String -> [SubExp] -> InternaliseM [SubExp]
traceRes (Name -> String
nameToString 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)
mapM (String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
desc (ExpT SOACS -> InternaliseM SubExp)
-> (SubExp -> ExpT SOACS) -> SubExp -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT SOACS)
-> (SubExp -> BasicOp) -> SubExp -> ExpT SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpaqueOp -> SubExp -> BasicOp
Opaque OpaqueOp
OpaqueNil) [SubExp]
e'
    Attr
_ ->
      [SubExp] -> InternaliseM [SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
e'
  where
    traceRes :: String -> [SubExp] -> InternaliseM [SubExp]
traceRes String
tag' [SubExp]
e' =
      (SubExp -> InternaliseM SubExp)
-> [SubExp] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
desc (ExpT SOACS -> InternaliseM SubExp)
-> (SubExp -> ExpT SOACS) -> SubExp -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT SOACS)
-> (SubExp -> BasicOp) -> SubExp -> ExpT SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpaqueOp -> SubExp -> BasicOp
Opaque (String -> OpaqueOp
OpaqueTrace String
tag')) [SubExp]
e'
    attr' :: Attr
attr' = AttrInfo -> Attr
internaliseAttr AttrInfo
attr
    f :: InternaliseEnv -> InternaliseEnv
f 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 :: Bool
envDoBoundsChecks = Bool
False}
      | Bool
otherwise =
        InternaliseEnv
env {envAttrs :: Attrs
envAttrs = InternaliseEnv -> Attrs
envAttrs InternaliseEnv
env Attrs -> Attrs -> Attrs
forall a. Semigroup a => a -> a -> a
<> Attr -> Attrs
oneAttr Attr
attr'}
internaliseExp String
desc (E.Assert ExpBase Info VName
e1 ExpBase Info VName
e2 (Info String
check) SrcLoc
loc) = do
  SubExp
e1' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"assert_cond" ExpBase Info VName
e1
  Certs
c <- String -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert String
"assert_c" SubExp
e1' ([ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [String -> ErrorMsgPart SubExp
forall a. String -> ErrorMsgPart a
ErrorString (String -> ErrorMsgPart SubExp) -> String -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ String
"Assertion is false: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
check]) SrcLoc
loc
  -- Make sure there are some bindings to certify.
  Certs -> InternaliseM [SubExp] -> InternaliseM [SubExp]
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)
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
=<< String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
e2
  where
    rebind :: SubExp -> m SubExp
rebind SubExp
v = do
      VName
v' <- String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"assert_res"
      [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 -> ExpT 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 (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> m SubExp) -> SubExp -> m SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
v'
internaliseExp String
_ (E.Constr Name
c [ExpBase Info VName]
es (Info (E.Scalar (E.Sum Map Name [PatType]
fs))) SrcLoc
_) = do
  ([DeclExtType]
ts, Map Name (Int, [Int])
constr_map) <- Map Name [StructType]
-> InternaliseM ([DeclExtType], Map Name (Int, [Int]))
internaliseSumType (Map Name [StructType]
 -> InternaliseM ([DeclExtType], Map Name (Int, [Int])))
-> Map Name [StructType]
-> InternaliseM ([DeclExtType], Map Name (Int, [Int]))
forall a b. (a -> b) -> a -> b
$ ([PatType] -> [StructType])
-> Map Name [PatType] -> Map Name [StructType]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((PatType -> StructType) -> [PatType] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map PatType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct) Map Name [PatType]
fs
  [SubExp]
es' <- [[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[SubExp]] -> [SubExp])
-> InternaliseM [[SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ExpBase Info VName -> InternaliseM [SubExp])
-> [ExpBase Info VName] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"payload") [ExpBase Info VName]
es

  let noExt :: p -> m SubExp
noExt p
_ = SubExp -> m SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> m SubExp) -> SubExp -> m SubExp
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
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 (m :: * -> *) p. Monad m => p -> m SubExp
noExt ([ExtType] -> InternaliseM [TypeBase Shape NoUniqueness])
-> [ExtType] -> InternaliseM [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> a -> b
$ (DeclExtType -> ExtType) -> [DeclExtType] -> [ExtType]
forall a b. (a -> b) -> [a] -> [b]
map DeclExtType -> ExtType
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
fromDecl [DeclExtType]
ts

  case Name -> Map Name (Int, [Int]) -> Maybe (Int, [Int])
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
c Map Name (Int, [Int])
constr_map of
    Just (Int
i, [Int]
js) ->
      (IntType -> Integer -> SubExp
intConst IntType
Int8 (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
i) SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
:) ([SubExp] -> [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int
-> [TypeBase Shape NoUniqueness]
-> [(Int, SubExp)]
-> InternaliseM [SubExp]
forall (f :: * -> *) a.
(Num a, MonadBuilder f, Eq a) =>
a -> [TypeBase Shape NoUniqueness] -> [(a, 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 ->
      String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error String
"internaliseExp Constr: missing constructor"
  where
    clauses :: a -> [TypeBase Shape NoUniqueness] -> [(a, SubExp)] -> f [SubExp]
clauses a
j (TypeBase Shape NoUniqueness
t : [TypeBase Shape NoUniqueness]
ts) [(a, SubExp)]
js_to_es
      | Just SubExp
e <- a
j a -> [(a, SubExp)] -> Maybe SubExp
forall a b. Eq a => a -> [(a, b)] -> Maybe b
`lookup` [(a, SubExp)]
js_to_es =
        (SubExp
e SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
:) ([SubExp] -> [SubExp]) -> f [SubExp] -> f [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> [TypeBase Shape NoUniqueness] -> [(a, SubExp)] -> f [SubExp]
clauses (a
j a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) [TypeBase Shape NoUniqueness]
ts [(a, SubExp)]
js_to_es
      | Bool
otherwise = do
        SubExp
blank <- String -> Exp (Rep f) -> f SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"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
=<< 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] -> [SubExp]
forall a. a -> [a] -> [a]
:) ([SubExp] -> [SubExp]) -> f [SubExp] -> f [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> [TypeBase Shape NoUniqueness] -> [(a, SubExp)] -> f [SubExp]
clauses (a
j a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) [TypeBase Shape NoUniqueness]
ts [(a, SubExp)]
js_to_es
    clauses a
_ [] [(a, SubExp)]
_ =
      [SubExp] -> f [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return []
internaliseExp String
_ (E.Constr Name
_ [ExpBase Info VName]
_ (Info PatType
t) SrcLoc
loc) =
  String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error (String -> InternaliseM [SubExp])
-> String -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp: constructor with type " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PatType -> String
forall a. Pretty a => a -> String
pretty PatType
t String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" at " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SrcLoc -> String
forall a. Located a => a -> String
locStr SrcLoc
loc
-- The "interesting" cases are over, now it's mostly boilerplate.

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

internaliseArg :: String -> (E.Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg :: String
-> (ExpBase Info VName, Maybe VName) -> InternaliseM [SubExp]
internaliseArg String
desc (ExpBase Info VName
arg, Maybe VName
argdim) = do
  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 (f :: * -> *) a. Applicative f => a -> f a
pure [VName -> SubExp
I.Var VName
d]
    Maybe VName
_ -> do
      [SubExp]
arg' <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
arg
      case ([SubExp]
arg', Maybe VName
argdim) of
        ([SubExp
se], Just VName
d) -> 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 -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
        ([SubExp], Maybe VName)
_ -> () -> InternaliseM ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
      [SubExp] -> InternaliseM [SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
arg'

subExpPrimType :: I.SubExp -> InternaliseM I.PrimType
subExpPrimType :: SubExp -> InternaliseM PrimType
subExpPrimType = (TypeBase Shape NoUniqueness -> PrimType)
-> InternaliseM (TypeBase Shape NoUniqueness)
-> InternaliseM PrimType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase Shape NoUniqueness -> PrimType
forall shape u. TypeBase shape u -> PrimType
I.elemType (InternaliseM (TypeBase Shape NoUniqueness)
 -> InternaliseM PrimType)
-> (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> SubExp
-> InternaliseM PrimType
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

generateCond :: E.Pat -> [I.SubExp] -> InternaliseM (I.SubExp, [I.SubExp])
generateCond :: PatBase Info VName -> [SubExp] -> InternaliseM (SubExp, [SubExp])
generateCond PatBase Info VName
orig_p [SubExp]
orig_ses = do
  ([SubExp]
cmps, [SubExp]
pertinent, [SubExp]
_) <- PatBase Info VName
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
forall vn.
(Eq vn, IsName vn) =>
PatBase Info vn
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
compares PatBase Info VName
orig_p [SubExp]
orig_ses
  SubExp
cmp <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"matches" (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [SubExp] -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
eAll [SubExp]
cmps
  (SubExp, [SubExp]) -> InternaliseM (SubExp, [SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
cmp, [SubExp]
pertinent)
  where
    -- Literals are always primitive values.
    compares :: PatBase Info vn
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
compares (E.PatLit PatLit
l Info PatType
t SrcLoc
_) (SubExp
se : [SubExp]
ses) = do
      SubExp
e' <- case PatLit
l of
        PatLitPrim PrimValue
v -> SubExp -> InternaliseM SubExp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> InternaliseM SubExp) -> SubExp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
forall v. IsValue v => v -> SubExp
constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimValue
internalisePrimValue PrimValue
v
        PatLitInt Integer
x -> String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"constant" (ExpBase Info VName -> InternaliseM SubExp)
-> ExpBase Info VName -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ Integer -> Info PatType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
Integer -> f PatType -> SrcLoc -> ExpBase f vn
E.IntLit Integer
x Info PatType
t SrcLoc
forall a. Monoid a => a
mempty
        PatLitFloat Double
x -> String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"constant" (ExpBase Info VName -> InternaliseM SubExp)
-> ExpBase Info VName -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ Double -> Info PatType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
Double -> f PatType -> SrcLoc -> ExpBase f vn
E.FloatLit Double
x Info PatType
t SrcLoc
forall a. Monoid a => a
mempty
      PrimType
t' <- SubExp -> InternaliseM PrimType
subExpPrimType SubExp
se
      SubExp
cmp <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"match_lit" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
t') SubExp
e' SubExp
se
      ([SubExp], [SubExp], [SubExp])
-> InternaliseM ([SubExp], [SubExp], [SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return ([SubExp
cmp], [SubExp
se], [SubExp]
ses)
    compares (E.PatConstr Name
c (Info (E.Scalar (E.Sum Map Name [PatType]
fs))) [PatBase Info vn]
pats SrcLoc
_) (SubExp
se : [SubExp]
ses) = do
      ([DeclExtType]
payload_ts, Map Name (Int, [Int])
m) <- Map Name [StructType]
-> InternaliseM ([DeclExtType], Map Name (Int, [Int]))
internaliseSumType (Map Name [StructType]
 -> InternaliseM ([DeclExtType], Map Name (Int, [Int])))
-> Map Name [StructType]
-> InternaliseM ([DeclExtType], Map Name (Int, [Int]))
forall a b. (a -> b) -> a -> b
$ ([PatType] -> [StructType])
-> Map Name [PatType] -> Map Name [StructType]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((PatType -> StructType) -> [PatType] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map PatType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct) Map Name [PatType]
fs
      case Name -> Map Name (Int, [Int]) -> Maybe (Int, [Int])
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
c Map Name (Int, [Int])
m of
        Just (Int
i, [Int]
payload_is) -> do
          let i' :: SubExp
i' = IntType -> Integer -> SubExp
intConst IntType
Int8 (Integer -> SubExp) -> Integer -> SubExp
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
i
          let ([SubExp]
payload_ses, [SubExp]
ses') = Int -> [SubExp] -> ([SubExp], [SubExp])
forall a. Int -> [a] -> ([a], [a])
splitAt ([DeclExtType] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DeclExtType]
payload_ts) [SubExp]
ses
          SubExp
cmp <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"match_constr" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
int8) SubExp
i' SubExp
se
          ([SubExp]
cmps, [SubExp]
pertinent, [SubExp]
_) <- [PatBase Info vn]
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
comparesMany [PatBase Info vn]
pats ([SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp]))
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
forall a b. (a -> b) -> a -> b
$ (Int -> SubExp) -> [Int] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map ([SubExp]
payload_ses [SubExp] -> Int -> SubExp
forall a. [a] -> Int -> a
!!) [Int]
payload_is
          ([SubExp], [SubExp], [SubExp])
-> InternaliseM ([SubExp], [SubExp], [SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
cmp SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: [SubExp]
cmps, [SubExp]
pertinent, [SubExp]
ses')
        Maybe (Int, [Int])
Nothing ->
          String -> InternaliseM ([SubExp], [SubExp], [SubExp])
forall a. HasCallStack => String -> a
error String
"generateCond: missing constructor"
    compares (E.PatConstr Name
_ (Info PatType
t) [PatBase Info vn]
_ SrcLoc
_) [SubExp]
_ =
      String -> InternaliseM ([SubExp], [SubExp], [SubExp])
forall a. HasCallStack => String -> a
error (String -> InternaliseM ([SubExp], [SubExp], [SubExp]))
-> String -> InternaliseM ([SubExp], [SubExp], [SubExp])
forall a b. (a -> b) -> a -> b
$ String
"generateCond: PatConstr has nonsensical type: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PatType -> String
forall a. Pretty a => a -> String
pretty PatType
t
    compares (E.Id vn
_ Info PatType
t SrcLoc
loc) [SubExp]
ses =
      PatBase Info vn
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
compares (Info PatType -> SrcLoc -> PatBase Info vn
forall (f :: * -> *) vn. f PatType -> SrcLoc -> PatBase f vn
E.Wildcard Info PatType
t SrcLoc
loc) [SubExp]
ses
    compares (E.Wildcard (Info PatType
t) SrcLoc
_) [SubExp]
ses = do
      Int
n <- StructType -> InternaliseM Int
forall als. TypeBase (DimDecl VName) als -> InternaliseM Int
internalisedTypeSize (StructType -> InternaliseM Int) -> StructType -> InternaliseM Int
forall a b. (a -> b) -> a -> b
$ PatType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatType
t
      let ([SubExp]
id_ses, [SubExp]
rest_ses) = Int -> [SubExp] -> ([SubExp], [SubExp])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n [SubExp]
ses
      ([SubExp], [SubExp], [SubExp])
-> InternaliseM ([SubExp], [SubExp], [SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return ([], [SubExp]
id_ses, [SubExp]
rest_ses)
    compares (E.PatParens PatBase Info vn
pat SrcLoc
_) [SubExp]
ses =
      PatBase Info vn
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
compares PatBase Info vn
pat [SubExp]
ses
    -- XXX: treat empty tuples and records as bool.
    compares (E.TuplePat [] SrcLoc
loc) [SubExp]
ses =
      PatBase Info vn
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
compares (Info PatType -> SrcLoc -> PatBase Info vn
forall (f :: * -> *) vn. f PatType -> SrcLoc -> PatBase f vn
E.Wildcard (PatType -> Info PatType
forall a. a -> Info a
Info (PatType -> Info PatType) -> PatType -> Info PatType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) Aliasing -> PatType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar (ScalarTypeBase (DimDecl VName) Aliasing -> PatType)
-> ScalarTypeBase (DimDecl VName) Aliasing -> PatType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) Aliasing
forall dim as. PrimType -> ScalarTypeBase dim as
E.Prim PrimType
E.Bool) SrcLoc
loc) [SubExp]
ses
    compares (E.RecordPat [] SrcLoc
loc) [SubExp]
ses =
      PatBase Info vn
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
compares (Info PatType -> SrcLoc -> PatBase Info vn
forall (f :: * -> *) vn. f PatType -> SrcLoc -> PatBase f vn
E.Wildcard (PatType -> Info PatType
forall a. a -> Info a
Info (PatType -> Info PatType) -> PatType -> Info PatType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) Aliasing -> PatType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar (ScalarTypeBase (DimDecl VName) Aliasing -> PatType)
-> ScalarTypeBase (DimDecl VName) Aliasing -> PatType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) Aliasing
forall dim as. PrimType -> ScalarTypeBase dim as
E.Prim PrimType
E.Bool) SrcLoc
loc) [SubExp]
ses
    compares (E.TuplePat [PatBase Info vn]
pats SrcLoc
_) [SubExp]
ses =
      [PatBase Info vn]
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
comparesMany [PatBase Info vn]
pats [SubExp]
ses
    compares (E.RecordPat [(Name, PatBase Info vn)]
fs SrcLoc
_) [SubExp]
ses =
      [PatBase Info vn]
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
comparesMany (((Name, PatBase Info vn) -> PatBase Info vn)
-> [(Name, PatBase Info vn)] -> [PatBase Info vn]
forall a b. (a -> b) -> [a] -> [b]
map (Name, PatBase Info vn) -> PatBase Info vn
forall a b. (a, b) -> b
snd ([(Name, PatBase Info vn)] -> [PatBase Info vn])
-> [(Name, PatBase Info vn)] -> [PatBase Info vn]
forall a b. (a -> b) -> a -> b
$ Map Name (PatBase Info vn) -> [(Name, PatBase Info vn)]
forall a. Map Name a -> [(Name, a)]
E.sortFields (Map Name (PatBase Info vn) -> [(Name, PatBase Info vn)])
-> Map Name (PatBase Info vn) -> [(Name, PatBase Info vn)]
forall a b. (a -> b) -> a -> b
$ [(Name, PatBase Info vn)] -> Map Name (PatBase Info vn)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, PatBase Info vn)]
fs) [SubExp]
ses
    compares (E.PatAscription PatBase Info vn
pat TypeDeclBase Info vn
_ SrcLoc
_) [SubExp]
ses =
      PatBase Info vn
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
compares PatBase Info vn
pat [SubExp]
ses
    compares PatBase Info vn
pat [] =
      String -> InternaliseM ([SubExp], [SubExp], [SubExp])
forall a. HasCallStack => String -> a
error (String -> InternaliseM ([SubExp], [SubExp], [SubExp]))
-> String -> InternaliseM ([SubExp], [SubExp], [SubExp])
forall a b. (a -> b) -> a -> b
$ String
"generateCond: No values left for pattern " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PatBase Info vn -> String
forall a. Pretty a => a -> String
pretty PatBase Info vn
pat

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

generateCaseIf :: [I.SubExp] -> Case -> I.Body -> InternaliseM I.Exp
generateCaseIf :: [SubExp]
-> CaseBase Info VName -> BodyT SOACS -> InternaliseM (ExpT SOACS)
generateCaseIf [SubExp]
ses (CasePat PatBase Info VName
p ExpBase Info VName
eCase SrcLoc
_) BodyT SOACS
bFail = do
  (SubExp
cond, [SubExp]
pertinent) <- PatBase Info VName -> [SubExp] -> InternaliseM (SubExp, [SubExp])
generateCond PatBase Info VName
p [SubExp]
ses
  BodyT SOACS
eCase' <- [SizeBinder VName]
-> PatBase Info VName
-> [SubExp]
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM (BodyT SOACS))
-> InternaliseM (BodyT SOACS)
forall a.
[SizeBinder VName]
-> PatBase Info VName
-> [SubExp]
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat' [] PatBase Info VName
p [SubExp]
pertinent ExpBase Info VName
eCase (String -> ExpBase Info VName -> InternaliseM (BodyT SOACS)
internaliseBody String
"case")
  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
cond) (BodyT SOACS -> InternaliseM (BodyT SOACS)
forall (m :: * -> *) a. Monad m => a -> m a
return BodyT SOACS
eCase') (BodyT SOACS -> InternaliseM (BodyT SOACS)
forall (m :: * -> *) a. Monad m => a -> m a
return BodyT SOACS
bFail)

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

internalisePat' ::
  [E.SizeBinder VName] ->
  E.Pat ->
  [I.SubExp] ->
  E.Exp ->
  (E.Exp -> InternaliseM a) ->
  InternaliseM a
internalisePat' :: [SizeBinder VName]
-> PatBase Info VName
-> [SubExp]
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat' [SizeBinder VName]
sizes PatBase Info VName
p [SubExp]
ses ExpBase Info VName
body ExpBase Info VName -> 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)
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
ses
  PatBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([VName] -> InternaliseM a)
-> InternaliseM a
forall a.
PatBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([VName] -> InternaliseM a)
-> InternaliseM a
stmPat PatBase Info VName
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 (PatType -> [VName] -> AppRes
AppRes (PatBase Info VName -> PatType
E.patternType PatBase Info VName
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 -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
se
    ExpBase Info VName -> InternaliseM a
m ExpBase Info VName
body

internaliseSlice ::
  SrcLoc ->
  [SubExp] ->
  [E.DimIndex] ->
  InternaliseM ([I.DimIndex SubExp], 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 <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"index_ok" (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [SubExp] -> InternaliseM (Exp (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 (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 <- String -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert String
"index_certs" SubExp
ok ErrorMsg SubExp
msg SrcLoc
loc
  ([DimIndex SubExp], Certs)
-> InternaliseM ([DimIndex SubExp], Certs)
forall (m :: * -> *) a. Monad m => a -> m a
return ([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 ExpBase Info VName
i) = do
  (SubExp
i', IntType
_) <- String -> ExpBase Info VName -> InternaliseM (SubExp, IntType)
internaliseDimExp String
"i" ExpBase Info VName
i
  let lowerBound :: ExpT SOACS
lowerBound =
        BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
          CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
I.Int64) (Int64 -> SubExp
forall v. IsValue v => v -> SubExp
I.constant (Int64
0 :: I.Int64)) SubExp
i'
      upperBound :: ExpT SOACS
upperBound =
        BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
          CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSlt IntType
I.Int64) SubExp
i' SubExp
w
  SubExp
ok <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"bounds_check" (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinOp
-> InternaliseM (Exp (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 (ExpT SOACS -> InternaliseM (ExpT SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ExpT SOACS
lowerBound) (ExpT SOACS -> InternaliseM (ExpT SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ExpT SOACS
upperBound)
  (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
I.DimFix SubExp
i', SubExp
ok, [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 (ExpBase Info VName)
Nothing
      Maybe (ExpBase Info VName)
Nothing
      (Just (E.Negate (E.IntLit Integer
1 Info PatType
_ SrcLoc
_) SrcLoc
_))
    ) = do
    SubExp
w_minus_1 <-
      String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"w_minus_1" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
        BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) SubExp
w SubExp
one
    (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return
      ( SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
I.DimSlice SubExp
w_minus_1 SubExp
w (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
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 (ExpBase Info VName)
i Maybe (ExpBase Info VName)
j Maybe (ExpBase Info VName)
s) = do
  SubExp
s' <- InternaliseM SubExp
-> (ExpBase Info VName -> InternaliseM SubExp)
-> Maybe (ExpBase Info VName)
-> InternaliseM SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (SubExp -> InternaliseM SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return SubExp
one) (((SubExp, IntType) -> SubExp)
-> InternaliseM (SubExp, IntType) -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SubExp, IntType) -> SubExp
forall a b. (a, b) -> a
fst (InternaliseM (SubExp, IntType) -> InternaliseM SubExp)
-> (ExpBase Info VName -> InternaliseM (SubExp, IntType))
-> ExpBase Info VName
-> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ExpBase Info VName -> InternaliseM (SubExp, IntType)
internaliseDimExp String
"s") Maybe (ExpBase Info VName)
s
  SubExp
s_sign <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"s_sign" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (IntType -> UnOp
I.SSignum IntType
Int64) SubExp
s'
  SubExp
backwards <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"backwards" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
int64) SubExp
s_sign SubExp
negone
  SubExp
w_minus_1 <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"w_minus_1" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) SubExp
w SubExp
one
  let i_def :: InternaliseM SubExp
i_def =
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"i_def" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          SubExp
-> BodyT SOACS
-> BodyT SOACS
-> IfDec (BranchType SOACS)
-> ExpT SOACS
forall rep.
SubExp
-> BodyT rep -> BodyT rep -> IfDec (BranchType rep) -> ExpT rep
I.If
            SubExp
backwards
            ([SubExp] -> BodyT SOACS
forall rep. Buildable rep => [SubExp] -> Body rep
resultBody [SubExp
w_minus_1])
            ([SubExp] -> BodyT SOACS
forall rep. Buildable rep => [SubExp] -> Body rep
resultBody [SubExp
zero])
            (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ [TypeBase Shape NoUniqueness] -> IfDec ExtType
ifCommon [PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64]
      j_def :: InternaliseM SubExp
j_def =
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"j_def" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          SubExp
-> BodyT SOACS
-> BodyT SOACS
-> IfDec (BranchType SOACS)
-> ExpT SOACS
forall rep.
SubExp
-> BodyT rep -> BodyT rep -> IfDec (BranchType rep) -> ExpT rep
I.If
            SubExp
backwards
            ([SubExp] -> BodyT SOACS
forall rep. Buildable rep => [SubExp] -> Body rep
resultBody [SubExp
negone])
            ([SubExp] -> BodyT SOACS
forall rep. Buildable rep => [SubExp] -> Body rep
resultBody [SubExp
w])
            (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ [TypeBase Shape NoUniqueness] -> IfDec ExtType
ifCommon [PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64]
  SubExp
i' <- InternaliseM SubExp
-> (ExpBase Info VName -> InternaliseM SubExp)
-> Maybe (ExpBase Info VName)
-> InternaliseM SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe InternaliseM SubExp
i_def (((SubExp, IntType) -> SubExp)
-> InternaliseM (SubExp, IntType) -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SubExp, IntType) -> SubExp
forall a b. (a, b) -> a
fst (InternaliseM (SubExp, IntType) -> InternaliseM SubExp)
-> (ExpBase Info VName -> InternaliseM (SubExp, IntType))
-> ExpBase Info VName
-> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ExpBase Info VName -> InternaliseM (SubExp, IntType)
internaliseDimExp String
"i") Maybe (ExpBase Info VName)
i
  SubExp
j' <- InternaliseM SubExp
-> (ExpBase Info VName -> InternaliseM SubExp)
-> Maybe (ExpBase Info VName)
-> InternaliseM SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe InternaliseM SubExp
j_def (((SubExp, IntType) -> SubExp)
-> InternaliseM (SubExp, IntType) -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SubExp, IntType) -> SubExp
forall a b. (a, b) -> a
fst (InternaliseM (SubExp, IntType) -> InternaliseM SubExp)
-> (ExpBase Info VName -> InternaliseM (SubExp, IntType))
-> ExpBase Info VName
-> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ExpBase Info VName -> InternaliseM (SubExp, IntType)
internaliseDimExp String
"j") Maybe (ExpBase Info VName)
j
  SubExp
j_m_i <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"j_m_i" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) SubExp
j' SubExp
i'
  -- Something like a division-rounding-up, but accomodating negative
  -- operands.
  let divRounding :: InternaliseM (ExpT SOACS)
-> InternaliseM (ExpT SOACS)
-> InternaliseM (Exp (Rep InternaliseM))
divRounding InternaliseM (ExpT SOACS)
x InternaliseM (ExpT SOACS)
y =
        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
          (IntType -> Safety -> BinOp
SQuot IntType
Int64 Safety
Safe)
          ( 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
              (IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
I.OverflowWrap)
              InternaliseM (Exp (Rep InternaliseM))
InternaliseM (ExpT SOACS)
x
              (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 (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) InternaliseM (Exp (Rep InternaliseM))
InternaliseM (ExpT SOACS)
y (InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
m (Exp (Rep m)) -> m (Exp (Rep m))
eSignum (InternaliseM (Exp (Rep InternaliseM))
 -> InternaliseM (Exp (Rep InternaliseM)))
-> InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
forall a b. (a -> b) -> a -> b
$ SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
toExp SubExp
s'))
          )
          InternaliseM (Exp (Rep InternaliseM))
InternaliseM (ExpT SOACS)
y
  SubExp
n <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"n" (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (ExpT SOACS)
-> InternaliseM (ExpT SOACS)
-> InternaliseM (Exp (Rep InternaliseM))
divRounding (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> 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))
toExp SubExp
s')

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

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

  SubExp
m <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"m" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) SubExp
n SubExp
one
  SubExp
m_t_s <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"m_t_s" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Mul IntType
Int64 Overflow
I.OverflowWrap) SubExp
m SubExp
s'
  SubExp
i_p_m_t_s <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"i_p_m_t_s" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
I.OverflowWrap) SubExp
i' SubExp
m_t_s
  SubExp
zero_leq_i_p_m_t_s <-
    String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"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 -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int64) SubExp
zero SubExp
i_p_m_t_s
  SubExp
i_p_m_t_s_leq_w <-
    String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"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 -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int64) SubExp
i_p_m_t_s SubExp
w
  SubExp
i_p_m_t_s_lth_w <-
    String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"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 -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSlt IntType
Int64) SubExp
i_p_m_t_s SubExp
w

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

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

  SubExp
slice_ok <-
    String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"slice_ok" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      SubExp
-> BodyT SOACS
-> BodyT SOACS
-> IfDec (BranchType SOACS)
-> ExpT SOACS
forall rep.
SubExp
-> BodyT rep -> BodyT rep -> IfDec (BranchType rep) -> ExpT rep
I.If
        SubExp
backwards
        ([SubExp] -> BodyT SOACS
forall rep. Buildable rep => [SubExp] -> Body rep
resultBody [SubExp
backwards_ok])
        ([SubExp] -> BodyT SOACS
forall rep. Buildable rep => [SubExp] -> Body rep
resultBody [SubExp
forwards_ok])
        (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ [TypeBase Shape NoUniqueness] -> IfDec ExtType
ifCommon [PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.Bool]

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

  SubExp
acceptable <-
    String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"slice_acceptable" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
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 (ExpBase Info VName)
i, Maybe (ExpBase Info VName)
j, Maybe (ExpBase Info VName)
s) of
        (Maybe (ExpBase Info VName)
_, Maybe (ExpBase Info VName)
_, Just {}) ->
          [ ErrorMsgPart SubExp
-> (ExpBase Info VName -> ErrorMsgPart SubExp)
-> Maybe (ExpBase Info VName)
-> ErrorMsgPart SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp
forall a b. a -> b -> a
const (ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp)
-> ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
i') Maybe (ExpBase Info VName)
i,
            ErrorMsgPart SubExp
":",
            ErrorMsgPart SubExp
-> (ExpBase Info VName -> ErrorMsgPart SubExp)
-> Maybe (ExpBase Info VName)
-> ErrorMsgPart SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp
forall a b. a -> b -> a
const (ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp)
-> ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
j') Maybe (ExpBase Info VName)
j,
            ErrorMsgPart SubExp
":",
            PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
s'
          ]
        (Maybe (ExpBase Info VName)
_, Just {}, Maybe (ExpBase Info VName)
_) ->
          [ ErrorMsgPart SubExp
-> (ExpBase Info VName -> ErrorMsgPart SubExp)
-> Maybe (ExpBase Info VName)
-> ErrorMsgPart SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp
forall a b. a -> b -> a
const (ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp)
-> ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
i') Maybe (ExpBase Info VName)
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]
-> (ExpBase Info VName -> [ErrorMsgPart SubExp])
-> Maybe (ExpBase Info VName)
-> [ErrorMsgPart SubExp]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [ErrorMsgPart SubExp]
forall a. Monoid a => a
mempty ([ErrorMsgPart SubExp]
-> ExpBase Info VName -> [ErrorMsgPart SubExp]
forall a b. a -> b -> a
const [ErrorMsgPart SubExp
":", PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
s']) Maybe (ExpBase Info VName)
s
        (Maybe (ExpBase Info VName)
_, Maybe (ExpBase Info VName)
Nothing, Maybe (ExpBase Info VName)
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 (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
I.DimSlice SubExp
i' SubExp
n SubExp
s', SubExp
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 -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)) ->
  (E.Exp, E.Exp, E.Exp, SrcLoc) ->
  InternaliseM [SubExp]
internaliseScanOrReduce :: String
-> String
-> (SubExp
    -> Lambda -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS))
-> (ExpBase Info VName, ExpBase Info VName, ExpBase Info VName,
    SrcLoc)
-> InternaliseM [SubExp]
internaliseScanOrReduce String
desc String
what SubExp
-> Lambda -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)
f (ExpBase Info VName
lam, ExpBase Info VName
ne, ExpBase Info VName
arr, SrcLoc
loc) = do
  [VName]
arrs <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars (String
what String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_arr") ExpBase Info VName
arr
  [SubExp]
nes <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp (String
what String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_ne") ExpBase Info VName
ne
  [SubExp]
nes' <- [(SubExp, VName)]
-> ((SubExp, VName) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([SubExp] -> [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 shape u.
ArrayShape shape =>
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
-> String
-> SubExp
-> InternaliseM SubExp
ensureShape
      ErrorMsg SubExp
"Row shape of input array does not match shape of neutral element"
      SrcLoc
loc
      TypeBase Shape NoUniqueness
rowtype
      (String
what String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_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)
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)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
  Lambda
lam' <- InternaliseLambda
-> ExpBase Info VName
-> [TypeBase Shape NoUniqueness]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM Lambda
internaliseFoldLambda InternaliseLambda
internaliseLambda ExpBase Info VName
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)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
  String -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [SubExp]
letTupExp' String
desc (ExpT SOACS -> InternaliseM [SubExp])
-> (SOAC SOACS -> ExpT SOACS)
-> SOAC SOACS
-> InternaliseM [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SOAC SOACS -> ExpT SOACS
forall rep. Op rep -> ExpT 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 -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)
f SubExp
w Lambda
lam' [SubExp]
nes' [VName]
arrs

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

  -- reshape neutral element to have same size as the destination array
  [SubExp]
ne_shp <- [(SubExp, VName)]
-> ((SubExp, VName) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([SubExp] -> [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 shape u.
ArrayShape shape =>
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
-> String
-> SubExp
-> InternaliseM SubExp
ensureShape
      ErrorMsg SubExp
"Row shape of destination array does not match shape of neutral element"
      SrcLoc
loc
      TypeBase Shape NoUniqueness
rowtype
      String
"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)
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)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
hist'
  Lambda
op' <- InternaliseLambda
-> ExpBase Info VName
-> [TypeBase Shape NoUniqueness]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM Lambda
internaliseFoldLambda InternaliseLambda
internaliseLambda ExpBase Info VName
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_param <- String
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (m :: * -> *) dec.
MonadFreshNames m =>
String -> dec -> m (Param dec)
newParam String
"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)
mapM (String
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (m :: * -> *) dec.
MonadFreshNames m =>
String -> dec -> m (Param dec)
newParam String
"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 shape u.
ArrayShape shape =>
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)
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_param 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 = 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 :: BodyT SOACS
body = Stms SOACS -> Result -> BodyT SOACS
forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody Stms SOACS
forall a. Monoid a => a
mempty (Result -> BodyT SOACS) -> Result -> BodyT 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
paramName [Param (TypeBase Shape NoUniqueness)]
params
  Lambda
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"
        (ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
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)
BodyT SOACS
body

  -- get sizes of histogram and image arrays
  SubExp
w_hist <- 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)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
hist'
  SubExp
w_img <- 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)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
img'

  -- Generate an assertion and reshapes to ensure that buckets' and
  -- img' are the same size.
  Shape
b_shape <- TypeBase Shape NoUniqueness -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape (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
buckets'
  let b_w :: SubExp
b_w = Int -> Shape -> SubExp
shapeSize Int
0 Shape
b_shape
  SubExp
cmp <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"bucket_cmp" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
I.int64) SubExp
b_w SubExp
w_img
  Certs
c <-
    String -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert
      String
"bucket_cert"
      SubExp
cmp
      ErrorMsg SubExp
"length of index and value array does not match"
      SrcLoc
loc
  VName
buckets'' <-
    Certs -> InternaliseM VName -> InternaliseM VName
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
$
      String -> Exp (Rep InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m VName
letExp (VName -> String
baseString VName
buckets') (Exp (Rep InternaliseM) -> InternaliseM VName)
-> Exp (Rep InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$
        BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
I.Reshape (ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp
reshapeOuter [SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimCoercion SubExp
w_img] Int
1 Shape
b_shape) VName
buckets'

  String -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [SubExp]
letTupExp' String
desc (ExpT SOACS -> InternaliseM [SubExp])
-> (SOAC SOACS -> ExpT SOACS)
-> SOAC SOACS
-> InternaliseM [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SOAC SOACS -> ExpT SOACS
forall rep. Op rep -> ExpT rep
I.Op (SOAC SOACS -> InternaliseM [SubExp])
-> SOAC SOACS -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    SubExp -> [HistOp SOACS] -> Lambda -> [VName] -> SOAC SOACS
forall rep.
SubExp -> [HistOp rep] -> Lambda rep -> [VName] -> SOAC rep
I.Hist SubExp
w_img [SubExp -> SubExp -> [VName] -> [SubExp] -> Lambda -> HistOp SOACS
forall rep.
SubExp -> SubExp -> [VName] -> [SubExp] -> Lambda rep -> HistOp rep
HistOp SubExp
w_hist SubExp
rf' [VName]
hist' [SubExp]
ne_shp Lambda
op'] Lambda
lam' ([VName] -> SOAC SOACS) -> [VName] -> SOAC SOACS
forall a b. (a -> b) -> a -> b
$ VName
buckets'' VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: [VName]
img'

internaliseStreamMap ::
  String ->
  StreamOrd ->
  E.Exp ->
  E.Exp ->
  InternaliseM [SubExp]
internaliseStreamMap :: String
-> StreamOrd
-> ExpBase Info VName
-> ExpBase Info VName
-> InternaliseM [SubExp]
internaliseStreamMap String
desc StreamOrd
o ExpBase Info VName
lam ExpBase Info VName
arr = do
  [VName]
arrs <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"stream_input" ExpBase Info VName
arr
  Lambda
lam' <- InternaliseLambda
-> ExpBase Info VName -> [SubExp] -> InternaliseM Lambda
internaliseStreamMapLambda InternaliseLambda
internaliseLambda ExpBase Info VName
lam ([SubExp] -> InternaliseM Lambda)
-> [SubExp] -> InternaliseM Lambda
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var [VName]
arrs
  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)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
  let form :: StreamForm SOACS
form = StreamOrd -> Commutativity -> Lambda -> StreamForm SOACS
forall rep.
StreamOrd -> Commutativity -> Lambda rep -> StreamForm rep
I.Parallel StreamOrd
o Commutativity
Commutative ([LParam] -> BodyT SOACS -> [TypeBase Shape NoUniqueness] -> Lambda
forall rep.
[LParam rep]
-> BodyT rep -> [TypeBase Shape NoUniqueness] -> LambdaT rep
I.Lambda [] (Stms SOACS -> Result -> BodyT SOACS
forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody Stms SOACS
forall a. Monoid a => a
mempty []) [])
  String -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [SubExp]
letTupExp' String
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Op SOACS -> ExpT SOACS
forall rep. Op rep -> ExpT rep
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp
-> [VName] -> StreamForm SOACS -> [SubExp] -> Lambda -> SOAC SOACS
forall rep.
SubExp
-> [VName] -> StreamForm rep -> [SubExp] -> Lambda rep -> SOAC rep
I.Stream SubExp
w [VName]
arrs StreamForm SOACS
form [] Lambda
lam'

internaliseStreamRed ::
  String ->
  StreamOrd ->
  Commutativity ->
  E.Exp ->
  E.Exp ->
  E.Exp ->
  InternaliseM [SubExp]
internaliseStreamRed :: String
-> StreamOrd
-> Commutativity
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> InternaliseM [SubExp]
internaliseStreamRed String
desc StreamOrd
o Commutativity
comm ExpBase Info VName
lam0 ExpBase Info VName
lam ExpBase Info VName
arr = do
  [VName]
arrs <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"stream_input" ExpBase Info VName
arr
  [TypeBase Shape NoUniqueness]
rowts <- (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)
mapM ((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
fmap TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
I.rowType (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]
arrs
  ([Param (TypeBase Shape NoUniqueness)]
lam_params, BodyT SOACS
lam_body) <-
    InternaliseLambda
-> ExpBase Info VName
-> [TypeBase Shape NoUniqueness]
-> InternaliseM ([LParam], BodyT SOACS)
internaliseStreamLambda InternaliseLambda
internaliseLambda ExpBase Info VName
lam [TypeBase Shape NoUniqueness]
rowts
  let (Param (TypeBase Shape NoUniqueness)
chunk_param, [Param (TypeBase Shape NoUniqueness)]
_, [Param (TypeBase Shape NoUniqueness)]
lam_val_params) =
        Int
-> [Param (TypeBase Shape NoUniqueness)]
-> (Param (TypeBase Shape NoUniqueness),
    [Param (TypeBase Shape NoUniqueness)],
    [Param (TypeBase Shape NoUniqueness)])
forall dec.
Int -> [Param dec] -> (Param dec, [Param dec], [Param dec])
partitionChunkedFoldParameters Int
0 [Param (TypeBase Shape NoUniqueness)]
lam_params

  -- Synthesize neutral elements by applying the fold function
  -- to an empty chunk.
  [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)
chunk_param] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
    BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)
  [Param (TypeBase Shape NoUniqueness)]
-> (Param (TypeBase Shape NoUniqueness) -> InternaliseM ())
-> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Param (TypeBase Shape NoUniqueness)]
lam_val_params ((Param (TypeBase Shape NoUniqueness) -> InternaliseM ())
 -> InternaliseM ())
-> (Param (TypeBase Shape NoUniqueness) -> InternaliseM ())
-> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \Param (TypeBase Shape NoUniqueness)
p ->
    [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 (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS)
-> ([SubExp] -> BasicOp) -> [SubExp] -> ExpT SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimType -> [SubExp] -> BasicOp
I.Scratch (TypeBase Shape NoUniqueness -> PrimType
forall shape u. TypeBase shape u -> PrimType
I.elemType (TypeBase Shape NoUniqueness -> PrimType)
-> TypeBase Shape NoUniqueness -> PrimType
forall a b. (a -> b) -> a -> b
$ Param (TypeBase Shape NoUniqueness) -> TypeBase Shape NoUniqueness
forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
I.paramType Param (TypeBase Shape NoUniqueness)
p) ([SubExp] -> ExpT SOACS) -> [SubExp] -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
        TypeBase Shape NoUniqueness -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
I.arrayDims (TypeBase Shape NoUniqueness -> [SubExp])
-> TypeBase Shape NoUniqueness -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Param (TypeBase Shape NoUniqueness) -> TypeBase Shape NoUniqueness
forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
I.paramType Param (TypeBase Shape NoUniqueness)
p
  Result
nes <- BodyT SOACS -> InternaliseM Result
forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind (BodyT SOACS -> InternaliseM Result)
-> InternaliseM (BodyT SOACS) -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BodyT SOACS -> InternaliseM (BodyT SOACS)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Body rep -> m (Body rep)
renameBody BodyT SOACS
lam_body

  [TypeBase Shape NoUniqueness]
nes_ts <- (SubExpRes -> InternaliseM (TypeBase Shape NoUniqueness))
-> Result -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExpRes -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExpRes -> m (TypeBase Shape NoUniqueness)
I.subExpResType Result
nes
  SubExp
outsz <- 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)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
  let acc_arr_tps :: [TypeBase Shape NoUniqueness]
acc_arr_tps = [TypeBase Shape NoUniqueness
-> Shape -> NoUniqueness -> TypeBase Shape NoUniqueness
forall shape u_unused u.
ArrayShape shape =>
TypeBase shape u_unused -> shape -> u -> TypeBase shape u
I.arrayOf TypeBase Shape NoUniqueness
t ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
I.Shape [SubExp
outsz]) NoUniqueness
NoUniqueness | TypeBase Shape NoUniqueness
t <- [TypeBase Shape NoUniqueness]
nes_ts]
  Lambda
lam0' <- InternaliseLambda
-> ExpBase Info VName
-> [TypeBase Shape NoUniqueness]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM Lambda
internaliseFoldLambda InternaliseLambda
internaliseLambda ExpBase Info VName
lam0 [TypeBase Shape NoUniqueness]
nes_ts [TypeBase Shape NoUniqueness]
acc_arr_tps

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

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

  Lambda
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)]
lam_params' (InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM)))
-> InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM))
forall a b. (a -> b) -> a -> b
$ do
    Result
lam_res <- Body (Rep InternaliseM) -> InternaliseM Result
forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind Body (Rep InternaliseM)
BodyT SOACS
lam_body
    [SubExp]
lam_res' <-
      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 chunk function result does not match shape of initial value"
        (ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
lam)
        []
        ((Param (TypeBase Shape NoUniqueness)
 -> TypeBase Shape NoUniqueness)
-> [Param (TypeBase Shape NoUniqueness)]
-> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map Param (TypeBase Shape NoUniqueness) -> TypeBase Shape NoUniqueness
forall t. Typed t => t -> TypeBase Shape NoUniqueness
I.typeOf ([Param (TypeBase Shape NoUniqueness)]
 -> [TypeBase Shape NoUniqueness])
-> [Param (TypeBase Shape NoUniqueness)]
-> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> a -> b
$ Lambda -> [LParam]
forall rep. LambdaT rep -> [LParam rep]
I.lambdaParams Lambda
lam0')
        ((SubExpRes -> SubExp) -> Result -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp Result
lam_res)
    ErrorMsg SubExp
-> SrcLoc
-> [TypeBase Shape NoUniqueness]
-> Result
-> InternaliseM Result
ensureResultShape
      ErrorMsg SubExp
"shape of result does not match shape of initial value"
      (ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
lam0)
      [TypeBase Shape NoUniqueness]
nes_ts
      (Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ( Lambda (Rep InternaliseM)
-> [InternaliseM (Exp (Rep InternaliseM))] -> InternaliseM Result
forall (m :: * -> *).
MonadBuilder m =>
Lambda (Rep m) -> [m (Exp (Rep m))] -> m Result
eLambda Lambda (Rep InternaliseM)
Lambda
lam0' ([InternaliseM (ExpT SOACS)] -> InternaliseM Result)
-> ([SubExp] -> [InternaliseM (ExpT SOACS)])
-> [SubExp]
-> InternaliseM Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SubExp -> InternaliseM (ExpT SOACS))
-> [SubExp] -> [InternaliseM (ExpT SOACS)]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> InternaliseM (ExpT SOACS)
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp ([SubExp] -> InternaliseM Result)
-> [SubExp] -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
              (Param (TypeBase Shape NoUniqueness) -> SubExp)
-> [Param (TypeBase Shape NoUniqueness)] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
I.Var (VName -> SubExp)
-> (Param (TypeBase Shape NoUniqueness) -> VName)
-> Param (TypeBase Shape NoUniqueness)
-> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param (TypeBase Shape NoUniqueness) -> VName
forall dec. Param dec -> VName
paramName) [Param (TypeBase Shape NoUniqueness)]
lam_acc_params [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
lam_res'
          )

  let form :: StreamForm SOACS
form = StreamOrd -> Commutativity -> Lambda -> StreamForm SOACS
forall rep.
StreamOrd -> Commutativity -> Lambda rep -> StreamForm rep
I.Parallel StreamOrd
o Commutativity
comm Lambda
lam0'
  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)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
  String -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [SubExp]
letTupExp' String
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Op SOACS -> ExpT SOACS
forall rep. Op rep -> ExpT rep
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp
-> [VName] -> StreamForm SOACS -> [SubExp] -> Lambda -> SOAC SOACS
forall rep.
SubExp
-> [VName] -> StreamForm rep -> [SubExp] -> Lambda rep -> SOAC rep
I.Stream SubExp
w [VName]
arrs StreamForm SOACS
form ((SubExpRes -> SubExp) -> Result -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp Result
nes) Lambda
lam'

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

  VName
acc_cert_v <- String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"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)
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
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 shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
rowType [TypeBase Shape NoUniqueness]
dest_ts) NoUniqueness
NoUniqueness
  Param (TypeBase Shape NoUniqueness)
acc_p <- String
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (m :: * -> *) dec.
MonadFreshNames m =>
String -> dec -> m (Param dec)
newParam String
"acc_p" TypeBase Shape NoUniqueness
acc_t
  Lambda
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 [VName
-> TypeBase Shape NoUniqueness
-> Param (TypeBase Shape NoUniqueness)
forall dec. VName -> dec -> Param dec
Param 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
    Lambda
lam' <-
      InternaliseLambda
-> ExpBase Info VName -> [SubExp] -> InternaliseM Lambda
internaliseMapLambda InternaliseLambda
internaliseLambda ExpBase Info VName
lam ([SubExp] -> InternaliseM Lambda)
-> [SubExp] -> InternaliseM Lambda
forall a b. (a -> b) -> a -> b
$
        (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map 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
paramName Param (TypeBase Shape NoUniqueness)
acc_p VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: [VName]
bs'
    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)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
bs'
    ([SubExp] -> Result)
-> InternaliseM [SubExp] -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [SubExp] -> Result
subExpsRes (InternaliseM [SubExp] -> InternaliseM Result)
-> (ExpT SOACS -> InternaliseM [SubExp])
-> ExpT SOACS
-> InternaliseM Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [SubExp]
letTupExp' String
"acc_res" (ExpT SOACS -> InternaliseM Result)
-> ExpT SOACS -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
      Op SOACS -> ExpT SOACS
forall rep. Op rep -> ExpT rep
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT 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
paramName Param (TypeBase Shape NoUniqueness)
acc_p VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: [VName]
bs') (Lambda -> ScremaForm SOACS
forall rep. Lambda rep -> ScremaForm rep
I.mapSOAC Lambda
lam')

  Maybe (Lambda, [SubExp])
op' <-
    case Maybe (ExpBase Info VName, ExpBase Info VName)
op of
      Just (ExpBase Info VName
op_lam, ExpBase Info VName
ne) -> do
        [SubExp]
ne' <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"hist_ne" ExpBase Info VName
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)
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, BodyT SOACS
lam_body, [TypeBase Shape NoUniqueness]
lam_rettype) <-
          InternaliseLambda
internaliseLambda ExpBase Info VName
op_lam ([TypeBase Shape NoUniqueness]
 -> InternaliseM
      ([LParam], BodyT SOACS, [TypeBase Shape NoUniqueness]))
-> [TypeBase Shape NoUniqueness]
-> InternaliseM
     ([LParam], BodyT 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 <- String
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (m :: * -> *) dec.
MonadFreshNames m =>
String -> dec -> m (Param dec)
newParam String
"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
op_lam' = [LParam] -> BodyT SOACS -> [TypeBase Shape NoUniqueness] -> Lambda
forall rep.
[LParam rep]
-> BodyT rep -> [TypeBase Shape NoUniqueness] -> LambdaT 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) BodyT SOACS
lam_body [TypeBase Shape NoUniqueness]
lam_rettype
        Maybe (Lambda, [SubExp]) -> InternaliseM (Maybe (Lambda, [SubExp]))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (Lambda, [SubExp])
 -> InternaliseM (Maybe (Lambda, [SubExp])))
-> Maybe (Lambda, [SubExp])
-> InternaliseM (Maybe (Lambda, [SubExp]))
forall a b. (a -> b) -> a -> b
$ (Lambda, [SubExp]) -> Maybe (Lambda, [SubExp])
forall a. a -> Maybe a
Just (Lambda
op_lam', [SubExp]
ne')
      Maybe (ExpBase Info VName, ExpBase Info VName)
Nothing ->
        Maybe (Lambda, [SubExp]) -> InternaliseM (Maybe (Lambda, [SubExp]))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Lambda, [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)
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 (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
$
    String -> Exp (Rep InternaliseM) -> InternaliseM [VName]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [VName]
letTupExp String
desc (Exp (Rep InternaliseM) -> InternaliseM [VName])
-> Exp (Rep InternaliseM) -> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$ [(Shape, [VName], Maybe (Lambda, [SubExp]))]
-> Lambda -> ExpT SOACS
forall rep.
[(Shape, [VName], Maybe (Lambda rep, [SubExp]))]
-> Lambda rep -> ExpT rep
WithAcc [([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
destw], [VName]
dest', Maybe (Lambda, [SubExp])
op')] Lambda
withacc_lam

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

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

internaliseExpToVars :: String -> E.Exp -> InternaliseM [I.VName]
internaliseExpToVars :: String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
desc ExpBase Info VName
e =
  (SubExp -> InternaliseM VName) -> [SubExp] -> InternaliseM [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t 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
=<< String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc ExpBase Info VName
e
  where
    asIdent :: SubExp -> InternaliseM VName
asIdent (I.Var VName
v) = VName -> InternaliseM VName
forall (m :: * -> *) a. Monad m => a -> m a
return VName
v
    asIdent SubExp
se = String -> Exp (Rep InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m VName
letExp String
desc (Exp (Rep InternaliseM) -> InternaliseM VName)
-> Exp (Rep InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
se

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

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

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

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

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

findFuncall ::
  E.AppExp ->
  InternaliseM
    ( E.QualName VName,
      [(E.Exp, Maybe VName)]
    )
findFuncall :: AppExp
-> InternaliseM
     (QualName VName, [(ExpBase Info VName, Maybe VName)])
findFuncall (E.Apply ExpBase Info VName
f ExpBase Info VName
arg (Info (Diet
_, Maybe VName
argext)) SrcLoc
_)
  | E.AppExp AppExp
f_e Info AppRes
_ <- ExpBase Info VName
f = do
    (QualName VName
fname, [(ExpBase Info VName, Maybe VName)]
args) <- AppExp
-> InternaliseM
     (QualName VName, [(ExpBase Info VName, Maybe VName)])
findFuncall AppExp
f_e
    (QualName VName, [(ExpBase Info VName, Maybe VName)])
-> InternaliseM
     (QualName VName, [(ExpBase Info VName, Maybe VName)])
forall (m :: * -> *) a. Monad m => a -> m a
return (QualName VName
fname, [(ExpBase Info VName, Maybe VName)]
args [(ExpBase Info VName, Maybe VName)]
-> [(ExpBase Info VName, Maybe VName)]
-> [(ExpBase Info VName, Maybe VName)]
forall a. [a] -> [a] -> [a]
++ [(ExpBase Info VName
arg, Maybe VName
argext)])
  | E.Var QualName VName
fname Info PatType
_ SrcLoc
_ <- ExpBase Info VName
f =
    (QualName VName, [(ExpBase Info VName, Maybe VName)])
-> InternaliseM
     (QualName VName, [(ExpBase Info VName, Maybe VName)])
forall (m :: * -> *) a. Monad m => a -> m a
return (QualName VName
fname, [(ExpBase Info VName
arg, Maybe VName
argext)])
findFuncall AppExp
e =
  String
-> InternaliseM
     (QualName VName, [(ExpBase Info VName, Maybe VName)])
forall a. HasCallStack => String -> a
error (String
 -> InternaliseM
      (QualName VName, [(ExpBase Info VName, Maybe VName)]))
-> String
-> InternaliseM
     (QualName VName, [(ExpBase Info VName, Maybe VName)])
forall a b. (a -> b) -> a -> b
$ String
"Invalid function expression in application: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ AppExp -> String
forall a. Pretty a => a -> String
pretty 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 -> InternaliseM [ExtType]
bodyExtType :: BodyT 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)
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 ExpBase Info VName
e SrcLoc
_) [TypeBase Shape NoUniqueness]
rowtypes =
  InternaliseLambda
internaliseLambda ExpBase Info VName
e [TypeBase Shape NoUniqueness]
rowtypes
internaliseLambda (E.Lambda [PatBase Info VName]
params ExpBase Info VName
body Maybe (TypeExp VName)
_ (Info (Aliasing
_, StructType
rettype)) SrcLoc
_) [TypeBase Shape NoUniqueness]
rowtypes =
  [PatBase Info VName]
-> [TypeBase Shape NoUniqueness]
-> ([LParam]
    -> InternaliseM
         ([Param (TypeBase Shape NoUniqueness)], BodyT SOACS,
          [TypeBase Shape NoUniqueness]))
-> InternaliseM
     ([Param (TypeBase Shape NoUniqueness)], BodyT SOACS,
      [TypeBase Shape NoUniqueness])
forall a.
[PatBase Info VName]
-> [TypeBase Shape NoUniqueness]
-> ([LParam] -> InternaliseM a)
-> InternaliseM a
bindingLambdaParams [PatBase Info VName]
params [TypeBase Shape NoUniqueness]
rowtypes (([LParam]
  -> InternaliseM
       ([Param (TypeBase Shape NoUniqueness)], BodyT SOACS,
        [TypeBase Shape NoUniqueness]))
 -> InternaliseM
      ([Param (TypeBase Shape NoUniqueness)], BodyT SOACS,
       [TypeBase Shape NoUniqueness]))
-> ([LParam]
    -> InternaliseM
         ([Param (TypeBase Shape NoUniqueness)], BodyT SOACS,
          [TypeBase Shape NoUniqueness]))
-> InternaliseM
     ([Param (TypeBase Shape NoUniqueness)], BodyT SOACS,
      [TypeBase Shape NoUniqueness])
forall a b. (a -> b) -> a -> b
$ \[LParam]
params' -> do
    BodyT SOACS
body' <- String -> ExpBase Info VName -> InternaliseM (BodyT SOACS)
internaliseBody String
"lam" ExpBase Info VName
body
    [TypeBase Shape NoUniqueness]
rettype' <- StructType
-> [ExtType] -> InternaliseM [TypeBase Shape NoUniqueness]
forall shape u.
StructType
-> [TypeBase shape u] -> InternaliseM [TypeBase Shape NoUniqueness]
internaliseLambdaReturnType StructType
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
=<< BodyT SOACS -> InternaliseM [ExtType]
bodyExtType BodyT SOACS
body'
    ([Param (TypeBase Shape NoUniqueness)], BodyT SOACS,
 [TypeBase Shape NoUniqueness])
-> InternaliseM
     ([Param (TypeBase Shape NoUniqueness)], BodyT SOACS,
      [TypeBase Shape NoUniqueness])
forall (m :: * -> *) a. Monad m => a -> m a
return ([Param (TypeBase Shape NoUniqueness)]
[LParam]
params', BodyT SOACS
body', [TypeBase Shape NoUniqueness]
rettype')
internaliseLambda ExpBase Info VName
e [TypeBase Shape NoUniqueness]
_ = String
-> InternaliseM
     ([Param (TypeBase Shape NoUniqueness)], BodyT SOACS,
      [TypeBase Shape NoUniqueness])
forall a. HasCallStack => String -> a
error (String
 -> InternaliseM
      ([Param (TypeBase Shape NoUniqueness)], BodyT SOACS,
       [TypeBase Shape NoUniqueness]))
-> String
-> InternaliseM
     ([Param (TypeBase Shape NoUniqueness)], BodyT SOACS,
      [TypeBase Shape NoUniqueness])
forall a b. (a -> b) -> a -> b
$ String
"internaliseLambda: unexpected expression:\n" String -> String -> String
forall a. [a] -> [a] -> [a]
++ ExpBase Info VName -> String
forall a. Pretty a => a -> String
pretty ExpBase Info VName
e

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

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

    -- Short-circuiting operators are magical.
    handleOps :: [ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM [SubExp])
handleOps [ExpBase Info VName
x, ExpBase Info VName
y] String
"&&" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc (ExpBase Info VName -> InternaliseM [SubExp])
-> ExpBase Info VName -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
        AppExp -> Info AppRes -> ExpBase Info VName
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
E.AppExp
          (ExpBase Info VName
-> ExpBase Info VName -> ExpBase Info VName -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
E.If ExpBase Info VName
x ExpBase Info VName
y (PrimValue -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
E.Literal (Bool -> PrimValue
E.BoolValue Bool
False) SrcLoc
forall a. Monoid a => a
mempty) 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
$ PatType -> [VName] -> AppRes
AppRes (ScalarTypeBase (DimDecl VName) Aliasing -> PatType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar (ScalarTypeBase (DimDecl VName) Aliasing -> PatType)
-> ScalarTypeBase (DimDecl VName) Aliasing -> PatType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) Aliasing
forall dim as. PrimType -> ScalarTypeBase dim as
E.Prim PrimType
E.Bool) [])
    handleOps [ExpBase Info VName
x, ExpBase Info VName
y] String
"||" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
desc (ExpBase Info VName -> InternaliseM [SubExp])
-> ExpBase Info VName -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
        AppExp -> Info AppRes -> ExpBase Info VName
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
E.AppExp
          (ExpBase Info VName
-> ExpBase Info VName -> ExpBase Info VName -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
E.If ExpBase Info VName
x (PrimValue -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
E.Literal (Bool -> PrimValue
E.BoolValue Bool
True) SrcLoc
forall a. Monoid a => a
mempty) ExpBase Info VName
y 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
$ PatType -> [VName] -> AppRes
AppRes (ScalarTypeBase (DimDecl VName) Aliasing -> PatType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar (ScalarTypeBase (DimDecl VName) Aliasing -> PatType)
-> ScalarTypeBase (DimDecl VName) Aliasing -> PatType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) Aliasing
forall dim as. PrimType -> ScalarTypeBase dim as
E.Prim PrimType
E.Bool) [])
    -- Handle equality and inequality specially, to treat the case of
    -- arrays.
    handleOps [ExpBase Info VName
xe, ExpBase Info VName
ye] String
op
      | Just String -> SubExp -> InternaliseM [SubExp]
cmp_f <- String -> Maybe (String -> SubExp -> InternaliseM [SubExp])
forall a (m :: * -> *).
(IsString a, MonadBuilder m, Eq a) =>
a -> Maybe (String -> SubExp -> m [SubExp])
isEqlOp String
op = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc -> do
        [SubExp]
xe' <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"x" ExpBase Info VName
xe
        [SubExp]
ye' <- String -> ExpBase Info VName -> InternaliseM [SubExp]
internaliseExp String
"y" ExpBase Info VName
ye
        [SubExp]
rs <- (SubExp -> SubExp -> InternaliseM SubExp)
-> [SubExp] -> [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM (String -> SubExp -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
(MonadBuilder m, Buildable (Rep m), BuilderOps (Rep m),
 Op (Rep m) ~ SOAC (Rep m)) =>
String -> SubExp -> SubExp -> m SubExp
doComparison String
desc) [SubExp]
xe' [SubExp]
ye'
        String -> SubExp -> InternaliseM [SubExp]
cmp_f String
desc (SubExp -> InternaliseM [SubExp])
-> InternaliseM SubExp -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"eq" (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [SubExp] -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
eAll [SubExp]
rs
      where
        isEqlOp :: a -> Maybe (String -> SubExp -> m [SubExp])
isEqlOp a
"!=" = (String -> SubExp -> m [SubExp])
-> Maybe (String -> SubExp -> m [SubExp])
forall a. a -> Maybe a
Just ((String -> SubExp -> m [SubExp])
 -> Maybe (String -> SubExp -> m [SubExp]))
-> (String -> SubExp -> m [SubExp])
-> Maybe (String -> SubExp -> m [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc SubExp
eq ->
          String -> Exp (Rep m) -> m [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [SubExp]
letTupExp' String
desc (Exp (Rep m) -> m [SubExp]) -> Exp (Rep m) -> m [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
I.Not SubExp
eq
        isEqlOp a
"==" = (String -> SubExp -> m [SubExp])
-> Maybe (String -> SubExp -> m [SubExp])
forall a. a -> Maybe a
Just ((String -> SubExp -> m [SubExp])
 -> Maybe (String -> SubExp -> m [SubExp]))
-> (String -> SubExp -> m [SubExp])
-> Maybe (String -> SubExp -> m [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
_ SubExp
eq ->
          [SubExp] -> m [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [SubExp
eq]
        isEqlOp a
_ = Maybe (String -> SubExp -> m [SubExp])
forall a. Maybe a
Nothing

        doComparison :: String -> SubExp -> SubExp -> m SubExp
doComparison String
desc SubExp
x SubExp
y = do
          TypeBase Shape NoUniqueness
x_t <- SubExp -> m (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
I.subExpType SubExp
x
          TypeBase Shape NoUniqueness
y_t <- SubExp -> m (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 -> String -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
desc (Exp (Rep m) -> m SubExp) -> Exp (Rep m) -> m SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
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) -> m SubExp) -> m [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([SubExp] -> [SubExp] -> [(SubExp, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
x_dims [SubExp]
y_dims) (((SubExp, SubExp) -> m SubExp) -> m [SubExp])
-> ((SubExp, SubExp) -> m SubExp) -> m [SubExp]
forall a b. (a -> b) -> a -> b
$ \(SubExp
x_dim, SubExp
y_dim) ->
                String -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"dim_eq" (Exp (Rep m) -> m SubExp) -> Exp (Rep m) -> m SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
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 <- String -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"shapes_match" (Exp (Rep m) -> m SubExp) -> m (Exp (Rep m)) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [SubExp] -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
eAll [SubExp]
dims_match
              Body (Rep m)
compare_elems_body <- Builder (Rep m) (Body (Rep m)) -> m (Body (Rep m))
forall rep (m :: * -> *) somerep.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
 SameScope somerep rep) =>
Builder rep (Body rep) -> m (Body rep)
runBodyBuilder (Builder (Rep m) (Body (Rep m)) -> m (Body (Rep m)))
-> Builder (Rep m) (Body (Rep m)) -> m (Body (Rep m))
forall a b. (a -> b) -> a -> b
$ do
                -- Flatten both x and y.
                SubExp
x_num_elems <-
                  String
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"x_num_elems"
                    (Exp (Rep m) -> BuilderT (Rep m) (State VNameSource) SubExp)
-> BuilderT (Rep m) (State VNameSource) (Exp (Rep m))
-> BuilderT (Rep m) (State VNameSource) SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinOp
-> SubExp
-> [SubExp]
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Exp (Rep (BuilderT (Rep m) (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' <- String
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m VName
letExp String
"x" (Exp (Rep (BuilderT (Rep m) (State VNameSource)))
 -> BuilderT (Rep m) (State VNameSource) VName)
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
x
                VName
y' <- String
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m VName
letExp String
"x" (Exp (Rep (BuilderT (Rep m) (State VNameSource)))
 -> BuilderT (Rep m) (State VNameSource) VName)
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
y
                VName
x_flat <- String
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m VName
letExp String
"x_flat" (Exp (Rep (BuilderT (Rep m) (State VNameSource)))
 -> BuilderT (Rep m) (State VNameSource) VName)
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
I.Reshape [SubExp -> DimChange SubExp
forall d. d -> DimChange d
I.DimNew SubExp
x_num_elems] VName
x'
                VName
y_flat <- String
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m VName
letExp String
"y_flat" (Exp (Rep (BuilderT (Rep m) (State VNameSource)))
 -> BuilderT (Rep m) (State VNameSource) VName)
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
I.Reshape [SubExp -> DimChange SubExp
forall d. d -> DimChange d
I.DimNew SubExp
x_num_elems] VName
y'

                -- Compare the elements.
                Lambda (Rep m)
cmp_lam <- CmpOp
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Lambda (Rep (BuilderT (Rep m) (State VNameSource))))
forall (m :: * -> *).
(MonadBuilder m, Buildable (Rep m)) =>
CmpOp -> m (Lambda (Rep m))
cmpOpLambda (CmpOp
 -> BuilderT
      (Rep m)
      (State VNameSource)
      (Lambda (Rep (BuilderT (Rep m) (State VNameSource)))))
-> CmpOp
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Lambda (Rep (BuilderT (Rep m) (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 <-
                  String
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m VName
letExp String
"cmps" (Exp (Rep (BuilderT (Rep m) (State VNameSource)))
 -> BuilderT (Rep m) (State VNameSource) VName)
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$
                    Op (Rep m) -> Exp (Rep m)
forall rep. Op rep -> ExpT rep
I.Op (Op (Rep m) -> Exp (Rep m)) -> Op (Rep m) -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$
                      SubExp -> [VName] -> ScremaForm (Rep m) -> SOAC (Rep m)
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
x_num_elems [VName
x_flat, VName
y_flat] (Lambda (Rep m) -> ScremaForm (Rep m)
forall rep. Lambda rep -> ScremaForm rep
I.mapSOAC Lambda (Rep m)
cmp_lam)

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

              String -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"arrays_equal" (Exp (Rep m) -> m SubExp) -> Exp (Rep m) -> m SubExp
forall a b. (a -> b) -> a -> b
$
                SubExp
-> Body (Rep m)
-> Body (Rep m)
-> IfDec (BranchType (Rep m))
-> Exp (Rep m)
forall rep.
SubExp
-> BodyT rep -> BodyT rep -> IfDec (BranchType rep) -> ExpT rep
I.If SubExp
shapes_match Body (Rep m)
compare_elems_body ([SubExp] -> Body (Rep m)
forall rep. Buildable rep => [SubExp] -> Body rep
resultBody [Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
False]) (IfDec (BranchType (Rep m)) -> Exp (Rep m))
-> IfDec (BranchType (Rep m)) -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$
                  [TypeBase Shape NoUniqueness] -> IfDec ExtType
ifCommon [PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.Bool]
    handleOps [ExpBase Info VName
x, ExpBase Info VName
y] String
name
      | Just BinOp
bop <- (BinOp -> Bool) -> [BinOp] -> Maybe BinOp
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((String
name String -> String -> Bool
forall a. Eq a => a -> a -> Bool
==) (String -> Bool) -> (BinOp -> String) -> BinOp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BinOp -> String
forall a. Pretty a => a -> String
pretty) [BinOp
forall a. Bounded a => a
minBound .. BinOp
forall a. Bounded a => a
maxBound :: E.BinOp] =
        (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc -> do
          SubExp
x' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"x" ExpBase Info VName
x
          SubExp
y' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"y" ExpBase Info VName
y
          case (ExpBase Info VName -> PatType
E.typeOf ExpBase Info VName
x, ExpBase Info VName -> PatType
E.typeOf ExpBase Info VName
y) of
            (E.Scalar (E.Prim PrimType
t1), E.Scalar (E.Prim PrimType
t2)) ->
              SrcLoc
-> String
-> BinOp
-> SubExp
-> SubExp
-> PrimType
-> PrimType
-> InternaliseM [SubExp]
internaliseBinOp SrcLoc
loc String
desc BinOp
bop SubExp
x' SubExp
y' PrimType
t1 PrimType
t2
            (PatType, PatType)
_ -> String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error String
"Futhark.Internalise.internaliseExp: non-primitive type in BinOp."
    handleOps [ExpBase Info VName]
_ String
_ = Maybe (String -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

    handleSOACs :: [ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM [SubExp])
handleSOACs [TupLit [ExpBase Info VName
lam, ExpBase Info VName
arr] SrcLoc
_] String
"map" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc -> do
      [VName]
arr' <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"map_arr" ExpBase Info VName
arr
      Lambda
lam' <- InternaliseLambda
-> ExpBase Info VName -> [SubExp] -> InternaliseM Lambda
internaliseMapLambda InternaliseLambda
internaliseLambda ExpBase Info VName
lam ([SubExp] -> InternaliseM Lambda)
-> [SubExp] -> InternaliseM Lambda
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var [VName]
arr'
      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)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arr'
      String -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [SubExp]
letTupExp' String
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
        Op SOACS -> ExpT SOACS
forall rep. Op rep -> ExpT rep
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT 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 [VName]
arr' (Lambda -> ScremaForm SOACS
forall rep. Lambda rep -> ScremaForm rep
I.mapSOAC Lambda
lam')
    handleSOACs [TupLit [ExpBase Info VName
k, ExpBase Info VName
lam, ExpBase Info VName
arr] SrcLoc
_] String
"partition" = do
      Int
k' <- Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int32 -> Int) -> Maybe Int32 -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> Maybe Int32
forall vn. ExpBase Info vn -> Maybe Int32
fromInt32 ExpBase Info VName
k
      (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
_desc -> do
        [VName]
arrs <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"partition_input" ExpBase Info VName
arr
        Lambda
lam' <- InternaliseLambda
-> Int -> ExpBase Info VName -> [SubExp] -> InternaliseM Lambda
internalisePartitionLambda InternaliseLambda
internaliseLambda Int
k' ExpBase Info VName
lam ([SubExp] -> InternaliseM Lambda)
-> [SubExp] -> InternaliseM Lambda
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> [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 -> [VName] -> InternaliseM ([SubExp], [SubExp])
partitionWithSOACS (Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k') Lambda
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 (Signed IntType
Int32)))) SrcLoc
_) = Int32 -> Maybe Int32
forall a. a -> Maybe a
Just (Int32 -> Maybe Int32) -> Int32 -> Maybe Int32
forall a b. (a -> b) -> a -> b
$ Integer -> Int32
forall a. Num a => Integer -> a
fromInteger Integer
k'
        fromInt32 ExpBase Info vn
_ = Maybe Int32
forall a. Maybe a
Nothing
    handleSOACs [TupLit [ExpBase Info VName
lam, ExpBase Info VName
ne, ExpBase Info VName
arr] SrcLoc
_] String
"reduce" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> String
-> (SubExp
    -> Lambda -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS))
-> (ExpBase Info VName, ExpBase Info VName, ExpBase Info VName,
    SrcLoc)
-> InternaliseM [SubExp]
internaliseScanOrReduce String
desc String
"reduce" SubExp
-> Lambda -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)
forall (f :: * -> *) rep.
(Buildable rep, MonadFreshNames f) =>
SubExp -> Lambda rep -> [SubExp] -> [VName] -> f (SOAC rep)
reduce (ExpBase Info VName
lam, ExpBase Info VName
ne, ExpBase Info VName
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 [TupLit [ExpBase Info VName
lam, ExpBase Info VName
ne, ExpBase Info VName
arr] SrcLoc
_] String
"reduce_comm" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> String
-> (SubExp
    -> Lambda -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS))
-> (ExpBase Info VName, ExpBase Info VName, ExpBase Info VName,
    SrcLoc)
-> InternaliseM [SubExp]
internaliseScanOrReduce String
desc String
"reduce" SubExp
-> Lambda -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)
forall (f :: * -> *) rep.
(Buildable rep, MonadFreshNames f) =>
SubExp -> Lambda rep -> [SubExp] -> [VName] -> f (SOAC rep)
reduce (ExpBase Info VName
lam, ExpBase Info VName
ne, ExpBase Info VName
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 [TupLit [ExpBase Info VName
lam, ExpBase Info VName
ne, ExpBase Info VName
arr] SrcLoc
_] String
"scan" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> String
-> (SubExp
    -> Lambda -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS))
-> (ExpBase Info VName, ExpBase Info VName, ExpBase Info VName,
    SrcLoc)
-> InternaliseM [SubExp]
internaliseScanOrReduce String
desc String
"scan" SubExp
-> Lambda -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)
forall (f :: * -> *) rep.
(Buildable rep, MonadFreshNames f) =>
SubExp -> Lambda rep -> [SubExp] -> [VName] -> f (SOAC rep)
reduce (ExpBase Info VName
lam, ExpBase Info VName
ne, ExpBase Info VName
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 [TupLit [ExpBase Info VName
op, ExpBase Info VName
f, ExpBase Info VName
arr] SrcLoc
_] String
"reduce_stream" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> StreamOrd
-> Commutativity
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> InternaliseM [SubExp]
internaliseStreamRed String
desc StreamOrd
InOrder Commutativity
Noncommutative ExpBase Info VName
op ExpBase Info VName
f ExpBase Info VName
arr
    handleSOACs [TupLit [ExpBase Info VName
op, ExpBase Info VName
f, ExpBase Info VName
arr] SrcLoc
_] String
"reduce_stream_per" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> StreamOrd
-> Commutativity
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> InternaliseM [SubExp]
internaliseStreamRed String
desc StreamOrd
Disorder Commutativity
Commutative ExpBase Info VName
op ExpBase Info VName
f ExpBase Info VName
arr
    handleSOACs [TupLit [ExpBase Info VName
f, ExpBase Info VName
arr] SrcLoc
_] String
"map_stream" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> StreamOrd
-> ExpBase Info VName
-> ExpBase Info VName
-> InternaliseM [SubExp]
internaliseStreamMap String
desc StreamOrd
InOrder ExpBase Info VName
f ExpBase Info VName
arr
    handleSOACs [TupLit [ExpBase Info VName
f, ExpBase Info VName
arr] SrcLoc
_] String
"map_stream_per" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> StreamOrd
-> ExpBase Info VName
-> ExpBase Info VName
-> InternaliseM [SubExp]
internaliseStreamMap String
desc StreamOrd
Disorder ExpBase Info VName
f ExpBase Info VName
arr
    handleSOACs [TupLit [ExpBase Info VName
rf, ExpBase Info VName
dest, ExpBase Info VName
op, ExpBase Info VName
ne, ExpBase Info VName
buckets, ExpBase Info VName
img] SrcLoc
_] String
"hist" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> SrcLoc
-> InternaliseM [SubExp]
internaliseHist String
desc ExpBase Info VName
rf ExpBase Info VName
dest ExpBase Info VName
op ExpBase Info VName
ne ExpBase Info VName
buckets ExpBase Info VName
img SrcLoc
loc
    handleSOACs [ExpBase Info VName]
_ String
_ = Maybe (String -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

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

    handleRest :: [ExpBase Info VName]
-> String -> Maybe (String -> InternaliseM [SubExp])
handleRest [E.TupLit [ExpBase Info VName
a, ExpBase Info VName
si, ExpBase Info VName
v] SrcLoc
_] String
"scatter" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ Int
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> String
-> InternaliseM [SubExp]
scatterF Int
1 ExpBase Info VName
a ExpBase Info VName
si ExpBase Info VName
v
    handleRest [E.TupLit [ExpBase Info VName
a, ExpBase Info VName
si, ExpBase Info VName
v] SrcLoc
_] String
"scatter_2d" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ Int
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> String
-> InternaliseM [SubExp]
scatterF Int
2 ExpBase Info VName
a ExpBase Info VName
si ExpBase Info VName
v
    handleRest [E.TupLit [ExpBase Info VName
a, ExpBase Info VName
si, ExpBase Info VName
v] SrcLoc
_] String
"scatter_3d" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ Int
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> String
-> InternaliseM [SubExp]
scatterF Int
3 ExpBase Info VName
a ExpBase Info VName
si ExpBase Info VName
v
    handleRest [E.TupLit [ExpBase Info VName
n, ExpBase Info VName
m, ExpBase Info VName
arr] SrcLoc
_] String
"unflatten" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc -> do
      [VName]
arrs <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"unflatten_arr" ExpBase Info VName
arr
      SubExp
n' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"n" ExpBase Info VName
n
      SubExp
m' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"m" ExpBase Info VName
m
      -- The unflattened dimension needs to have the same number of elements
      -- as the original dimension.
      SubExp
old_dim <- Int -> [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)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
      SubExp
dim_ok <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"dim_ok"
          (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< CmpOp
-> InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
CmpOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eCmpOp
            (PrimType -> CmpOp
I.CmpEq PrimType
I.int64)
            (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 (IntType -> Overflow -> BinOp
I.Mul IntType
Int64 Overflow
I.OverflowUndef) (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
n') (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
m'))
            (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
old_dim)
      Certs
dim_ok_cert <-
        String -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert
          String
"dim_ok_cert"
          SubExp
dim_ok
          ErrorMsg SubExp
"new shape has different number of elements than old shape"
          SrcLoc
loc
      Certs -> InternaliseM [SubExp] -> InternaliseM [SubExp]
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'
          String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
desc (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
            BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
              ShapeChange SubExp -> VName -> BasicOp
I.Reshape (ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp
reshapeOuter [SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimNew SubExp
n', SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimNew SubExp
m'] Int
1 (Shape -> ShapeChange SubExp) -> Shape -> ShapeChange SubExp
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoUniqueness -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape TypeBase Shape NoUniqueness
arr_t) VName
arr'
    handleRest [ExpBase Info VName
arr] String
"flatten" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc -> do
      [VName]
arrs <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"flatten_arr" ExpBase Info VName
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 <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"flat_dim" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Mul IntType
Int64 Overflow
I.OverflowUndef) SubExp
n SubExp
m
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
desc (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
            ShapeChange SubExp -> VName -> BasicOp
I.Reshape (ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp
reshapeOuter [SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimNew SubExp
k] Int
2 (Shape -> ShapeChange SubExp) -> Shape -> ShapeChange SubExp
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoUniqueness -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape TypeBase Shape NoUniqueness
arr_t) VName
arr'
    handleRest [TupLit [ExpBase Info VName
x, ExpBase Info VName
y] SrcLoc
_] String
"concat" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc -> do
      [VName]
xs <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"concat_x" ExpBase Info VName
x
      [VName]
ys <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"concat_y" ExpBase Info VName
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)
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 =
            String -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"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 -> ExpT 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)
mapM (([TypeBase Shape NoUniqueness] -> SubExp)
-> InternaliseM [TypeBase Shape NoUniqueness]
-> InternaliseM SubExp
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)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType) [[VName]
ys]

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

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

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

    scatterF :: Int
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> String
-> InternaliseM [SubExp]
scatterF Int
dim ExpBase Info VName
a ExpBase Info VName
si ExpBase Info VName
v String
desc = do
      [VName]
si' <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"write_arg_i" ExpBase Info VName
si
      [VName]
svs <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"write_arg_v" ExpBase Info VName
v
      [VName]
sas <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"write_arg_a" ExpBase Info VName
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)
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)
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 <-
          String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"write_cmp" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
            BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
              CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
I.int64) SubExp
si_w SubExp
sv_w
        Certs
c <-
          String -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert
            String
"write_cert"
            SubExp
cmp
            ErrorMsg SubExp
"length of index and value array does not match"
            SrcLoc
loc
        Certs -> InternaliseM VName -> InternaliseM VName
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
$
          String -> Exp (Rep InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m VName
letExp (VName -> String
baseString VName
sv String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_write_sv") (Exp (Rep InternaliseM) -> InternaliseM VName)
-> Exp (Rep InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$
            BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
I.Reshape (ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp
reshapeOuter [SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimCoercion SubExp
si_w] Int
1 Shape
sv_shape) VName
sv

      [TypeBase Shape NoUniqueness]
indexType <- (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
-> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall shape u.
ArrayShape shape =>
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)
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)
mapM (\TypeBase Shape NoUniqueness
_ -> String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"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 (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
$ String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"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)
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 (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 shape u.
ArrayShape shape =>
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 shape u.
ArrayShape shape =>
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 VName
-> TypeBase Shape NoUniqueness
-> Param (TypeBase Shape NoUniqueness)
forall dec. VName -> dec -> Param dec
I.Param [VName]
bodyNames [TypeBase Shape NoUniqueness]
paramTypes

      -- This body is pretty boring right now, as every input is exactly the output.
      -- But it can get funky later on if fused with something else.
      BodyT SOACS
body <- Scope SOACS
-> InternaliseM (BodyT SOACS) -> InternaliseM (BodyT SOACS)
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 (BodyT SOACS) -> InternaliseM (BodyT SOACS))
-> (InternaliseM Result -> InternaliseM (BodyT SOACS))
-> InternaliseM Result
-> InternaliseM (BodyT SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InternaliseM Result -> InternaliseM (BodyT SOACS)
forall (m :: * -> *).
MonadBuilder m =>
m Result -> m (Body (Rep m))
buildBody_ (InternaliseM Result -> InternaliseM (BodyT SOACS))
-> InternaliseM Result -> InternaliseM (BodyT 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 (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 ->
          String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"write_res" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
name
        ErrorMsg SubExp
-> SrcLoc
-> [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
lam =
            Lambda :: forall rep.
[LParam rep]
-> BodyT rep -> [TypeBase Shape NoUniqueness] -> LambdaT rep
I.Lambda
              { lambdaParams :: [LParam]
I.lambdaParams = [Param (TypeBase Shape NoUniqueness)]
[LParam]
bodyParams,
                lambdaReturnType :: [TypeBase Shape NoUniqueness]
I.lambdaReturnType = [TypeBase Shape NoUniqueness]
bodyTypes,
                lambdaBody :: BodyT SOACS
I.lambdaBody = BodyT 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
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
      String -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [SubExp]
letTupExp' String
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Op SOACS -> ExpT SOACS
forall rep. Op rep -> ExpT rep
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> Lambda -> [VName] -> [(Shape, Int, VName)] -> SOAC SOACS
forall rep.
SubExp
-> Lambda rep -> [VName] -> [(Shape, Int, VName)] -> SOAC rep
I.Scatter SubExp
si_w Lambda
lam [VName]
sivs ([(Shape, Int, VName)] -> SOAC SOACS)
-> [(Shape, Int, VName)] -> SOAC SOACS
forall a b. (a -> b) -> a -> b
$ [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

flatIndexHelper :: String -> SrcLoc -> E.Exp -> E.Exp -> [(E.Exp, E.Exp)] -> InternaliseM [SubExp]
flatIndexHelper :: String
-> SrcLoc
-> ExpBase Info VName
-> ExpBase Info VName
-> [(ExpBase Info VName, ExpBase Info VName)]
-> InternaliseM [SubExp]
flatIndexHelper String
desc SrcLoc
loc ExpBase Info VName
arr ExpBase Info VName
offset [(ExpBase Info VName, ExpBase Info VName)]
slices = do
  [VName]
arrs <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"arr" ExpBase Info VName
arr
  SubExp
offset' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"offset" ExpBase Info VName
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)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
  SubExp
offset_inbounds_down <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"offset_inbounds_down" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
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 <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"offset_inbounds_up" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
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' <-
    ((ExpBase Info VName, ExpBase Info VName)
 -> InternaliseM (SubExp, SubExp))
-> [(ExpBase Info VName, ExpBase Info VName)]
-> InternaliseM [(SubExp, SubExp)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM
      ( \(ExpBase Info VName
n, ExpBase Info VName
s) -> do
          SubExp
n' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"n" ExpBase Info VName
n
          SubExp
s' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"s" ExpBase Info VName
s
          (SubExp, SubExp) -> InternaliseM (SubExp, SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
n', SubExp
s')
      )
      [(ExpBase Info VName, ExpBase Info VName)]
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 <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"span" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
Int64 Overflow
I.OverflowUndef) SubExp
n (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1)
          SubExp
spn <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"span" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Mul IntType
Int64 Overflow
I.OverflowUndef) SubExp
n_m1 SubExp
s

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

          SubExp
lower' <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"minimum" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
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' <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"maximum" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
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 (m :: * -> *) a. Monad m => a -> m a
return (SubExp
lower', SubExp
upper')
      )
      (SubExp
offset', SubExp
offset')
      [(SubExp, SubExp)]
slices'
  SubExp
min_in_bounds <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"min_in_bounds" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
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 <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"max_in_bounds" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
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 -> String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"inBounds" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
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 <- String -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert String
"bounds_cert" SubExp
all_bounds ([ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg [String -> ErrorMsgPart SubExp
forall a. String -> ErrorMsgPart a
ErrorString (String -> ErrorMsgPart SubExp) -> String -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ String
"Flat slice out of bounds: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SubExp -> String
forall a. Pretty a => a -> String
pretty SubExp
old_dim String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" and " String -> String -> String
forall a. [a] -> [a] -> [a]
++ [(SubExp, SubExp)] -> String
forall a. Pretty a => a -> String
pretty [(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 (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' ->
      String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
desc (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
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 :: String
-> SrcLoc
-> ExpBase Info VName
-> ExpBase Info VName
-> [ExpBase Info VName]
-> ExpBase Info VName
-> InternaliseM [SubExp]
flatUpdateHelper String
desc SrcLoc
loc ExpBase Info VName
arr1 ExpBase Info VName
offset [ExpBase Info VName]
slices ExpBase Info VName
arr2 = do
  [VName]
arrs1 <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"arr" ExpBase Info VName
arr1
  SubExp
offset' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"offset" ExpBase Info VName
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)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs1
  SubExp
offset_inbounds_down <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"offset_inbounds_down" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
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 <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"offset_inbounds_up" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
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 <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"arr" ExpBase Info VName
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)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs2
  [(SubExp, SubExp)]
slices' <-
    ((ExpBase Info VName, Int) -> InternaliseM (SubExp, SubExp))
-> [(ExpBase Info VName, Int)] -> InternaliseM [(SubExp, SubExp)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM
      ( \(ExpBase Info VName
s, Int
i) -> do
          SubExp
s' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"s" ExpBase Info VName
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 (m :: * -> *) a. Monad m => a -> m a
return (SubExp
n, SubExp
s')
      )
      ([(ExpBase Info VName, Int)] -> InternaliseM [(SubExp, SubExp)])
-> [(ExpBase Info VName, Int)] -> InternaliseM [(SubExp, SubExp)]
forall a b. (a -> b) -> a -> b
$ [ExpBase Info VName] -> [Int] -> [(ExpBase Info VName, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [ExpBase Info VName]
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 <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"span" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
Int64 Overflow
I.OverflowUndef) SubExp
n (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1)
          SubExp
spn <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"span" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Mul IntType
Int64 Overflow
I.OverflowUndef) SubExp
n_m1 SubExp
s

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

          SubExp
lower' <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"minimum" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
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' <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"maximum" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
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 (m :: * -> *) a. Monad m => a -> m a
return (SubExp
lower', SubExp
upper')
      )
      (SubExp
offset', SubExp
offset')
      [(SubExp, SubExp)]
slices'
  SubExp
min_in_bounds <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"min_in_bounds" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
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 <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"max_in_bounds" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
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 -> String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"inBounds" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
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 <- String -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert String
"bounds_cert" SubExp
all_bounds ([ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg [String -> ErrorMsgPart SubExp
forall a. String -> ErrorMsgPart a
ErrorString (String -> ErrorMsgPart SubExp) -> String -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ String
"Flat slice out of bounds: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SubExp -> String
forall a. Pretty a => a -> String
pretty SubExp
old_dim String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" and " String -> String -> String
forall a. [a] -> [a] -> [a]
++ [(SubExp, SubExp)] -> String
forall a. Pretty a => a -> String
pretty [(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 (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') ->
      String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
desc (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ VName -> FlatSlice SubExp -> VName -> BasicOp
I.FlatUpdate VName
arr1' FlatSlice SubExp
slice VName
arr2'

funcall ::
  String ->
  QualName VName ->
  [SubExp] ->
  SrcLoc ->
  InternaliseM ([SubExp], [I.ExtType])
funcall :: String
-> QualName VName
-> [SubExp]
-> SrcLoc
-> InternaliseM ([SubExp], [ExtType])
funcall String
desc (QualName [VName]
_ VName
fname) [SubExp]
args SrcLoc
loc = do
  ([VName]
shapes, [DeclType]
value_paramts, [Param DeclType]
fun_params, [(SubExp, TypeBase Shape NoUniqueness)] -> Maybe [DeclExtType]
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)
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
args

  [SubExp]
shapeargs <- [VName]
-> [FParam]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM [SubExp]
argShapes [VName]
shapes [Param DeclType]
[FParam]
fun_params [TypeBase Shape NoUniqueness]
argts
  let diets :: [Diet]
diets =
        Int -> Diet -> [Diet]
forall a. Int -> a -> [a]
replicate ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
shapeargs) Diet
I.ObservePrim
          [Diet] -> [Diet] -> [Diet]
forall a. [a] -> [a] -> [a]
++ (DeclType -> Diet) -> [DeclType] -> [Diet]
forall a b. (a -> b) -> [a] -> [b]
map DeclType -> Diet
forall shape. TypeBase shape Uniqueness -> Diet
I.diet [DeclType]
value_paramts
  [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)
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 [DeclExtType]
rettype_fun ([(SubExp, TypeBase Shape NoUniqueness)] -> Maybe [DeclExtType])
-> [(SubExp, TypeBase Shape NoUniqueness)] -> Maybe [DeclExtType]
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 [DeclExtType]
Nothing ->
      String -> InternaliseM ([SubExp], [ExtType])
forall a. HasCallStack => String -> a
error (String -> InternaliseM ([SubExp], [ExtType]))
-> String -> InternaliseM ([SubExp], [ExtType])
forall a b. (a -> b) -> a -> b
$
        [String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
          [ String
"Cannot apply ",
            VName -> String
forall a. Pretty a => a -> String
pretty VName
fname,
            String
" to ",
            Int -> String
forall a. Show a => a -> String
show ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
args'),
            String
" arguments\n ",
            [SubExp] -> String
forall a. Pretty a => a -> String
pretty [SubExp]
args',
            String
"\nof types\n ",
            [TypeBase Shape NoUniqueness] -> String
forall a. Pretty a => a -> String
pretty [TypeBase Shape NoUniqueness]
argts',
            String
"\nFunction has ",
            Int -> String
forall a. Show a => a -> String
show ([Param DeclType] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Param DeclType]
fun_params),
            String
" parameters\n ",
            [Param DeclType] -> String
forall a. Pretty a => a -> String
pretty [Param DeclType]
fun_params
          ]
    Just [DeclExtType]
ts -> do
      Safety
safety <- InternaliseM Safety
askSafety
      Attrs
attrs <- (InternaliseEnv -> Attrs) -> InternaliseM Attrs
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks InternaliseEnv -> Attrs
envAttrs
      [SubExp]
ses <-
        Attrs -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. MonadBuilder m => Attrs -> m a -> m a
attributing Attrs
attrs (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
          String -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [SubExp]
letTupExp' String
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
            Name
-> [(SubExp, Diet)]
-> [RetType SOACS]
-> (Safety, SrcLoc, [SrcLoc])
-> ExpT SOACS
forall rep.
Name
-> [(SubExp, Diet)]
-> [RetType rep]
-> (Safety, SrcLoc, [SrcLoc])
-> ExpT rep
I.Apply (VName -> Name
internaliseFunName VName
fname) ([SubExp] -> [Diet] -> [(SubExp, Diet)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
args' [Diet]
diets) [DeclExtType]
[RetType SOACS]
ts (Safety
safety, SrcLoc
loc, [SrcLoc]
forall a. Monoid a => a
mempty)
      ([SubExp], [ExtType]) -> InternaliseM ([SubExp], [ExtType])
forall (m :: * -> *) a. Monad m => a -> m a
return ([SubExp]
ses, (DeclExtType -> ExtType) -> [DeclExtType] -> [ExtType]
forall a b. (a -> b) -> [a] -> [b]
map DeclExtType -> ExtType
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl [DeclExtType]
ts)

-- Bind existential names defined by an expression, based on the
-- concrete values that expression evaluated to.  This most
-- importantly should be done after function calls, but also
-- everything else that can produce existentials in the source
-- language.
bindExtSizes :: AppRes -> [SubExp] -> InternaliseM ()
bindExtSizes :: AppRes -> [SubExp] -> InternaliseM ()
bindExtSizes (AppRes PatType
ret [VName]
retext) [SubExp]
ses = do
  [DeclExtType]
ts <- StructType -> InternaliseM [DeclExtType]
internaliseType (StructType -> InternaliseM [DeclExtType])
-> StructType -> InternaliseM [DeclExtType]
forall a b. (a -> b) -> a -> b
$ PatType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatType
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)
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
ses

  let combine :: DeclExtType -> TypeBase Shape NoUniqueness -> Map VName SubExp
combine DeclExtType
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' (DeclExtType -> [Ext SubExp]
forall u. TypeBase ExtShape u -> [Ext SubExp]
arrayExtDims DeclExtType
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 (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
$ (DeclExtType -> TypeBase Shape NoUniqueness -> Map VName SubExp)
-> [DeclExtType]
-> [TypeBase Shape NoUniqueness]
-> [Map VName SubExp]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith DeclExtType -> TypeBase Shape NoUniqueness -> Map VName SubExp
combine [DeclExtType]
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 -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se

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

-- Implement partitioning using maps, scans and writes.
partitionWithSOACS :: Int -> I.Lambda -> [I.VName] -> InternaliseM ([I.SubExp], [I.SubExp])
partitionWithSOACS :: Int -> Lambda -> [VName] -> InternaliseM ([SubExp], [SubExp])
partitionWithSOACS Int
k Lambda
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)
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 <- String -> Exp (Rep InternaliseM) -> InternaliseM [VName]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [VName]
letTupExp String
"increments" (Exp (Rep InternaliseM) -> InternaliseM [VName])
-> Exp (Rep InternaliseM) -> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$ Op SOACS -> ExpT SOACS
forall rep. Op rep -> ExpT rep
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT 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 [VName]
arrs (Lambda -> ScremaForm SOACS
forall rep. Lambda rep -> ScremaForm rep
mapSOAC Lambda
lam)
  (VName
classes, [VName]
increments) <- case [VName]
classes_and_increments of
    VName
classes : [VName]
increments -> (VName, [VName]) -> InternaliseM (VName, [VName])
forall (m :: * -> *) a. Monad m => a -> m a
return (VName
classes, Int -> [VName] -> [VName]
forall a. Int -> [a] -> [a]
take Int
k [VName]
increments)
    [VName]
_ -> String -> InternaliseM (VName, [VName])
forall a. HasCallStack => String -> a
error String
"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
$ VName
-> TypeBase Shape NoUniqueness
-> Param (TypeBase Shape NoUniqueness)
forall dec. VName -> dec -> Param dec
I.Param (VName
 -> TypeBase Shape NoUniqueness
 -> Param (TypeBase Shape NoUniqueness))
-> InternaliseM VName
-> InternaliseM
     (TypeBase Shape NoUniqueness
      -> Param (TypeBase Shape NoUniqueness))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"x" InternaliseM
  (TypeBase Shape NoUniqueness
   -> Param (TypeBase Shape NoUniqueness))
-> InternaliseM (TypeBase Shape NoUniqueness)
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TypeBase Shape NoUniqueness
-> InternaliseM (TypeBase Shape NoUniqueness)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (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
$ VName
-> TypeBase Shape NoUniqueness
-> Param (TypeBase Shape NoUniqueness)
forall dec. VName -> dec -> Param dec
I.Param (VName
 -> TypeBase Shape NoUniqueness
 -> Param (TypeBase Shape NoUniqueness))
-> InternaliseM VName
-> InternaliseM
     (TypeBase Shape NoUniqueness
      -> Param (TypeBase Shape NoUniqueness))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"y" InternaliseM
  (TypeBase Shape NoUniqueness
   -> Param (TypeBase Shape NoUniqueness))
-> InternaliseM (TypeBase Shape NoUniqueness)
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TypeBase Shape NoUniqueness
-> InternaliseM (TypeBase Shape NoUniqueness)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
  BodyT SOACS
add_lam_body <- Builder SOACS (BodyT SOACS) -> InternaliseM (BodyT SOACS)
forall rep (m :: * -> *) somerep.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
 SameScope somerep rep) =>
Builder rep (Body rep) -> m (Body rep)
runBodyBuilder (Builder SOACS (BodyT SOACS) -> InternaliseM (BodyT SOACS))
-> Builder SOACS (BodyT SOACS) -> InternaliseM (BodyT SOACS)
forall a b. (a -> b) -> a -> b
$
    Scope SOACS
-> Builder SOACS (BodyT SOACS) -> Builder SOACS (BodyT SOACS)
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 (BodyT SOACS) -> Builder SOACS (BodyT SOACS))
-> Builder SOACS (BodyT SOACS) -> Builder SOACS (BodyT SOACS)
forall a b. (a -> b) -> a -> b
$
      ([SubExp] -> BodyT SOACS)
-> BuilderT SOACS (State VNameSource) [SubExp]
-> Builder SOACS (BodyT SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [SubExp] -> BodyT SOACS
forall rep. Buildable rep => [SubExp] -> Body rep
resultBody (BuilderT SOACS (State VNameSource) [SubExp]
 -> Builder SOACS (BodyT SOACS))
-> BuilderT SOACS (State VNameSource) [SubExp]
-> Builder SOACS (BodyT SOACS)
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) ->
          String
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"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 -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
              BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp
                (IntType -> Overflow -> BinOp
I.Add IntType
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
add_lam =
        Lambda :: forall rep.
[LParam rep]
-> BodyT rep -> [TypeBase Shape NoUniqueness] -> LambdaT rep
I.Lambda
          { lambdaBody :: BodyT SOACS
I.lambdaBody = BodyT SOACS
add_lam_body,
            lambdaParams :: [LParam]
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 (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 -> [SubExp] -> Scan SOACS
forall rep. Lambda rep -> [SubExp] -> Scan rep
I.Scan Lambda
add_lam [SubExp]
nes]
  [VName]
all_offsets <- String -> Exp (Rep InternaliseM) -> InternaliseM [VName]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [VName]
letTupExp String
"offsets" (Exp (Rep InternaliseM) -> InternaliseM [VName])
-> Exp (Rep InternaliseM) -> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$ Op SOACS -> ExpT SOACS
forall rep. Op rep -> ExpT rep
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT 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 [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 <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"last_index" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
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)
  BodyT SOACS
nonempty_body <- Builder SOACS (BodyT SOACS) -> InternaliseM (BodyT SOACS)
forall rep (m :: * -> *) somerep.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
 SameScope somerep rep) =>
Builder rep (Body rep) -> m (Body rep)
runBodyBuilder (Builder SOACS (BodyT SOACS) -> InternaliseM (BodyT SOACS))
-> Builder SOACS (BodyT SOACS) -> InternaliseM (BodyT SOACS)
forall a b. (a -> b) -> a -> b
$
    ([SubExp] -> BodyT SOACS)
-> BuilderT SOACS (State VNameSource) [SubExp]
-> Builder SOACS (BodyT SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [SubExp] -> BodyT SOACS
forall rep. Buildable rep => [SubExp] -> Body rep
resultBody (BuilderT SOACS (State VNameSource) [SubExp]
 -> Builder SOACS (BodyT SOACS))
-> BuilderT SOACS (State VNameSource) [SubExp]
-> Builder SOACS (BodyT SOACS)
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 ->
        String
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"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 -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
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]
  let empty_body :: BodyT SOACS
empty_body = [SubExp] -> BodyT SOACS
forall rep. Buildable rep => [SubExp] -> Body rep
resultBody ([SubExp] -> BodyT SOACS) -> [SubExp] -> BodyT SOACS
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 <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"is_empty" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
CmpEq PrimType
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 <-
    String -> Exp (Rep InternaliseM) -> InternaliseM [VName]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [VName]
letTupExp String
"partition_size" (Exp (Rep InternaliseM) -> InternaliseM [VName])
-> Exp (Rep InternaliseM) -> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$
      SubExp
-> BodyT SOACS
-> BodyT SOACS
-> IfDec (BranchType SOACS)
-> ExpT SOACS
forall rep.
SubExp
-> BodyT rep -> BodyT rep -> IfDec (BranchType rep) -> ExpT rep
I.If SubExp
is_empty BodyT SOACS
empty_body BodyT SOACS
nonempty_body (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
        [TypeBase Shape NoUniqueness] -> IfDec ExtType
ifCommon ([TypeBase Shape NoUniqueness] -> IfDec ExtType)
-> [TypeBase Shape NoUniqueness] -> IfDec ExtType
forall a b. (a -> b) -> a -> b
$ 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

  -- 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 ->
    String -> Exp (Rep InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m VName
letExp String
"partition_dest" (Exp (Rep InternaliseM) -> InternaliseM VName)
-> Exp (Rep InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
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
write_lam <- do
    Param (TypeBase Shape NoUniqueness)
c_param <- VName
-> TypeBase Shape NoUniqueness
-> Param (TypeBase Shape NoUniqueness)
forall dec. VName -> dec -> Param dec
I.Param (VName
 -> TypeBase Shape NoUniqueness
 -> Param (TypeBase Shape NoUniqueness))
-> InternaliseM VName
-> InternaliseM
     (TypeBase Shape NoUniqueness
      -> Param (TypeBase Shape NoUniqueness))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"c" InternaliseM
  (TypeBase Shape NoUniqueness
   -> Param (TypeBase Shape NoUniqueness))
-> InternaliseM (TypeBase Shape NoUniqueness)
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TypeBase Shape NoUniqueness
-> InternaliseM (TypeBase Shape NoUniqueness)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (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
$ VName
-> TypeBase Shape NoUniqueness
-> Param (TypeBase Shape NoUniqueness)
forall dec. VName -> dec -> Param dec
I.Param (VName
 -> TypeBase Shape NoUniqueness
 -> Param (TypeBase Shape NoUniqueness))
-> InternaliseM VName
-> InternaliseM
     (TypeBase Shape NoUniqueness
      -> Param (TypeBase Shape NoUniqueness))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"offset" InternaliseM
  (TypeBase Shape NoUniqueness
   -> Param (TypeBase Shape NoUniqueness))
-> InternaliseM (TypeBase Shape NoUniqueness)
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TypeBase Shape NoUniqueness
-> InternaliseM (TypeBase Shape NoUniqueness)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
    [Param (TypeBase Shape NoUniqueness)]
value_params <- [TypeBase Shape NoUniqueness]
-> (TypeBase Shape NoUniqueness
    -> InternaliseM (Param (TypeBase Shape NoUniqueness)))
-> InternaliseM [Param (TypeBase Shape NoUniqueness)]
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 (Param (TypeBase Shape NoUniqueness)))
 -> InternaliseM [Param (TypeBase Shape NoUniqueness)])
-> (TypeBase Shape NoUniqueness
    -> InternaliseM (Param (TypeBase Shape NoUniqueness)))
-> InternaliseM [Param (TypeBase Shape NoUniqueness)]
forall a b. (a -> b) -> a -> b
$ \TypeBase Shape NoUniqueness
arr_t ->
      VName
-> TypeBase Shape NoUniqueness
-> Param (TypeBase Shape NoUniqueness)
forall dec. VName -> dec -> Param dec
I.Param (VName
 -> TypeBase Shape NoUniqueness
 -> Param (TypeBase Shape NoUniqueness))
-> InternaliseM VName
-> InternaliseM
     (TypeBase Shape NoUniqueness
      -> Param (TypeBase Shape NoUniqueness))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"v" InternaliseM
  (TypeBase Shape NoUniqueness
   -> Param (TypeBase Shape NoUniqueness))
-> InternaliseM (TypeBase Shape NoUniqueness)
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TypeBase Shape NoUniqueness
-> InternaliseM (TypeBase Shape NoUniqueness)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
I.rowType TypeBase Shape NoUniqueness
arr_t)
    (SubExp
offset, Stms SOACS
offset_stms) <-
      InternaliseM SubExp
-> InternaliseM (SubExp, 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] -> 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]
offset_params
    Lambda -> InternaliseM Lambda
forall (m :: * -> *) a. Monad m => a -> m a
return
      Lambda :: forall rep.
[LParam rep]
-> BodyT rep -> [TypeBase Shape NoUniqueness] -> LambdaT rep
I.Lambda
        { lambdaParams :: [LParam]
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 (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 shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
I.rowType [TypeBase Shape NoUniqueness]
arr_ts,
          lambdaBody :: BodyT SOACS
I.lambdaBody =
            Stms SOACS -> Result -> BodyT SOACS
forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody Stms SOACS
offset_stms (Result -> BodyT SOACS) -> Result -> BodyT SOACS
forall a b. (a -> b) -> a -> b
$
              Int -> SubExpRes -> Result
forall a. Int -> a -> [a]
replicate ([TypeBase Shape NoUniqueness] -> 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)
        }
  [VName]
results <-
    String -> Exp (Rep InternaliseM) -> InternaliseM [VName]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [VName]
letTupExp String
"partition_res" (Exp (Rep InternaliseM) -> InternaliseM [VName])
-> Exp (Rep InternaliseM) -> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$
      Op SOACS -> ExpT SOACS
forall rep. Op rep -> ExpT rep
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
        SubExp -> Lambda -> [VName] -> [(Shape, Int, VName)] -> SOAC SOACS
forall rep.
SubExp
-> Lambda rep -> [VName] -> [(Shape, Int, VName)] -> SOAC rep
I.Scatter
          SubExp
w
          Lambda
write_lam
          (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)] -> SOAC SOACS)
-> [(Shape, Int, VName)] -> SOAC SOACS
forall a b. (a -> b) -> a -> b
$ [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
Shape [SubExp
w]) (Int -> [Int]
forall a. a -> [a]
repeat Int
1) [VName]
blanks
  SubExp
sizes' <-
    String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"partition_sizes" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
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 (m :: * -> *) a. Monad m => a -> m a
return ((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] ->
      InternaliseM SubExp
    mkOffsetLambdaBody :: [SubExp] -> SubExp -> Int -> [LParam] -> InternaliseM SubExp
mkOffsetLambdaBody [SubExp]
_ SubExp
_ Int
_ [] =
      SubExp -> InternaliseM SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> InternaliseM SubExp) -> SubExp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (-Int64
1 :: Int64)
    mkOffsetLambdaBody [SubExp]
sizes SubExp
c Int
i (LParam
p : [LParam]
ps) = do
      SubExp
is_this_one <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"is_this_one" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
CmpEq PrimType
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] -> InternaliseM SubExp
mkOffsetLambdaBody [SubExp]
sizes SubExp
c (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [LParam]
ps
      SubExp
this_one <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"this_offset"
          (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinOp
-> SubExp -> [SubExp] -> InternaliseM (Exp (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
p) SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
take Int
i [SubExp]
sizes)
      String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"total_res" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
        SubExp
-> BodyT SOACS
-> BodyT SOACS
-> IfDec (BranchType SOACS)
-> ExpT SOACS
forall rep.
SubExp
-> BodyT rep -> BodyT rep -> IfDec (BranchType rep) -> ExpT rep
I.If
          SubExp
is_this_one
          ([SubExp] -> BodyT SOACS
forall rep. Buildable rep => [SubExp] -> Body rep
resultBody [SubExp
this_one])
          ([SubExp] -> BodyT SOACS
forall rep. Buildable rep => [SubExp] -> Body rep
resultBody [SubExp
next_one])
          (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ [TypeBase Shape NoUniqueness] -> IfDec ExtType
ifCommon [PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64]

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

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

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