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

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

import Control.Monad.Reader
import 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 (stmPattern)
import Futhark.Internalise.AccurateSizes
import Futhark.Internalise.Bindings
import Futhark.Internalise.Defunctionalise as Defunctionalise
import Futhark.Internalise.Defunctorise as Defunctorise
import Futhark.Internalise.Lambdas
import Futhark.Internalise.LiftLambdas as LiftLambdas
import Futhark.Internalise.Monad as I
import Futhark.Internalise.Monomorphise as Monomorphise
import Futhark.Internalise.TypesValues
import Futhark.Transform.Rename as I
import Futhark.Util (splitAt3)
import Futhark.Util.Pretty (prettyOneLine)
import Language.Futhark as E hiding (TypeArg)
import Language.Futhark.Semantic (Imports)

-- | Convert a program in source Futhark to a program in the Futhark
-- core language.
internaliseProg ::
  MonadFreshNames m =>
  Bool ->
  Imports ->
  m (I.Prog SOACS)
internaliseProg :: forall (m :: * -> *).
MonadFreshNames m =>
Bool -> Imports -> m (Prog SOACS)
internaliseProg Bool
always_safe Imports
prog = do
  [Dec]
prog_decs <- Imports -> m [Dec]
forall (m :: * -> *). MonadFreshNames m => Imports -> m [Dec]
Defunctorise.transformProg Imports
prog
  [ValBind]
prog_decs' <- [Dec] -> m [ValBind]
forall (m :: * -> *). MonadFreshNames m => [Dec] -> m [ValBind]
Monomorphise.transformProg [Dec]
prog_decs
  [ValBind]
prog_decs'' <- [ValBind] -> m [ValBind]
forall (m :: * -> *). MonadFreshNames m => [ValBind] -> m [ValBind]
LiftLambdas.transformProg [ValBind]
prog_decs'
  [ValBind]
prog_decs''' <- [ValBind] -> m [ValBind]
forall (m :: * -> *). MonadFreshNames m => [ValBind] -> m [ValBind]
Defunctionalise.transformProg [ValBind]
prog_decs''
  (Stms SOACS
consts, [FunDef SOACS]
funs) <-
    Bool -> InternaliseM () -> m (Stms SOACS, [FunDef SOACS])
forall (m :: * -> *).
MonadFreshNames m =>
Bool -> InternaliseM () -> m (Stms SOACS, [FunDef SOACS])
runInternaliseM Bool
always_safe ([ValBind] -> InternaliseM ()
internaliseValBinds [ValBind]
prog_decs''')
  Prog SOACS -> m (Prog SOACS)
forall lore (m :: * -> *).
(Renameable lore, MonadFreshNames m) =>
Prog lore -> m (Prog lore)
I.renameProg (Prog SOACS -> m (Prog SOACS)) -> Prog SOACS -> m (Prog SOACS)
forall a b. (a -> b) -> a -> b
$ Stms SOACS -> [FunDef SOACS] -> Prog SOACS
forall lore. Stms lore -> [FunDef lore] -> Prog lore
I.Prog Stms SOACS
consts [FunDef SOACS]
funs

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

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

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

internaliseFunName :: VName -> Name
internaliseFunName :: VName -> Name
internaliseFunName = [Char] -> Name
nameFromString ([Char] -> Name) -> (VName -> [Char]) -> VName -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> [Char]
forall a. Pretty a => a -> [Char]
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 [PatternBase Info VName]
params Exp
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]
-> [PatternBase Info VName]
-> ([FParam] -> [[FParam]] -> InternaliseM ())
-> InternaliseM ()
forall a.
[TypeParamBase VName]
-> [PatternBase Info VName]
-> ([FParam] -> [[FParam]] -> InternaliseM a)
-> InternaliseM a
bindingFParams [TypeParamBase VName]
tparams [PatternBase 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 ([SubExp], [DeclExtType])
-> InternaliseM (Body (Lore InternaliseM), [DeclExtType])
forall (m :: * -> *) a.
MonadBinder m =>
m ([SubExp], a) -> m (Body (Lore m), a)
buildBody (InternaliseM ([SubExp], [DeclExtType])
 -> InternaliseM (Body (Lore InternaliseM), [DeclExtType]))
-> InternaliseM ([SubExp], [DeclExtType])
-> InternaliseM (Body (Lore InternaliseM), [DeclExtType])
forall a b. (a -> b) -> a -> b
$ do
        [SubExp]
body_res <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp (VName -> [Char]
baseString VName
fname [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"_res") Exp
body
        [DeclExtType]
rettype_bad <-
          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
        let rettype' :: [DeclExtType]
rettype' = [DeclExtType] -> [DeclExtType]
forall {u}.
[TypeBase (ShapeBase ExtSize) u]
-> [TypeBase (ShapeBase ExtSize) u]
zeroExts [DeclExtType]
rettype_bad
        [SubExp]
body_res' <-
          ErrorMsg SubExp
-> SrcLoc -> [ExtType] -> [SubExp] -> InternaliseM [SubExp]
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') [SubExp]
body_res
        ([SubExp], [DeclExtType]) -> InternaliseM ([SubExp], [DeclExtType])
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([SubExp]
body_res', [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 lore.
Maybe EntryPoint
-> Attrs
-> Name
-> [RetType lore]
-> [FParam lore]
-> BodyT lore
-> FunDef lore
I.FunDef
              Maybe EntryPoint
forall a. Maybe a
Nothing
              ([AttrInfo] -> Attrs
internaliseAttrs [AttrInfo]
attrs)
              (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 (ShapeBase ExtSize) u]
-> [TypeBase (ShapeBase ExtSize) u]
zeroExts [TypeBase (ShapeBase ExtSize) u]
ts = [TypeBase (ShapeBase ExtSize) u]
-> [TypeBase (ShapeBase ExtSize) u]
-> [TypeBase (ShapeBase ExtSize) u]
forall u.
[TypeBase (ShapeBase ExtSize) u]
-> [TypeBase (ShapeBase ExtSize) u]
-> [TypeBase (ShapeBase ExtSize) u]
generaliseExtTypes [TypeBase (ShapeBase ExtSize) u]
ts [TypeBase (ShapeBase ExtSize) 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 [PatternBase Info VName]
params Exp
_ Maybe DocComment
_ [AttrInfo]
attrs SrcLoc
loc) = ValBind
vb
  [TypeParamBase VName]
-> [PatternBase Info VName]
-> ([FParam] -> [[FParam]] -> InternaliseM ())
-> InternaliseM ()
forall a.
[TypeParamBase VName]
-> [PatternBase Info VName]
-> ([FParam] -> [[FParam]] -> InternaliseM a)
-> InternaliseM a
bindingFParams [TypeParamBase VName]
tparams [PatternBase 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' = [(EntryType, [FParam])]
-> (EntryType, [[DeclExtType]]) -> EntryPoint
entryPoint ([EntryType]
-> [[Param DeclType]] -> [(EntryType, [Param DeclType])]
forall a b. [a] -> [b] -> [(a, b)]
zip [EntryType]
e_paramts [[Param DeclType]]
[[FParam]]
params') (EntryType
e_rettype, [[DeclExtType]]
entry_rettype)
        args :: [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 <- InternaliseM [SubExp] -> InternaliseM (Body (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
m [SubExp] -> m (Body (Lore m))
buildBody_ (InternaliseM [SubExp] -> InternaliseM (Body (Lore InternaliseM)))
-> InternaliseM [SubExp] -> InternaliseM (Body (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ do
      -- Special case the (rare) situation where the entry point is
      -- not a function.
      Maybe [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
<$> [Char]
-> QualName VName
-> [SubExp]
-> SrcLoc
-> InternaliseM ([SubExp], [ExtType])
funcall [Char]
"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 (ShapeBase ExtSize) 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
      [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
$ [SubExp]
ctx [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
vals

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

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

    entryPointType :: (EntryType, t a) -> [EntryPointType]
entryPointType (EntryType
t, t a
ts)
      | E.Scalar (E.Prim E.Unsigned {}) <- EntryType -> StructType
E.entryType EntryType
t =
        [EntryPointType
I.TypeUnsigned]
      | E.Array ()
_ Uniqueness
_ (E.Prim E.Unsigned {}) ShapeDecl (DimDecl VName)
_ <- EntryType -> StructType
E.entryType EntryType
t =
        [EntryPointType
I.TypeUnsigned]
      | E.Scalar E.Prim {} <- EntryType -> StructType
E.entryType EntryType
t =
        [EntryPointType
I.TypeDirect]
      | E.Array ()
_ Uniqueness
_ E.Prim {} ShapeDecl (DimDecl VName)
_ <- EntryType -> StructType
E.entryType EntryType
t =
        [EntryPointType
I.TypeDirect]
      | Bool
otherwise =
        [[Char] -> Int -> EntryPointType
I.TypeOpaque [Char]
desc (Int -> EntryPointType) -> Int -> EntryPointType
forall a b. (a -> b) -> a -> b
$ t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length t a
ts]
      where
        desc :: [Char]
desc = [Char]
-> (TypeExp VName -> [Char]) -> Maybe (TypeExp VName) -> [Char]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (TypeBase () () -> [Char]
forall a. Pretty a => a -> [Char]
prettyOneLine TypeBase () ()
t') TypeExp VName -> [Char]
forall {vn}. (Eq vn, IsName vn) => TypeExp vn -> [Char]
typeExpOpaqueName (Maybe (TypeExp VName) -> [Char])
-> Maybe (TypeExp VName) -> [Char]
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 -> [Char]
typeExpOpaqueName (TEApply TypeExp vn
te TypeArgExpDim {} SrcLoc
_) =
      TypeExp vn -> [Char]
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 [Char]
"arr_" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ TypeExp vn -> [Char]
typeExpOpaqueName TypeExp vn
te'
            [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"_"
            [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
d)
            [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"d"
    typeExpOpaqueName TypeExp vn
te = TypeExp vn -> [Char]
forall a. Pretty a => a -> [Char]
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 :: [Char] -> Exp -> InternaliseM (BodyT SOACS)
internaliseBody [Char]
desc Exp
e =
  InternaliseM [SubExp] -> InternaliseM (Body (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
m [SubExp] -> m (Body (Lore m))
buildBody_ (InternaliseM [SubExp] -> InternaliseM (Body (Lore InternaliseM)))
-> InternaliseM [SubExp] -> InternaliseM (Body (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp ([Char]
desc [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"_res") Exp
e

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

internaliseAppExp :: String -> E.AppExp -> InternaliseM [I.SubExp]
internaliseAppExp :: [Char] -> AppExp -> InternaliseM [SubExp]
internaliseAppExp [Char]
desc (E.Index Exp
e [DimIndexBase Info VName]
idxs SrcLoc
loc) = do
  [VName]
vs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"indexed" Exp
e
  [SubExp]
dims <- case [VName]
vs of
    [] -> [SubExp] -> InternaliseM [SubExp]
forall (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 lore (m :: * -> *).
HasScope lore m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
v
  ([DimIndex SubExp]
idxs', Certificates
cs) <- SrcLoc
-> [SubExp]
-> [DimIndexBase Info VName]
-> InternaliseM ([DimIndex SubExp], Certificates)
internaliseSlice SrcLoc
loc [SubExp]
dims [DimIndexBase Info VName]
idxs
  let index :: VName -> InternaliseM (ExpT SOACS)
index VName
v = do
        TypeBase Shape NoUniqueness
v_t <- VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall lore (m :: * -> *).
HasScope lore 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 lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ VName -> [DimIndex SubExp] -> BasicOp
I.Index VName
v ([DimIndex SubExp] -> BasicOp) -> [DimIndex SubExp] -> BasicOp
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoUniqueness
-> [DimIndex SubExp] -> [DimIndex SubExp]
fullSlice TypeBase Shape NoUniqueness
v_t [DimIndex SubExp]
idxs'
  Certificates -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
cs (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char] -> [Exp (Lore InternaliseM)] -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> [Exp (Lore m)] -> m [SubExp]
letSubExps [Char]
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 [Char]
desc (E.Range Exp
start Maybe Exp
maybe_second Inclusiveness Exp
end SrcLoc
loc) = do
  SubExp
start' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"range_start" Exp
start
  SubExp
end' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"range_end" (Exp -> InternaliseM SubExp) -> Exp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ case Inclusiveness Exp
end of
    DownToExclusive Exp
e -> Exp
e
    ToInclusive Exp
e -> Exp
e
    UpToExclusive Exp
e -> Exp
e
  Maybe SubExp
maybe_second' <-
    (Exp -> InternaliseM SubExp)
-> Maybe Exp -> InternaliseM (Maybe SubExp)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ([Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"range_second") Maybe Exp
maybe_second

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

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

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

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

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

  SubExp
bounds_invalid_downwards <-
    [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"bounds_invalid_downwards" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp CmpOp
le_op SubExp
start' SubExp
end'
  SubExp
bounds_invalid_upwards <-
    [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"bounds_invalid_upwards" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp CmpOp
lt_op SubExp
end' SubExp
start'

  (SubExp
distance, SubExp
step_wrong_dir, SubExp
bounds_invalid) <- case Inclusiveness Exp
end of
    DownToExclusive {} -> do
      SubExp
step_wrong_dir <-
        [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"step_wrong_dir" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
one
      SubExp
distance <-
        [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"distance" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
it Overflow
I.OverflowWrap) SubExp
start' SubExp
end'
      SubExp
distance_i64 <- IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder 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 <-
        [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"step_wrong_dir" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
negone
      SubExp
distance <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"distance" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
it Overflow
I.OverflowWrap) SubExp
end' SubExp
start'
      SubExp
distance_i64 <- IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder 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 <-
        [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"downwards" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
negone
      SubExp
distance_downwards_exclusive <-
        [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"distance_downwards_exclusive" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
it Overflow
I.OverflowWrap) SubExp
start' SubExp
end'
      SubExp
distance_upwards_exclusive <-
        [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"distance_upwards_exclusive" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
it Overflow
I.OverflowWrap) SubExp
end' SubExp
start'

      SubExp
bounds_invalid <-
        [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"bounds_invalid" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          SubExp
-> BodyT SOACS
-> BodyT SOACS
-> IfDec (BranchType SOACS)
-> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If
            SubExp
downwards
            ([SubExp] -> BodyT SOACS
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [SubExp
bounds_invalid_downwards])
            ([SubExp] -> BodyT SOACS
forall lore. Bindable lore => [SubExp] -> Body lore
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 <-
        [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"distance_exclusive" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          SubExp
-> BodyT SOACS
-> BodyT SOACS
-> IfDec (BranchType SOACS)
-> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If
            SubExp
downwards
            ([SubExp] -> BodyT SOACS
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [SubExp
distance_downwards_exclusive])
            ([SubExp] -> BodyT SOACS
forall lore. Bindable lore => [SubExp] -> Body lore
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 :: * -> *).
MonadBinder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
distance_exclusive
      SubExp
distance <-
        [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"distance" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
            BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp
              (IntType -> Overflow -> BinOp
Add IntType
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 <-
    [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"step_invalid" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogOr SubExp
step_wrong_dir SubExp
step_zero

  SubExp
invalid <-
    [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"range_invalid" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogOr SubExp
step_invalid SubExp
bounds_invalid
  SubExp
valid <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"valid" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
I.Not SubExp
invalid
  Certificates
cs <- [Char]
-> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certificates
assert [Char]
"range_valid_c" SubExp
valid ErrorMsg SubExp
errmsg SrcLoc
loc

  SubExp
step_i64 <- IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
step
  SubExp
pos_step <-
    [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"pos_step" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Mul IntType
Int64 Overflow
I.OverflowWrap) SubExp
step_i64 SubExp
step_sign_i64

  SubExp
num_elems <-
    Certificates -> InternaliseM SubExp -> InternaliseM SubExp
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
cs (InternaliseM SubExp -> InternaliseM SubExp)
-> InternaliseM SubExp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"num_elems" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
        BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Safety -> BinOp
SDivUp IntType
Int64 Safety
I.Unsafe) SubExp
distance SubExp
pos_step

  SubExp
se <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
desc (BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> SubExp -> SubExp -> IntType -> BasicOp
I.Iota SubExp
num_elems SubExp
start' SubExp
step IntType
it)
  [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [SubExp
se]
internaliseAppExp [Char]
desc (E.Coerce Exp
e (TypeDecl TypeExp VName
dt (Info StructType
et)) SrcLoc
loc) = do
  [SubExp]
ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
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 SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt64 [SubExp]
dims)
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
") cannot match shape of type `"]
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
dt'
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"`."]
    ErrorMsg SubExp
-> SrcLoc -> ExtType -> [Char] -> SubExp -> InternaliseM SubExp
ensureExtShape ([ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [ErrorMsgPart SubExp]
parts) SrcLoc
loc (DeclExtType -> ExtType
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl DeclExtType
t') [Char]
desc SubExp
e'
internaliseAppExp [Char]
desc e :: AppExp
e@E.Apply {} = do
  (QualName VName
qfname, [(Exp, Maybe VName)]
args) <- AppExp -> InternaliseM (QualName VName, [(Exp, 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 = [Char] -> Name
nameFromString ([Char] -> Name) -> [Char] -> Name
forall a b. (a -> b) -> a -> b
$ Name -> [Char]
forall a. Pretty a => a -> [Char]
pretty (Name -> [Char]) -> Name -> [Char]
forall a b. (a -> b) -> a -> b
$ VName -> Name
baseName (VName -> Name) -> VName -> Name
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname
      loc :: SrcLoc
loc = AppExp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf AppExp
e
      arg_desc :: [Char]
arg_desc = Name -> [Char]
nameToString Name
fname [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"_arg"

  -- Some functions are magical (overloaded) and we handle that here.
  case () of
    -- Overloaded functions never take array arguments (except
    -- equality, but those cannot be existential), so we can safely
    -- ignore the existential dimensions.
    ()
      | Just [Char] -> InternaliseM [SubExp]
internalise <- QualName VName
-> [Exp] -> SrcLoc -> Maybe ([Char] -> InternaliseM [SubExp])
isOverloadedFunction QualName VName
qfname (((Exp, Maybe VName) -> Exp) -> [(Exp, Maybe VName)] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (Exp, Maybe VName) -> Exp
forall a b. (a, b) -> a
fst [(Exp, Maybe VName)]
args) SrcLoc
loc ->
        [Char] -> InternaliseM [SubExp]
internalise [Char]
desc
      | VName -> Int
baseTag (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag,
        Just (PrimType
rettype, [PrimType]
_) <- Name
-> Map Name (PrimType, [PrimType]) -> Maybe (PrimType, [PrimType])
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
fname Map Name (PrimType, [PrimType])
I.builtInFunctions -> do
        let tag :: [a] -> [(a, Diet)]
tag [a]
ses = [(a
se, Diet
I.Observe) | a
se <- [a]
ses]
        [[SubExp]]
args' <- [[SubExp]] -> [[SubExp]]
forall a. [a] -> [a]
reverse ([[SubExp]] -> [[SubExp]])
-> InternaliseM [[SubExp]] -> InternaliseM [[SubExp]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Exp, Maybe VName) -> InternaliseM [SubExp])
-> [(Exp, Maybe VName)] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Char] -> (Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg [Char]
arg_desc) ([(Exp, Maybe VName)] -> [(Exp, Maybe VName)]
forall a. [a] -> [a]
reverse [(Exp, Maybe VName)]
args)
        let args'' :: [(SubExp, Diet)]
args'' = ([SubExp] -> [(SubExp, Diet)]) -> [[SubExp]] -> [(SubExp, Diet)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap [SubExp] -> [(SubExp, Diet)]
forall {a}. [a] -> [(a, Diet)]
tag [[SubExp]]
args'
        [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
          Name
-> [(SubExp, Diet)]
-> [RetType SOACS]
-> (Safety, SrcLoc, [SrcLoc])
-> ExpT SOACS
forall lore.
Name
-> [(SubExp, Diet)]
-> [RetType lore]
-> (Safety, SrcLoc, [SrcLoc])
-> ExpT lore
I.Apply
            Name
fname
            [(SubExp, Diet)]
args''
            [PrimType -> DeclExtType
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
rettype]
            (Safety
Safe, SrcLoc
loc, [])
      | Bool
otherwise -> do
        [SubExp]
args' <- [[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[SubExp]] -> [SubExp])
-> ([[SubExp]] -> [[SubExp]]) -> [[SubExp]] -> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[SubExp]] -> [[SubExp]]
forall a. [a] -> [a]
reverse ([[SubExp]] -> [SubExp])
-> InternaliseM [[SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Exp, Maybe VName) -> InternaliseM [SubExp])
-> [(Exp, Maybe VName)] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Char] -> (Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg [Char]
arg_desc) ([(Exp, Maybe VName)] -> [(Exp, Maybe VName)]
forall a. [a] -> [a]
reverse [(Exp, Maybe VName)]
args)
        ([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
<$> [Char]
-> QualName VName
-> [SubExp]
-> SrcLoc
-> InternaliseM ([SubExp], [ExtType])
funcall [Char]
desc QualName VName
qfname [SubExp]
args' SrcLoc
loc
internaliseAppExp [Char]
desc (E.LetPat [SizeBinder VName]
sizes PatternBase Info VName
pat Exp
e Exp
body SrcLoc
_) =
  [Char]
-> [SizeBinder VName]
-> PatternBase Info VName
-> Exp
-> Exp
-> (Exp -> InternaliseM [SubExp])
-> InternaliseM [SubExp]
forall a.
[Char]
-> [SizeBinder VName]
-> PatternBase Info VName
-> Exp
-> Exp
-> (Exp -> InternaliseM a)
-> InternaliseM a
internalisePat [Char]
desc [SizeBinder VName]
sizes PatternBase Info VName
pat Exp
e Exp
body ([Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc)
internaliseAppExp [Char]
_ (E.LetFun VName
ofname ([TypeParamBase VName], [PatternBase Info VName],
 Maybe (TypeExp VName), Info StructType, Exp)
_ Exp
_ SrcLoc
_) =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"Unexpected LetFun " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ VName -> [Char]
forall a. Pretty a => a -> [Char]
pretty VName
ofname
internaliseAppExp [Char]
desc (E.DoLoop [VName]
sparams PatternBase Info VName
mergepat Exp
mergeexp LoopFormBase Info VName
form Exp
loopbody SrcLoc
loc) = do
  [SubExp]
ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"loop_init" Exp
mergeexp
  ((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 (Lore InternaliseM))
forall (m :: * -> *) a.
MonadBinder m =>
m a -> m (a, Stms (Lore m))
collectStms (InternaliseM
   (BodyT SOACS,
    (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
 -> InternaliseM
      ((BodyT SOACS,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])),
       Stms (Lore InternaliseM)))
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     ((BodyT SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])),
      Stms (Lore 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 (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *). MonadBinder m => Stms (Lore m) -> m ()
addStms Stms (Lore 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]
forall (m :: * -> *).
(HasScope SOACS m, Monad m) =>
[VName] -> [FParam] -> [TypeBase Shape NoUniqueness] -> m [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 lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param DeclType] -> Scope SOACS
forall lore dec.
(FParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
scopeOfFParams ([Param DeclType] -> Scope SOACS)
-> [Param DeclType] -> Scope SOACS
forall a b. (a -> b) -> a -> b
$ ((Param DeclType, SubExp) -> Param DeclType)
-> [(Param DeclType, SubExp)] -> [Param DeclType]
forall a b. (a -> b) -> [a] -> [b]
map (Param DeclType, SubExp) -> Param DeclType
forall a b. (a, b) -> a
fst [(Param DeclType, SubExp)]
merge) (InternaliseM (BodyT SOACS) -> InternaliseM (BodyT SOACS))
-> (InternaliseM [SubExp] -> InternaliseM (BodyT SOACS))
-> InternaliseM [SubExp]
-> InternaliseM (BodyT SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LoopForm SOACS
-> InternaliseM (BodyT SOACS) -> InternaliseM (BodyT SOACS)
forall lore a (m :: * -> *) b.
(Scoped lore a, LocalScope lore m) =>
a -> m b -> m b
inScopeOf LoopForm SOACS
form' (InternaliseM (BodyT SOACS) -> InternaliseM (BodyT SOACS))
-> (InternaliseM [SubExp] -> InternaliseM (BodyT SOACS))
-> InternaliseM [SubExp]
-> InternaliseM (BodyT SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InternaliseM [SubExp] -> InternaliseM (BodyT SOACS)
forall (m :: * -> *).
MonadBinder m =>
m [SubExp] -> m (Body (Lore m))
buildBody_ (InternaliseM [SubExp] -> InternaliseM (BodyT SOACS))
-> InternaliseM [SubExp] -> InternaliseM (BodyT SOACS)
forall a b. (a -> b) -> a -> b
$
      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])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Body (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *). MonadBinder m => Body (Lore m) -> m [SubExp]
bodyBind Body (Lore 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. MonadBinder m => Attrs -> m a -> m a
attributing
      Attrs
attrs
      ([Char] -> Exp (Lore InternaliseM) -> InternaliseM [VName]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [VName]
letTupExp [Char]
desc ([(FParam, SubExp)]
-> [(FParam, SubExp)]
-> LoopForm SOACS
-> BodyT SOACS
-> ExpT SOACS
forall lore.
[(FParam lore, SubExp)]
-> [(FParam lore, SubExp)]
-> LoopForm lore
-> BodyT lore
-> ExpT lore
I.DoLoop [(Param DeclType, SubExp)]
[(FParam, SubExp)]
ctxmerge [(Param DeclType, SubExp)]
[(FParam, SubExp)]
valmerge LoopForm SOACS
form' 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
  ([SubExp],
   (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a.
InternaliseM ([SubExp], a) -> InternaliseM (BodyT SOACS, a)
bodyFromStms (InternaliseM
   ([SubExp],
    (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
 -> InternaliseM
      (BodyT SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     ([SubExp],
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$
        LoopForm SOACS
-> InternaliseM
     ([SubExp],
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     ([SubExp],
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall lore a (m :: * -> *) b.
(Scoped lore a, LocalScope lore m) =>
a -> m b -> m b
inScopeOf LoopForm SOACS
form' (InternaliseM
   ([SubExp],
    (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
 -> InternaliseM
      ([SubExp],
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     ([SubExp],
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     ([SubExp],
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ do
          [SubExp]
ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"loopres" Exp
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]
forall (m :: * -> *).
(HasScope SOACS m, Monad m) =>
[VName] -> [FParam] -> [TypeBase Shape NoUniqueness] -> m [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
          ([SubExp],
 (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     ([SubExp],
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall (m :: * -> *) a. Monad m => a -> m a
return
            ( [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 PatternBase Info VName
x Exp
arr) = do
      [VName]
arr' <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"for_in_arr" Exp
arr
      [TypeBase Shape NoUniqueness]
arr_ts <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall lore (m :: * -> *).
HasScope lore m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arr'
      let w :: SubExp
w = Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [TypeBase Shape NoUniqueness]
arr_ts

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

      [TypeBase Shape NoUniqueness]
ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
mergeinit
      [TypeParamBase VName]
-> PatternBase 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]
-> PatternBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([FParam] -> [FParam] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' PatternBase 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' ->
          [PatternBase 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.
[PatternBase Info VName]
-> [TypeBase Shape NoUniqueness]
-> ([LParam] -> InternaliseM a)
-> InternaliseM a
bindingLambdaParams [PatternBase 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 lore.
VName
-> IntType -> SubExp -> [(LParam lore, VName)] -> LoopForm lore
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 Exp
num_iterations) = do
      SubExp
num_iterations' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"upper_bound" Exp
num_iterations
      TypeBase Shape NoUniqueness
num_iterations_t <- SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
I.subExpType SubExp
num_iterations'
      IntType
it <- case TypeBase Shape NoUniqueness
num_iterations_t of
        I.Prim (IntType IntType
it) -> IntType -> InternaliseM IntType
forall (m :: * -> *) a. Monad m => a -> m a
return IntType
it
        TypeBase Shape NoUniqueness
_ -> [Char] -> InternaliseM IntType
forall a. HasCallStack => [Char] -> a
error [Char]
"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]
-> PatternBase 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]
-> PatternBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([FParam] -> [FParam] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' PatternBase 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 lore.
VName
-> IntType -> SubExp -> [(LParam lore, VName)] -> LoopForm lore
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 Exp
cond) = do
      [TypeBase Shape NoUniqueness]
ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
mergeinit
      [TypeParamBase VName]
-> PatternBase 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]
-> PatternBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([FParam] -> [FParam] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' PatternBase 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]
forall (m :: * -> *).
(HasScope SOACS m, Monad m) =>
[VName] -> [FParam] -> [TypeBase Shape NoUniqueness] -> m [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_bnds) <- InternaliseM SubExp
-> InternaliseM (SubExp, Stms (Lore InternaliseM))
forall (m :: * -> *) a.
MonadBinder m =>
m a -> m (a, Stms (Lore m))
collectStms (InternaliseM SubExp
 -> InternaliseM (SubExp, Stms (Lore InternaliseM)))
-> InternaliseM SubExp
-> InternaliseM (SubExp, Stms (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ do
          [(Param DeclType, SubExp)]
-> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param DeclType] -> [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 (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
          [(Param DeclType, SubExp)]
-> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param DeclType] -> [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 (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
                BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
                  case SubExp
se of
                    I.Var VName
v
                      | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ 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
          [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"loop_cond" Exp
cond

        Stms (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *). MonadBinder m => Stms (Lore m) -> m ()
addStms Stms (Lore InternaliseM)
Stms SOACS
init_loop_cond_bnds

        InternaliseM
  ([SubExp],
   (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (BodyT SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a.
InternaliseM ([SubExp], a) -> InternaliseM (BodyT SOACS, a)
bodyFromStms (InternaliseM
   ([SubExp],
    (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
 -> InternaliseM
      (BodyT SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     ([SubExp],
      (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 <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"loopres" Exp
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 <- [Char] -> DeclType -> InternaliseM (Param DeclType)
forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"loop_while" (DeclType -> InternaliseM (Param DeclType))
-> DeclType -> InternaliseM (Param DeclType)
forall a b. (a -> b) -> a -> b
$ PrimType -> DeclType
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.Bool
          [SubExp]
shapeargs <- [VName]
-> [FParam]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM [SubExp]
forall (m :: * -> *).
(HasScope SOACS m, Monad m) =>
[VName] -> [FParam] -> [TypeBase Shape NoUniqueness] -> m [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 lore (m :: * -> *).
(Renameable lore, MonadFreshNames m) =>
Body lore -> m (Body lore)
renameBody (BodyT SOACS -> InternaliseM (BodyT SOACS))
-> (InternaliseM [SubExp] -> InternaliseM (BodyT SOACS))
-> InternaliseM [SubExp]
-> InternaliseM (BodyT SOACS)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< InternaliseM [SubExp] -> InternaliseM (BodyT SOACS)
forall (m :: * -> *).
MonadBinder m =>
m [SubExp] -> m (Body (Lore m))
buildBody_ (InternaliseM [SubExp] -> InternaliseM (BodyT SOACS))
-> InternaliseM [SubExp] -> 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 (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
            [(Param DeclType, SubExp)]
-> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param DeclType] -> [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 (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
                  BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
                    case SubExp
se of
                      I.Var VName
v
                        | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ 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
            [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"loop_cond" Exp
cond
          [SubExp]
loop_end_cond <- Body (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *). MonadBinder m => Body (Lore m) -> m [SubExp]
bodyBind Body (Lore InternaliseM)
BodyT SOACS
loop_end_cond_body

          ([SubExp],
 (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     ([SubExp],
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall (m :: * -> *) a. Monad m => a -> m a
return
            ( [SubExp]
shapeargs [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
loop_end_cond [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
ses,
              ( VName -> LoopForm SOACS
forall lore. VName -> LoopForm lore
I.WhileLoop (VName -> LoopForm SOACS) -> VName -> LoopForm SOACS
forall a b. (a -> b) -> a -> b
$ Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName Param DeclType
loop_while,
                [Param DeclType]
[FParam]
shapepat,
                Param DeclType
loop_while Param DeclType -> [Param DeclType] -> [Param DeclType]
forall a. a -> [a] -> [a]
: [Param DeclType]
[FParam]
mergepat',
                SubExp
loop_initial_cond SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: [SubExp]
mergeinit
              )
            )
internaliseAppExp [Char]
desc (E.LetWith IdentBase Info VName
name IdentBase Info VName
src [DimIndexBase Info VName]
idxs Exp
ve Exp
body SrcLoc
loc) = do
  let pat :: PatternBase Info VName
pat = VName -> Info PatternType -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase 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 PatternType
forall (f :: * -> *) vn. IdentBase f vn -> f PatternType
E.identType IdentBase Info VName
name) SrcLoc
loc
      src_t :: Info PatternType
src_t = PatternType -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
E.fromStruct (PatternType -> PatternType)
-> Info PatternType -> Info PatternType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IdentBase Info VName -> Info PatternType
forall (f :: * -> *) vn. IdentBase f vn -> f PatternType
E.identType IdentBase Info VName
src
      e :: Exp
e = Exp -> [DimIndexBase Info VName] -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> [DimIndexBase f vn] -> ExpBase f vn -> SrcLoc -> ExpBase f vn
E.Update (QualName VName -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
E.Var (VName -> QualName VName
forall v. v -> QualName v
E.qualName (VName -> QualName VName) -> VName -> QualName VName
forall a b. (a -> b) -> a -> b
$ IdentBase Info VName -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
E.identName IdentBase Info VName
src) Info PatternType
src_t SrcLoc
loc) [DimIndexBase Info VName]
idxs Exp
ve SrcLoc
loc
  [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc (Exp -> InternaliseM [SubExp]) -> Exp -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
E.AppExp
      ([SizeBinder VName]
-> PatternBase Info VName -> Exp -> Exp -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatternBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
E.LetPat [] PatternBase Info VName
pat Exp
e Exp
body SrcLoc
loc)
      (AppRes -> Info AppRes
forall a. a -> Info a
Info (PatternType -> [VName] -> AppRes
AppRes (Exp -> PatternType
E.typeOf Exp
body) [VName]
forall a. Monoid a => a
mempty))
internaliseAppExp [Char]
desc (E.Match Exp
e NonEmpty (CaseBase Info VName)
cs SrcLoc
_) = do
  [SubExp]
ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp ([Char]
desc [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"_scrutinee") Exp
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 PatternBase Info VName
pCase Exp
eCase SrcLoc
_, Maybe (NonEmpty (CaseBase Info VName))
Nothing) -> do
      (SubExp
_, [SubExp]
pertinent) <- PatternBase Info VName
-> [SubExp] -> InternaliseM (SubExp, [SubExp])
generateCond PatternBase Info VName
pCase [SubExp]
ses
      [SizeBinder VName]
-> PatternBase Info VName
-> [SubExp]
-> Exp
-> (Exp -> InternaliseM [SubExp])
-> InternaliseM [SubExp]
forall a.
[SizeBinder VName]
-> PatternBase Info VName
-> [SubExp]
-> Exp
-> (Exp -> InternaliseM a)
-> InternaliseM a
internalisePat' [] PatternBase Info VName
pCase [SubExp]
pertinent Exp
eCase ([Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc)
    (CaseBase Info VName
c, Just NonEmpty (CaseBase Info VName)
cs') -> do
      let CasePat PatternBase Info VName
pLast Exp
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) <- PatternBase Info VName
-> [SubExp] -> InternaliseM (SubExp, [SubExp])
generateCond PatternBase Info VName
pLast [SubExp]
ses
        BodyT SOACS
eLast' <- [SizeBinder VName]
-> PatternBase Info VName
-> [SubExp]
-> Exp
-> (Exp -> InternaliseM (BodyT SOACS))
-> InternaliseM (BodyT SOACS)
forall a.
[SizeBinder VName]
-> PatternBase Info VName
-> [SubExp]
-> Exp
-> (Exp -> InternaliseM a)
-> InternaliseM a
internalisePat' [] PatternBase Info VName
pLast [SubExp]
pertinent Exp
eLast ([Char] -> Exp -> InternaliseM (BodyT SOACS)
internaliseBody [Char]
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 (Lore InternaliseM))]
-> InternaliseM (Body (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
[m (Exp (Lore m))] -> m (Body (Lore m))
eBody ([InternaliseM (Exp (Lore InternaliseM))]
 -> InternaliseM (Body (Lore InternaliseM)))
-> [InternaliseM (Exp (Lore InternaliseM))]
-> InternaliseM (Body (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ InternaliseM (ExpT SOACS) -> [InternaliseM (ExpT SOACS)]
forall (m :: * -> *) a. Monad m => a -> m a
return (InternaliseM (ExpT SOACS) -> [InternaliseM (ExpT SOACS)])
-> InternaliseM (ExpT SOACS) -> [InternaliseM (ExpT SOACS)]
forall a b. (a -> b) -> a -> b
$ [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'
      [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
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 [Char]
desc (E.If Exp
ce Exp
te Exp
fe SrcLoc
_) =
  [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
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 (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
(MonadBinder m, BranchType (Lore m) ~ ExtType) =>
m (Exp (Lore m))
-> m (Body (Lore m)) -> m (Body (Lore m)) -> m (Exp (Lore m))
eIf
      (BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS)
-> (SubExp -> BasicOp) -> SubExp -> ExpT SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
SubExp (SubExp -> ExpT SOACS)
-> InternaliseM SubExp -> InternaliseM (ExpT SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"cond" Exp
ce)
      ([Char] -> Exp -> InternaliseM (BodyT SOACS)
internaliseBody ([Char]
desc [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"_t") Exp
te)
      ([Char] -> Exp -> InternaliseM (BodyT SOACS)
internaliseBody ([Char]
desc [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"_f") Exp
fe)
internaliseAppExp [Char]
_ e :: AppExp
e@E.BinOp {} =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseAppExp: Unexpected BinOp " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ AppExp -> [Char]
forall a. Pretty a => a -> [Char]
pretty AppExp
e

internaliseExp :: String -> E.Exp -> InternaliseM [I.SubExp]
internaliseExp :: [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc (E.Parens Exp
e SrcLoc
_) =
  [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
internaliseExp [Char]
desc (E.QualParens (QualName VName, SrcLoc)
_ Exp
e SrcLoc
_) =
  [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
internaliseExp [Char]
desc (E.StringLit [Word8]
vs SrcLoc
_) =
  (SubExp -> [SubExp])
-> InternaliseM SubExp -> InternaliseM [SubExp]
forall (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
$
    [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ [SubExp] -> TypeBase Shape NoUniqueness -> BasicOp
I.ArrayLit ((Word8 -> SubExp) -> [Word8] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map Word8 -> SubExp
forall v. IsValue v => v -> SubExp
constant [Word8]
vs) (TypeBase Shape NoUniqueness -> BasicOp)
-> TypeBase Shape NoUniqueness -> BasicOp
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int8
internaliseExp [Char]
_ (E.Var (E.QualName [VName]
_ VName
name) Info PatternType
_ 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 [Char]
desc (E.AppExp AppExp
e (Info AppRes
appres)) = do
  [SubExp]
ses <- [Char] -> AppExp -> InternaliseM [SubExp]
internaliseAppExp [Char]
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 [Char]
_ (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 [Char]
_ (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 [Char]
desc (E.TupLit [Exp]
es SrcLoc
_) = [[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[SubExp]] -> [SubExp])
-> InternaliseM [[SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> InternaliseM [SubExp]) -> [Exp] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc) [Exp]
es
internaliseExp [Char]
desc (E.RecordLit [FieldBase Info VName]
orig_fields SrcLoc
_) =
  ((Name, [SubExp]) -> [SubExp]) -> [(Name, [SubExp])] -> [SubExp]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Name, [SubExp]) -> [SubExp]
forall a b. (a, b) -> b
snd ([(Name, [SubExp])] -> [SubExp])
-> ([Map Name [SubExp]] -> [(Name, [SubExp])])
-> [Map Name [SubExp]]
-> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name [SubExp] -> [(Name, [SubExp])]
forall a. Map Name a -> [(Name, a)]
sortFields (Map Name [SubExp] -> [(Name, [SubExp])])
-> ([Map Name [SubExp]] -> Map Name [SubExp])
-> [Map Name [SubExp]]
-> [(Name, [SubExp])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Map Name [SubExp]] -> Map Name [SubExp]
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
M.unions ([Map Name [SubExp]] -> [SubExp])
-> InternaliseM [Map Name [SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (FieldBase Info VName -> InternaliseM (Map Name [SubExp]))
-> [FieldBase Info VName] -> InternaliseM [Map Name [SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM FieldBase Info VName -> InternaliseM (Map Name [SubExp])
internaliseField [FieldBase Info VName]
orig_fields
  where
    internaliseField :: FieldBase Info VName -> InternaliseM (Map Name [SubExp])
internaliseField (E.RecordFieldExplicit Name
name Exp
e SrcLoc
_) =
      Name -> [SubExp] -> Map Name [SubExp]
forall k a. k -> a -> Map k a
M.singleton Name
name ([SubExp] -> Map Name [SubExp])
-> InternaliseM [SubExp] -> InternaliseM (Map Name [SubExp])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
    internaliseField (E.RecordFieldImplicit VName
name Info PatternType
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 -> Exp -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
E.RecordFieldExplicit
          (VName -> Name
baseName VName
name)
          (QualName VName -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
E.Var (VName -> QualName VName
forall v. v -> QualName v
E.qualName VName
name) Info PatternType
t SrcLoc
loc)
          SrcLoc
loc
internaliseExp [Char]
desc (E.ArrayLit [Exp]
es (Info PatternType
arr_t) SrcLoc
loc)
  -- If this is a multidimensional array literal of primitives, we
  -- treat it specially by flattening it out followed by a reshape.
  -- This cuts down on the amount of statements that are produced, and
  -- thus allows us to efficiently handle huge array literals - a
  -- corner case, but an important one.
  | Just (([Int]
eshape, [Exp]
e') : [([Int], [Exp])]
es') <- (Exp -> Maybe ([Int], [Exp])) -> [Exp] -> Maybe [([Int], [Exp])]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> Maybe ([Int], [Exp])
isArrayLiteral [Exp]
es,
    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [Int] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
eshape,
    (([Int], [Exp]) -> Bool) -> [([Int], [Exp])] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (([Int]
eshape [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
==) ([Int] -> Bool)
-> (([Int], [Exp]) -> [Int]) -> ([Int], [Exp]) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Int], [Exp]) -> [Int]
forall a b. (a, b) -> a
fst) [([Int], [Exp])]
es',
    Just PatternType
basetype <- Int -> PatternType -> Maybe PatternType
forall dim as. Int -> TypeBase dim as -> Maybe (TypeBase dim as)
E.peelArray ([Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
eshape) PatternType
arr_t = do
    let flat_lit :: Exp
flat_lit = [Exp] -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
[ExpBase f vn] -> f PatternType -> SrcLoc -> ExpBase f vn
E.ArrayLit ([Exp]
e' [Exp] -> [Exp] -> [Exp]
forall a. [a] -> [a] -> [a]
++ (([Int], [Exp]) -> [Exp]) -> [([Int], [Exp])] -> [Exp]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Int], [Exp]) -> [Exp]
forall a b. (a, b) -> b
snd [([Int], [Exp])]
es') (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
basetype) SrcLoc
loc
        new_shape :: [Int]
new_shape = [Exp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Exp]
es Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
eshape
    [VName]
flat_arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"flat_literal" Exp
flat_lit
    [VName] -> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
flat_arrs ((VName -> InternaliseM SubExp) -> InternaliseM [SubExp])
-> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
flat_arr -> do
      TypeBase Shape NoUniqueness
flat_arr_t <- VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall lore (m :: * -> *).
HasScope lore 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
      [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
I.Reshape ShapeChange SubExp
new_shape' VName
flat_arr
  | Bool
otherwise = do
    [[SubExp]]
es' <- (Exp -> InternaliseM [SubExp]) -> [Exp] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"arr_elem") [Exp]
es
    [DeclExtType]
arr_t_ext <- StructType -> InternaliseM [DeclExtType]
internaliseType (StructType -> InternaliseM [DeclExtType])
-> StructType -> InternaliseM [DeclExtType]
forall a b. (a -> b) -> a -> b
$ PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
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 (ShapeBase ExtSize) 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
            [] -> [Char] -> InternaliseM [TypeBase Shape NoUniqueness]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [TypeBase Shape NoUniqueness])
-> [Char] -> InternaliseM [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp ArrayLit: existential type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ PatternType -> [Char]
forall a. Pretty a => a -> [Char]
pretty PatternType
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
-> [Char]
-> SubExp
-> InternaliseM SubExp
ensureShape
                  ErrorMsg SubExp
"shape of element differs from shape of first element"
                  SrcLoc
loc
                  TypeBase Shape NoUniqueness
rt
                  [Char]
"elem_reshaped"
              )
              [SubExp]
ks
          ExpT SOACS -> InternaliseM (ExpT SOACS)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpT SOACS -> InternaliseM (ExpT SOACS))
-> ExpT SOACS -> InternaliseM (ExpT SOACS)
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ [SubExp] -> TypeBase Shape NoUniqueness -> BasicOp
I.ArrayLit [SubExp]
ks' TypeBase Shape NoUniqueness
rt

    [Char] -> [Exp (Lore InternaliseM)] -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> [Exp (Lore m)] -> m [SubExp]
letSubExps [Char]
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 :: Exp -> Maybe ([Int], [Exp])
isArrayLiteral (E.ArrayLit [Exp]
inner_es Info PatternType
_ SrcLoc
_) = do
      ([Int]
eshape, [Exp]
e) : [([Int], [Exp])]
inner_es' <- (Exp -> Maybe ([Int], [Exp])) -> [Exp] -> Maybe [([Int], [Exp])]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> Maybe ([Int], [Exp])
isArrayLiteral [Exp]
inner_es
      Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ (([Int], [Exp]) -> Bool) -> [([Int], [Exp])] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (([Int]
eshape [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
==) ([Int] -> Bool)
-> (([Int], [Exp]) -> [Int]) -> ([Int], [Exp]) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Int], [Exp]) -> [Int]
forall a b. (a, b) -> a
fst) [([Int], [Exp])]
inner_es'
      ([Int], [Exp]) -> Maybe ([Int], [Exp])
forall (m :: * -> *) a. Monad m => a -> m a
return ([Exp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Exp]
inner_es Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
eshape, [Exp]
e [Exp] -> [Exp] -> [Exp]
forall a. [a] -> [a] -> [a]
++ (([Int], [Exp]) -> [Exp]) -> [([Int], [Exp])] -> [Exp]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Int], [Exp]) -> [Exp]
forall a b. (a, b) -> b
snd [([Int], [Exp])]
inner_es')
    isArrayLiteral Exp
e =
      ([Int], [Exp]) -> Maybe ([Int], [Exp])
forall a. a -> Maybe a
Just ([], [Exp
e])
internaliseExp [Char]
desc (E.Ascript Exp
e TypeDeclBase Info VName
_ SrcLoc
_) =
  [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
internaliseExp [Char]
desc (E.Negate Exp
e SrcLoc
_) = do
  SubExp
e' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"negate_arg" Exp
e
  TypeBase Shape NoUniqueness
et <- SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType SubExp
e'
  case TypeBase Shape NoUniqueness
et of
    I.Prim (I.IntType IntType
t) ->
      [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (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) ->
      [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (FloatType -> BinOp
I.FSub FloatType
t) (FloatType -> Double -> SubExp
I.floatConst FloatType
t Double
0) SubExp
e'
    TypeBase Shape NoUniqueness
_ -> [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise.internaliseExp: non-numeric type in Negate"
internaliseExp [Char]
desc (E.Update Exp
src [DimIndexBase Info VName]
slice Exp
ve SrcLoc
loc) = do
  [SubExp]
ves <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"lw_val" Exp
ve
  [VName]
srcs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"src" Exp
src
  [SubExp]
dims <- case [VName]
srcs of
    [] -> [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 lore (m :: * -> *).
HasScope lore m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
v
  ([DimIndex SubExp]
idxs', Certificates
cs) <- SrcLoc
-> [SubExp]
-> [DimIndexBase Info VName]
-> InternaliseM ([DimIndex SubExp], Certificates)
internaliseSlice SrcLoc
loc [SubExp]
dims [DimIndexBase 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 lore (m :: * -> *).
HasScope lore m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
sname
        let full_slice :: [DimIndex SubExp]
full_slice = TypeBase Shape NoUniqueness
-> [DimIndex SubExp] -> [DimIndex 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` [DimIndex SubExp] -> [SubExp]
forall d. Slice d -> [d]
sliceDims [DimIndex SubExp]
full_slice
        SubExp
ve'' <-
          ErrorMsg SubExp
-> SrcLoc
-> TypeBase Shape NoUniqueness
-> [Char]
-> SubExp
-> InternaliseM SubExp
ensureShape
            ErrorMsg SubExp
"shape of value does not match shape of source array"
            SrcLoc
loc
            TypeBase Shape NoUniqueness
rowtype
            [Char]
"lw_val_correct_shape"
            SubExp
ve'
        [Char]
-> VName
-> [DimIndex SubExp]
-> Exp (Lore InternaliseM)
-> InternaliseM VName
forall (m :: * -> *).
MonadBinder m =>
[Char] -> VName -> [DimIndex SubExp] -> Exp (Lore m) -> m VName
letInPlace [Char]
desc VName
sname [DimIndex SubExp]
full_slice (Exp (Lore InternaliseM) -> InternaliseM VName)
-> Exp (Lore InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
ve''
  Certificates -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
cs (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var ([VName] -> [SubExp])
-> InternaliseM [VName] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> SubExp -> InternaliseM VName)
-> [VName] -> [SubExp] -> InternaliseM [VName]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM VName -> SubExp -> InternaliseM VName
comb [VName]
srcs [SubExp]
ves
internaliseExp [Char]
desc (E.RecordUpdate Exp
src [Name]
fields Exp
ve Info PatternType
_ SrcLoc
_) = do
  [SubExp]
src' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
src
  [SubExp]
ve' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
ve
  StructType
-> [Name] -> [SubExp] -> [SubExp] -> InternaliseM [SubExp]
forall {a}. StructType -> [Name] -> [a] -> [a] -> InternaliseM [a]
replace (Exp -> PatternType
E.typeOf Exp
src PatternType -> () -> StructType
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` ()) [Name]
fields [SubExp]
ve' [SubExp]
src'
  where
    replace :: StructType -> [Name] -> [a] -> [a] -> InternaliseM [a]
replace (E.Scalar (E.Record Map Name StructType
m)) (Name
f : [Name]
fs) [a]
ve' [a]
src'
      | Just StructType
t <- Name -> Map Name StructType -> Maybe StructType
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f Map Name StructType
m = do
        Int
i <-
          ([Int] -> Int) -> InternaliseM [Int] -> InternaliseM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (InternaliseM [Int] -> InternaliseM Int)
-> InternaliseM [Int] -> InternaliseM Int
forall a b. (a -> b) -> a -> b
$
            ((Name, StructType) -> InternaliseM Int)
-> [(Name, StructType)] -> InternaliseM [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (StructType -> InternaliseM Int
internalisedTypeSize (StructType -> InternaliseM Int)
-> ((Name, StructType) -> StructType)
-> (Name, StructType)
-> InternaliseM Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, StructType) -> StructType
forall a b. (a, b) -> b
snd) ([(Name, StructType)] -> InternaliseM [Int])
-> [(Name, StructType)] -> InternaliseM [Int]
forall a b. (a -> b) -> a -> b
$
              ((Name, StructType) -> Bool)
-> [(Name, StructType)] -> [(Name, StructType)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
/= Name
f) (Name -> Bool)
-> ((Name, StructType) -> Name) -> (Name, StructType) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, StructType) -> Name
forall a b. (a, b) -> a
fst) ([(Name, StructType)] -> [(Name, StructType)])
-> [(Name, StructType)] -> [(Name, StructType)]
forall a b. (a -> b) -> a -> b
$ Map Name StructType -> [(Name, StructType)]
forall a. Map Name a -> [(Name, a)]
sortFields Map Name StructType
m
        Int
k <- StructType -> InternaliseM Int
internalisedTypeSize StructType
t
        let ([a]
bef, [a]
to_update, [a]
aft) = Int -> Int -> [a] -> ([a], [a], [a])
forall a. Int -> Int -> [a] -> ([a], [a], [a])
splitAt3 Int
i Int
k [a]
src'
        [a]
src'' <- StructType -> [Name] -> [a] -> [a] -> InternaliseM [a]
replace StructType
t [Name]
fs [a]
ve' [a]
to_update
        [a] -> InternaliseM [a]
forall (m :: * -> *) a. Monad m => a -> m a
return ([a] -> InternaliseM [a]) -> [a] -> InternaliseM [a]
forall a b. (a -> b) -> a -> b
$ [a]
bef [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
src'' [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
aft
    replace StructType
_ [Name]
_ [a]
ve' [a]
_ = [a] -> InternaliseM [a]
forall (m :: * -> *) a. Monad m => a -> m a
return [a]
ve'
internaliseExp [Char]
desc (E.Attr AttrInfo
attr Exp
e SrcLoc
_) =
  (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
$ [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
  where
    attrs :: Attrs
attrs = Attr -> Attrs
oneAttr (Attr -> Attrs) -> Attr -> Attrs
forall a b. (a -> b) -> a -> b
$ AttrInfo -> Attr
internaliseAttr AttrInfo
attr
    f :: InternaliseEnv -> InternaliseEnv
f InternaliseEnv
env
      | Attr
"unsafe" Attr -> Attrs -> Bool
`inAttrs` Attrs
attrs,
        Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ InternaliseEnv -> Bool
envSafe InternaliseEnv
env =
        InternaliseEnv
env {envDoBoundsChecks :: Bool
envDoBoundsChecks = Bool
False}
      | Bool
otherwise =
        InternaliseEnv
env {envAttrs :: Attrs
envAttrs = InternaliseEnv -> Attrs
envAttrs InternaliseEnv
env Attrs -> Attrs -> Attrs
forall a. Semigroup a => a -> a -> a
<> Attrs
attrs}
internaliseExp [Char]
desc (E.Assert Exp
e1 Exp
e2 (Info [Char]
check) SrcLoc
loc) = do
  SubExp
e1' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"assert_cond" Exp
e1
  Certificates
c <- [Char]
-> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certificates
assert [Char]
"assert_c" SubExp
e1' ([ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [[Char] -> ErrorMsgPart SubExp
forall a. [Char] -> ErrorMsgPart a
ErrorString ([Char] -> ErrorMsgPart SubExp) -> [Char] -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ [Char]
"Assertion is false: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
check]) SrcLoc
loc
  -- Make sure there are some bindings to certify.
  Certificates -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
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 :: * -> *}. MonadBinder m => SubExp -> m SubExp
rebind ([SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e2
  where
    rebind :: SubExp -> m SubExp
rebind SubExp
v = do
      VName
v' <- [Char] -> m VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"assert_res"
      [VName] -> Exp (Lore m) -> m ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [VName
v'] (Exp (Lore m) -> m ()) -> Exp (Lore m) -> m ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
v
      SubExp -> m SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> m SubExp) -> SubExp -> m SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
v'
internaliseExp [Char]
_ (E.Constr Name
c [Exp]
es (Info (E.Scalar (E.Sum Map Name [PatternType]
fs))) SrcLoc
_) = do
  ([DeclExtType]
ts, Map Name (Int, [Int])
constr_map) <- Map Name [StructType]
-> InternaliseM ([DeclExtType], Map Name (Int, [Int]))
internaliseSumType (Map Name [StructType]
 -> InternaliseM ([DeclExtType], Map Name (Int, [Int])))
-> Map Name [StructType]
-> InternaliseM ([DeclExtType], Map Name (Int, [Int]))
forall a b. (a -> b) -> a -> b
$ ([PatternType] -> [StructType])
-> Map Name [PatternType] -> Map Name [StructType]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((PatternType -> StructType) -> [PatternType] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct) Map Name [PatternType]
fs
  [SubExp]
es' <- [[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[SubExp]] -> [SubExp])
-> InternaliseM [[SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> InternaliseM [SubExp]) -> [Exp] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"payload") [Exp]
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 (ShapeBase ExtSize) 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, MonadBinder 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 ->
      [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error [Char]
"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 <- [Char] -> Exp (Lore f) -> f SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"zero" (Exp (Lore f) -> f SubExp) -> f (Exp (Lore f)) -> f SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TypeBase Shape NoUniqueness -> f (Exp (Lore f))
forall (m :: * -> *).
MonadBinder m =>
TypeBase Shape NoUniqueness -> m (Exp (Lore 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 [Char]
_ (E.Constr Name
_ [Exp]
_ (Info PatternType
t) SrcLoc
loc) =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: constructor with type " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ PatternType -> [Char]
forall a. Pretty a => a -> [Char]
pretty PatternType
t [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr SrcLoc
loc
-- The "interesting" cases are over, now it's mostly boilerplate.

internaliseExp [Char]
_ (E.Literal PrimValue
v SrcLoc
_) =
  [SubExp] -> InternaliseM [SubExp]
forall (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 [Char]
_ (E.IntLit Integer
v (Info PatternType
t) SrcLoc
_) =
  case PatternType
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]
    PatternType
_ -> [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: nonsensical type for integer literal: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ PatternType -> [Char]
forall a. Pretty a => a -> [Char]
pretty PatternType
t
internaliseExp [Char]
_ (E.FloatLit Double
v (Info PatternType
t) SrcLoc
_) =
  case PatternType
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]
    PatternType
_ -> [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: nonsensical type for float literal: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ PatternType -> [Char]
forall a. Pretty a => a -> [Char]
pretty PatternType
t
-- Builtin operators are handled specially because they are
-- overloaded.
internaliseExp [Char]
desc (E.Project Name
k Exp
e (Info PatternType
rt) SrcLoc
_) = do
  Int
n <- StructType -> InternaliseM Int
internalisedTypeSize (StructType -> InternaliseM Int) -> StructType -> InternaliseM Int
forall a b. (a -> b) -> a -> b
$ PatternType
rt PatternType -> () -> StructType
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` ()
  Int
i' <- ([Int] -> Int) -> InternaliseM [Int] -> InternaliseM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (InternaliseM [Int] -> InternaliseM Int)
-> InternaliseM [Int] -> InternaliseM Int
forall a b. (a -> b) -> a -> b
$
    (StructType -> InternaliseM Int)
-> [StructType] -> InternaliseM [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM StructType -> InternaliseM Int
internalisedTypeSize ([StructType] -> InternaliseM [Int])
-> [StructType] -> InternaliseM [Int]
forall a b. (a -> b) -> a -> b
$
      case Exp -> PatternType
E.typeOf Exp
e PatternType -> () -> StructType
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` () of
        E.Scalar (Record Map Name StructType
fs) ->
          ((Name, StructType) -> StructType)
-> [(Name, StructType)] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map (Name, StructType) -> StructType
forall a b. (a, b) -> b
snd ([(Name, StructType)] -> [StructType])
-> [(Name, StructType)] -> [StructType]
forall a b. (a -> b) -> a -> b
$ ((Name, StructType) -> Bool)
-> [(Name, StructType)] -> [(Name, StructType)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
/= Name
k) (Name -> Bool)
-> ((Name, StructType) -> Name) -> (Name, StructType) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, StructType) -> Name
forall a b. (a, b) -> a
fst) ([(Name, StructType)] -> [(Name, StructType)])
-> [(Name, StructType)] -> [(Name, StructType)]
forall a b. (a -> b) -> a -> b
$ Map Name StructType -> [(Name, StructType)]
forall a. Map Name a -> [(Name, a)]
sortFields Map Name StructType
fs
        StructType
t -> [StructType
t]
  Int -> [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
<$> [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
internaliseExp [Char]
_ e :: Exp
e@E.Lambda {} =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected lambda at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp [Char]
_ e :: Exp
e@E.OpSection {} =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected operator section at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp [Char]
_ e :: Exp
e@E.OpSectionLeft {} =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected left operator section at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp [Char]
_ e :: Exp
e@E.OpSectionRight {} =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected right operator section at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp [Char]
_ e :: Exp
e@E.ProjectSection {} =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected projection section at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp [Char]
_ e :: Exp
e@E.IndexSection {} =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected index section at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e)

internaliseArg :: String -> (E.Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg :: [Char] -> (Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg [Char]
desc (Exp
arg, Maybe VName
argdim) = do
  [SubExp]
arg' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
arg
  case ([SubExp]
arg', Maybe VName
argdim) of
    ([SubExp
se], Just VName
d) -> [VName] -> Exp (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [VName
d] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
    ([SubExp], Maybe VName)
_ -> () -> InternaliseM ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
  [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return [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.Pattern -> [I.SubExp] -> InternaliseM (I.SubExp, [I.SubExp])
generateCond :: PatternBase Info VName
-> [SubExp] -> InternaliseM (SubExp, [SubExp])
generateCond PatternBase Info VName
orig_p [SubExp]
orig_ses = do
  ([SubExp]
cmps, [SubExp]
pertinent, [SubExp]
_) <- PatternBase Info VName
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
forall {vn}.
(Eq vn, IsName vn) =>
PatternBase Info vn
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
compares PatternBase Info VName
orig_p [SubExp]
orig_ses
  SubExp
cmp <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"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 (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => [SubExp] -> m (Exp (Lore 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 :: PatternBase Info vn
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
compares (E.PatternLit PatLit
l Info PatternType
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 -> [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"constant" (Exp -> InternaliseM SubExp) -> Exp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ Integer -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Integer -> f PatternType -> SrcLoc -> ExpBase f vn
E.IntLit Integer
x Info PatternType
t SrcLoc
forall a. Monoid a => a
mempty
        PatLitFloat Double
x -> [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"constant" (Exp -> InternaliseM SubExp) -> Exp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ Double -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Double -> f PatternType -> SrcLoc -> ExpBase f vn
E.FloatLit Double
x Info PatternType
t SrcLoc
forall a. Monoid a => a
mempty
      PrimType
t' <- SubExp -> InternaliseM PrimType
subExpPrimType SubExp
se
      SubExp
cmp <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"match_lit" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
t') SubExp
e' SubExp
se
      ([SubExp], [SubExp], [SubExp])
-> InternaliseM ([SubExp], [SubExp], [SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return ([SubExp
cmp], [SubExp
se], [SubExp]
ses)
    compares (E.PatternConstr Name
c (Info (E.Scalar (E.Sum Map Name [PatternType]
fs))) [PatternBase 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
$ ([PatternType] -> [StructType])
-> Map Name [PatternType] -> Map Name [StructType]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((PatternType -> StructType) -> [PatternType] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct) Map Name [PatternType]
fs
      case Name -> Map Name (Int, [Int]) -> Maybe (Int, [Int])
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
c Map Name (Int, [Int])
m of
        Just (Int
i, [Int]
payload_is) -> do
          let i' :: SubExp
i' = IntType -> Integer -> SubExp
intConst IntType
Int8 (Integer -> SubExp) -> Integer -> SubExp
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
i
          let ([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 <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"match_constr" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
int8) SubExp
i' SubExp
se
          ([SubExp]
cmps, [SubExp]
pertinent, [SubExp]
_) <- [PatternBase Info vn]
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
comparesMany [PatternBase 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 ->
          [Char] -> InternaliseM ([SubExp], [SubExp], [SubExp])
forall a. HasCallStack => [Char] -> a
error [Char]
"generateCond: missing constructor"
    compares (E.PatternConstr Name
_ (Info PatternType
t) [PatternBase Info vn]
_ SrcLoc
_) [SubExp]
_ =
      [Char] -> InternaliseM ([SubExp], [SubExp], [SubExp])
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM ([SubExp], [SubExp], [SubExp]))
-> [Char] -> InternaliseM ([SubExp], [SubExp], [SubExp])
forall a b. (a -> b) -> a -> b
$ [Char]
"generateCond: PatternConstr has nonsensical type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ PatternType -> [Char]
forall a. Pretty a => a -> [Char]
pretty PatternType
t
    compares (E.Id vn
_ Info PatternType
t SrcLoc
loc) [SubExp]
ses =
      PatternBase Info vn
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
compares (Info PatternType -> SrcLoc -> PatternBase Info vn
forall (f :: * -> *) vn.
f PatternType -> SrcLoc -> PatternBase f vn
E.Wildcard Info PatternType
t SrcLoc
loc) [SubExp]
ses
    compares (E.Wildcard (Info PatternType
t) SrcLoc
_) [SubExp]
ses = do
      Int
n <- StructType -> InternaliseM Int
internalisedTypeSize (StructType -> InternaliseM Int) -> StructType -> InternaliseM Int
forall a b. (a -> b) -> a -> b
$ PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
t
      let ([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.PatternParens PatternBase Info vn
pat SrcLoc
_) [SubExp]
ses =
      PatternBase Info vn
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
compares PatternBase Info vn
pat [SubExp]
ses
    -- XXX: treat empty tuples and records as bool.
    compares (E.TuplePattern [] SrcLoc
loc) [SubExp]
ses =
      PatternBase Info vn
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
compares (Info PatternType -> SrcLoc -> PatternBase Info vn
forall (f :: * -> *) vn.
f PatternType -> SrcLoc -> PatternBase f vn
E.Wildcard (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) Aliasing -> PatternType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar (ScalarTypeBase (DimDecl VName) Aliasing -> PatternType)
-> ScalarTypeBase (DimDecl VName) Aliasing -> PatternType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) Aliasing
forall dim as. PrimType -> ScalarTypeBase dim as
E.Prim PrimType
E.Bool) SrcLoc
loc) [SubExp]
ses
    compares (E.RecordPattern [] SrcLoc
loc) [SubExp]
ses =
      PatternBase Info vn
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
compares (Info PatternType -> SrcLoc -> PatternBase Info vn
forall (f :: * -> *) vn.
f PatternType -> SrcLoc -> PatternBase f vn
E.Wildcard (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) Aliasing -> PatternType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar (ScalarTypeBase (DimDecl VName) Aliasing -> PatternType)
-> ScalarTypeBase (DimDecl VName) Aliasing -> PatternType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) Aliasing
forall dim as. PrimType -> ScalarTypeBase dim as
E.Prim PrimType
E.Bool) SrcLoc
loc) [SubExp]
ses
    compares (E.TuplePattern [PatternBase Info vn]
pats SrcLoc
_) [SubExp]
ses =
      [PatternBase Info vn]
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
comparesMany [PatternBase Info vn]
pats [SubExp]
ses
    compares (E.RecordPattern [(Name, PatternBase Info vn)]
fs SrcLoc
_) [SubExp]
ses =
      [PatternBase Info vn]
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
comparesMany (((Name, PatternBase Info vn) -> PatternBase Info vn)
-> [(Name, PatternBase Info vn)] -> [PatternBase Info vn]
forall a b. (a -> b) -> [a] -> [b]
map (Name, PatternBase Info vn) -> PatternBase Info vn
forall a b. (a, b) -> b
snd ([(Name, PatternBase Info vn)] -> [PatternBase Info vn])
-> [(Name, PatternBase Info vn)] -> [PatternBase Info vn]
forall a b. (a -> b) -> a -> b
$ Map Name (PatternBase Info vn) -> [(Name, PatternBase Info vn)]
forall a. Map Name a -> [(Name, a)]
E.sortFields (Map Name (PatternBase Info vn) -> [(Name, PatternBase Info vn)])
-> Map Name (PatternBase Info vn) -> [(Name, PatternBase Info vn)]
forall a b. (a -> b) -> a -> b
$ [(Name, PatternBase Info vn)] -> Map Name (PatternBase Info vn)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, PatternBase Info vn)]
fs) [SubExp]
ses
    compares (E.PatternAscription PatternBase Info vn
pat TypeDeclBase Info vn
_ SrcLoc
_) [SubExp]
ses =
      PatternBase Info vn
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
compares PatternBase Info vn
pat [SubExp]
ses
    compares PatternBase Info vn
pat [] =
      [Char] -> InternaliseM ([SubExp], [SubExp], [SubExp])
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM ([SubExp], [SubExp], [SubExp]))
-> [Char] -> InternaliseM ([SubExp], [SubExp], [SubExp])
forall a b. (a -> b) -> a -> b
$ [Char]
"generateCond: No values left for pattern " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ PatternBase Info vn -> [Char]
forall a. Pretty a => a -> [Char]
pretty PatternBase Info vn
pat

    comparesMany :: [PatternBase 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 (PatternBase Info vn
pat : [PatternBase Info vn]
pats) [SubExp]
ses = do
      ([SubExp]
cmps1, [SubExp]
pertinent1, [SubExp]
ses') <- PatternBase Info vn
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
compares PatternBase Info vn
pat [SubExp]
ses
      ([SubExp]
cmps2, [SubExp]
pertinent2, [SubExp]
ses'') <- [PatternBase Info vn]
-> [SubExp] -> InternaliseM ([SubExp], [SubExp], [SubExp])
comparesMany [PatternBase 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 PatternBase Info VName
p Exp
eCase SrcLoc
_) BodyT SOACS
bFail = do
  (SubExp
cond, [SubExp]
pertinent) <- PatternBase Info VName
-> [SubExp] -> InternaliseM (SubExp, [SubExp])
generateCond PatternBase Info VName
p [SubExp]
ses
  BodyT SOACS
eCase' <- [SizeBinder VName]
-> PatternBase Info VName
-> [SubExp]
-> Exp
-> (Exp -> InternaliseM (BodyT SOACS))
-> InternaliseM (BodyT SOACS)
forall a.
[SizeBinder VName]
-> PatternBase Info VName
-> [SubExp]
-> Exp
-> (Exp -> InternaliseM a)
-> InternaliseM a
internalisePat' [] PatternBase Info VName
p [SubExp]
pertinent Exp
eCase ([Char] -> Exp -> InternaliseM (BodyT SOACS)
internaliseBody [Char]
"case")
  InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
(MonadBinder m, BranchType (Lore m) ~ ExtType) =>
m (Exp (Lore m))
-> m (Body (Lore m)) -> m (Body (Lore m)) -> m (Exp (Lore m))
eIf (SubExp -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => SubExp -> m (Exp (Lore m))
eSubExp SubExp
cond) (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.Pattern ->
  E.Exp ->
  E.Exp ->
  (E.Exp -> InternaliseM a) ->
  InternaliseM a
internalisePat :: forall a.
[Char]
-> [SizeBinder VName]
-> PatternBase Info VName
-> Exp
-> Exp
-> (Exp -> InternaliseM a)
-> InternaliseM a
internalisePat [Char]
desc [SizeBinder VName]
sizes PatternBase Info VName
p Exp
e Exp
body Exp -> InternaliseM a
m = do
  [SubExp]
ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc' Exp
e
  [SizeBinder VName]
-> PatternBase Info VName
-> [SubExp]
-> Exp
-> (Exp -> InternaliseM a)
-> InternaliseM a
forall a.
[SizeBinder VName]
-> PatternBase Info VName
-> [SubExp]
-> Exp
-> (Exp -> InternaliseM a)
-> InternaliseM a
internalisePat' [SizeBinder VName]
sizes PatternBase Info VName
p [SubExp]
ses Exp
body Exp -> InternaliseM a
m
  where
    desc' :: [Char]
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
$ PatternBase Info VName -> Set (IdentBase Info VName)
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set (IdentBase f vn)
E.patternIdents PatternBase Info VName
p of
      [IdentBase Info VName
v] -> VName -> [Char]
baseString (VName -> [Char]) -> VName -> [Char]
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]
_ -> [Char]
desc

internalisePat' ::
  [E.SizeBinder VName] ->
  E.Pattern ->
  [I.SubExp] ->
  E.Exp ->
  (E.Exp -> InternaliseM a) ->
  InternaliseM a
internalisePat' :: forall a.
[SizeBinder VName]
-> PatternBase Info VName
-> [SubExp]
-> Exp
-> (Exp -> InternaliseM a)
-> InternaliseM a
internalisePat' [SizeBinder VName]
sizes PatternBase Info VName
p [SubExp]
ses Exp
body Exp -> 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
  PatternBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([VName] -> InternaliseM a)
-> InternaliseM a
forall a.
PatternBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([VName] -> InternaliseM a)
-> InternaliseM a
stmPattern PatternBase 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 (PatternType -> [VName] -> AppRes
AppRes (PatternBase Info VName -> PatternType
E.patternType PatternBase 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 (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [VName
v] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
se
    Exp -> InternaliseM a
m Exp
body

internaliseSlice ::
  SrcLoc ->
  [SubExp] ->
  [E.DimIndex] ->
  InternaliseM ([I.DimIndex SubExp], Certificates)
internaliseSlice :: SrcLoc
-> [SubExp]
-> [DimIndexBase Info VName]
-> InternaliseM ([DimIndex SubExp], Certificates)
internaliseSlice SrcLoc
loc [SubExp]
dims [DimIndexBase 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
 -> DimIndexBase Info VName
 -> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp]))
-> [SubExp]
-> [DimIndexBase Info VName]
-> InternaliseM [(DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM SubExp
-> DimIndexBase Info VName
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
internaliseDimIndex [SubExp]
dims [DimIndexBase Info VName]
idxs
  SubExp
ok <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"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 (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => [SubExp] -> m (Exp (Lore 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 SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt64 ([SubExp] -> [ErrorMsgPart SubExp])
-> [SubExp] -> [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
take ([DimIndexBase Info VName] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimIndexBase Info VName]
idxs) [SubExp]
dims)
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"]."]
  Certificates
c <- [Char]
-> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certificates
assert [Char]
"index_certs" SubExp
ok ErrorMsg SubExp
msg SrcLoc
loc
  ([DimIndex SubExp], Certificates)
-> InternaliseM ([DimIndex SubExp], Certificates)
forall (m :: * -> *) a. Monad m => a -> m a
return ([DimIndex SubExp]
idxs', Certificates
c)

internaliseDimIndex ::
  SubExp ->
  E.DimIndex ->
  InternaliseM (I.DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
internaliseDimIndex :: SubExp
-> DimIndexBase Info VName
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
internaliseDimIndex SubExp
w (E.DimFix Exp
i) = do
  (SubExp
i', IntType
_) <- [Char] -> Exp -> InternaliseM (SubExp, IntType)
internaliseDimExp [Char]
"i" Exp
i
  let lowerBound :: ExpT SOACS
lowerBound =
        BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
          CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
I.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 lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
          CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSlt IntType
I.Int64) SubExp
i' SubExp
w
  SubExp
ok <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"bounds_check" (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinOp
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> m (Exp (Lore m)) -> m (Exp (Lore m)) -> m (Exp (Lore m))
eBinOp BinOp
I.LogAnd (ExpT SOACS -> InternaliseM (ExpT SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ExpT SOACS
lowerBound) (ExpT SOACS -> InternaliseM (ExpT SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ExpT SOACS
upperBound)
  (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
I.DimFix SubExp
i', SubExp
ok, [SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt64 SubExp
i'])

-- Special-case an important common case that otherwise leads to horrible code.
internaliseDimIndex
  SubExp
w
  ( E.DimSlice
      Maybe Exp
Nothing
      Maybe Exp
Nothing
      (Just (E.Negate (E.IntLit Integer
1 Info PatternType
_ SrcLoc
_) SrcLoc
_))
    ) = do
    SubExp
w_minus_1 <-
      [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"w_minus_1" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
        BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
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 Exp
i Maybe Exp
j Maybe Exp
s) = do
  SubExp
s' <- InternaliseM SubExp
-> (Exp -> InternaliseM SubExp) -> Maybe Exp -> InternaliseM SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (SubExp -> InternaliseM SubExp
forall (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)
-> (Exp -> InternaliseM (SubExp, IntType))
-> Exp
-> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp -> InternaliseM (SubExp, IntType)
internaliseDimExp [Char]
"s") Maybe Exp
s
  SubExp
s_sign <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"s_sign" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (IntType -> UnOp
I.SSignum IntType
Int64) SubExp
s'
  SubExp
backwards <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"backwards" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
int64) SubExp
s_sign SubExp
negone
  SubExp
w_minus_1 <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"w_minus_1" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) SubExp
w SubExp
one
  let i_def :: InternaliseM SubExp
i_def =
        [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"i_def" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          SubExp
-> BodyT SOACS
-> BodyT SOACS
-> IfDec (BranchType SOACS)
-> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If
            SubExp
backwards
            ([SubExp] -> BodyT SOACS
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [SubExp
w_minus_1])
            ([SubExp] -> BodyT SOACS
forall lore. Bindable lore => [SubExp] -> Body lore
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 =
        [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"j_def" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          SubExp
-> BodyT SOACS
-> BodyT SOACS
-> IfDec (BranchType SOACS)
-> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If
            SubExp
backwards
            ([SubExp] -> BodyT SOACS
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [SubExp
negone])
            ([SubExp] -> BodyT SOACS
forall lore. Bindable lore => [SubExp] -> Body lore
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
-> (Exp -> InternaliseM SubExp) -> Maybe Exp -> InternaliseM SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe InternaliseM SubExp
i_def (((SubExp, IntType) -> SubExp)
-> InternaliseM (SubExp, IntType) -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SubExp, IntType) -> SubExp
forall a b. (a, b) -> a
fst (InternaliseM (SubExp, IntType) -> InternaliseM SubExp)
-> (Exp -> InternaliseM (SubExp, IntType))
-> Exp
-> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp -> InternaliseM (SubExp, IntType)
internaliseDimExp [Char]
"i") Maybe Exp
i
  SubExp
j' <- InternaliseM SubExp
-> (Exp -> InternaliseM SubExp) -> Maybe Exp -> InternaliseM SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe InternaliseM SubExp
j_def (((SubExp, IntType) -> SubExp)
-> InternaliseM (SubExp, IntType) -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SubExp, IntType) -> SubExp
forall a b. (a, b) -> a
fst (InternaliseM (SubExp, IntType) -> InternaliseM SubExp)
-> (Exp -> InternaliseM (SubExp, IntType))
-> Exp
-> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp -> InternaliseM (SubExp, IntType)
internaliseDimExp [Char]
"j") Maybe Exp
j
  SubExp
j_m_i <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"j_m_i" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
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 (Lore InternaliseM))
divRounding InternaliseM (ExpT SOACS)
x InternaliseM (ExpT SOACS)
y =
        BinOp
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> m (Exp (Lore m)) -> m (Exp (Lore m)) -> m (Exp (Lore m))
eBinOp
          (IntType -> Safety -> BinOp
SQuot IntType
Int64 Safety
Unsafe)
          ( BinOp
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> m (Exp (Lore m)) -> m (Exp (Lore m)) -> m (Exp (Lore m))
eBinOp
              (IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
I.OverflowWrap)
              InternaliseM (Exp (Lore InternaliseM))
InternaliseM (ExpT SOACS)
x
              (BinOp
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> m (Exp (Lore m)) -> m (Exp (Lore m)) -> m (Exp (Lore m))
eBinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) InternaliseM (Exp (Lore InternaliseM))
InternaliseM (ExpT SOACS)
y (InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
m (Exp (Lore m)) -> m (Exp (Lore m))
eSignum (InternaliseM (Exp (Lore InternaliseM))
 -> InternaliseM (Exp (Lore InternaliseM)))
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ SubExp -> InternaliseM (Exp (Lore InternaliseM))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Lore m))
toExp SubExp
s'))
          )
          InternaliseM (Exp (Lore InternaliseM))
InternaliseM (ExpT SOACS)
y
  SubExp
n <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"n" (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (ExpT SOACS)
-> InternaliseM (ExpT SOACS)
-> InternaliseM (Exp (Lore InternaliseM))
divRounding (SubExp -> InternaliseM (Exp (Lore InternaliseM))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Lore m))
toExp SubExp
j_m_i) (SubExp -> InternaliseM (Exp (Lore InternaliseM))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Lore m))
toExp SubExp
s')

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

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

  SubExp
zero_lte_i <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"zero_lte_i" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int64) SubExp
zero SubExp
i'
  SubExp
i_lte_j <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"i_lte_j" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int64) SubExp
i' SubExp
j'
  SubExp
forwards_ok <-
    [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"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 (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => [SubExp] -> m (Exp (Lore m))
eAll [SubExp
zero_lte_i, SubExp
zero_lte_i, SubExp
i_lte_j, SubExp
zero_leq_i_p_m_t_s, SubExp
i_p_m_t_s_lth_w]

  SubExp
negone_lte_j <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"negone_lte_j" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int64) SubExp
negone SubExp
j'
  SubExp
j_lte_i <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"j_lte_i" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int64) SubExp
j' SubExp
i'
  SubExp
backwards_ok <-
    [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"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 (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => [SubExp] -> m (Exp (Lore m))
eAll
        [SubExp
negone_lte_j, SubExp
negone_lte_j, SubExp
j_lte_i, SubExp
zero_leq_i_p_m_t_s, SubExp
i_p_m_t_s_leq_w]

  SubExp
slice_ok <-
    [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"slice_ok" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      SubExp
-> BodyT SOACS
-> BodyT SOACS
-> IfDec (BranchType SOACS)
-> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If
        SubExp
backwards
        ([SubExp] -> BodyT SOACS
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [SubExp
backwards_ok])
        ([SubExp] -> BodyT SOACS
forall lore. Bindable lore => [SubExp] -> Body lore
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 <-
    [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"ok_or_empty" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogOr SubExp
empty_slice SubExp
slice_ok

  let parts :: [ErrorMsgPart SubExp]
parts = case (Maybe Exp
i, Maybe Exp
j, Maybe Exp
s) of
        (Maybe Exp
_, Maybe Exp
_, Just {}) ->
          [ ErrorMsgPart SubExp
-> (Exp -> ErrorMsgPart SubExp) -> Maybe Exp -> ErrorMsgPart SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp
forall a b. a -> b -> a
const (ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp)
-> ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt64 SubExp
i') Maybe Exp
i,
            ErrorMsgPart SubExp
":",
            ErrorMsgPart SubExp
-> (Exp -> ErrorMsgPart SubExp) -> Maybe Exp -> ErrorMsgPart SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp
forall a b. a -> b -> a
const (ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp)
-> ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt64 SubExp
j') Maybe Exp
j,
            ErrorMsgPart SubExp
":",
            SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt64 SubExp
s'
          ]
        (Maybe Exp
_, Just {}, Maybe Exp
_) ->
          [ ErrorMsgPart SubExp
-> (Exp -> ErrorMsgPart SubExp) -> Maybe Exp -> ErrorMsgPart SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp
forall a b. a -> b -> a
const (ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp)
-> ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt64 SubExp
i') Maybe Exp
i,
            ErrorMsgPart SubExp
":",
            SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt64 SubExp
j'
          ]
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
-> (Exp -> [ErrorMsgPart SubExp])
-> Maybe Exp
-> [ErrorMsgPart SubExp]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [ErrorMsgPart SubExp]
forall a. Monoid a => a
mempty ([ErrorMsgPart SubExp] -> Exp -> [ErrorMsgPart SubExp]
forall a b. a -> b -> a
const [ErrorMsgPart SubExp
":", SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt64 SubExp
s']) Maybe Exp
s
        (Maybe Exp
_, Maybe Exp
Nothing, Maybe Exp
Nothing) ->
          [SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt64 SubExp
i', ErrorMsgPart SubExp
":"]
  (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
I.DimSlice SubExp
i' SubExp
n SubExp
s', SubExp
ok_or_empty, [ErrorMsgPart SubExp]
parts)
  where
    zero :: SubExp
zero = 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 :: [Char]
-> [Char]
-> (SubExp
    -> Lambda -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS))
-> (Exp, Exp, Exp, SrcLoc)
-> InternaliseM [SubExp]
internaliseScanOrReduce [Char]
desc [Char]
what SubExp
-> Lambda -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)
f (Exp
lam, Exp
ne, Exp
arr, SrcLoc
loc) = do
  [VName]
arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars ([Char]
what [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"_arr") Exp
arr
  [SubExp]
nes <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp ([Char]
what [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"_ne") Exp
ne
  [SubExp]
nes' <- [(SubExp, VName)]
-> ((SubExp, VName) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([SubExp] -> [VName] -> [(SubExp, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
nes [VName]
arrs) (((SubExp, VName) -> InternaliseM SubExp) -> InternaliseM [SubExp])
-> ((SubExp, VName) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \(SubExp
ne', VName
arr') -> do
    TypeBase Shape NoUniqueness
rowtype <- Int -> TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall 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 lore (m :: * -> *).
HasScope lore m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
arr'
    ErrorMsg SubExp
-> SrcLoc
-> TypeBase Shape NoUniqueness
-> [Char]
-> SubExp
-> InternaliseM SubExp
ensureShape
      ErrorMsg SubExp
"Row shape of input array does not match shape of neutral element"
      SrcLoc
loc
      TypeBase Shape NoUniqueness
rowtype
      ([Char]
what [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"_ne_right_shape")
      SubExp
ne'
  [TypeBase Shape NoUniqueness]
nests <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
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 lore (m :: * -> *).
HasScope lore m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
  Lambda
lam' <- InternaliseLambda
-> Exp
-> [TypeBase Shape NoUniqueness]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM Lambda
internaliseFoldLambda InternaliseLambda
internaliseLambda Exp
lam [TypeBase Shape NoUniqueness]
nests [TypeBase Shape NoUniqueness]
arrts
  SubExp
w <- Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 ([TypeBase Shape NoUniqueness] -> SubExp)
-> InternaliseM [TypeBase Shape NoUniqueness]
-> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall lore (m :: * -> *).
HasScope lore m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
  [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
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 lore. Op lore -> ExpT lore
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 :: [Char]
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> SrcLoc
-> InternaliseM [SubExp]
internaliseHist [Char]
desc Exp
rf Exp
hist Exp
op Exp
ne Exp
buckets Exp
img SrcLoc
loc = do
  SubExp
rf' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"hist_rf" Exp
rf
  [SubExp]
ne' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"hist_ne" Exp
ne
  [VName]
hist' <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"hist_hist" Exp
hist
  VName
buckets' <-
    [Char] -> Exp (Lore InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m VName
letExp [Char]
"hist_buckets" (ExpT SOACS -> InternaliseM VName)
-> (SubExp -> ExpT SOACS) -> SubExp -> InternaliseM VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS)
-> (SubExp -> BasicOp) -> SubExp -> ExpT SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
SubExp
      (SubExp -> InternaliseM VName)
-> InternaliseM SubExp -> InternaliseM VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"hist_buckets" Exp
buckets
  [VName]
img' <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"hist_img" Exp
img

  -- reshape neutral element to have same size as the destination array
  [SubExp]
ne_shp <- [(SubExp, VName)]
-> ((SubExp, VName) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([SubExp] -> [VName] -> [(SubExp, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
ne' [VName]
hist') (((SubExp, VName) -> InternaliseM SubExp) -> InternaliseM [SubExp])
-> ((SubExp, VName) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \(SubExp
n, VName
h) -> do
    TypeBase Shape NoUniqueness
rowtype <- Int -> TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall 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 lore (m :: * -> *).
HasScope lore m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
h
    ErrorMsg SubExp
-> SrcLoc
-> TypeBase Shape NoUniqueness
-> [Char]
-> SubExp
-> InternaliseM SubExp
ensureShape
      ErrorMsg SubExp
"Row shape of destination array does not match shape of neutral element"
      SrcLoc
loc
      TypeBase Shape NoUniqueness
rowtype
      [Char]
"hist_ne_right_shape"
      SubExp
n
  [TypeBase Shape NoUniqueness]
ne_ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
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 lore (m :: * -> *).
HasScope lore m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
hist'
  Lambda
op' <- InternaliseLambda
-> Exp
-> [TypeBase Shape NoUniqueness]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM Lambda
internaliseFoldLambda InternaliseLambda
internaliseLambda Exp
op [TypeBase Shape NoUniqueness]
ne_ts [TypeBase Shape NoUniqueness]
his_ts

  -- reshape return type of bucket function to have same size as neutral element
  -- (modulo the index)
  Param (TypeBase Shape NoUniqueness)
bucket_param <- [Char]
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"bucket_p" (TypeBase Shape NoUniqueness
 -> InternaliseM (Param (TypeBase Shape NoUniqueness)))
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64
  [Param (TypeBase Shape NoUniqueness)]
img_params <- (TypeBase Shape NoUniqueness
 -> InternaliseM (Param (TypeBase Shape NoUniqueness)))
-> [TypeBase Shape NoUniqueness]
-> InternaliseM [Param (TypeBase Shape NoUniqueness)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Char]
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"img_p" (TypeBase Shape NoUniqueness
 -> InternaliseM (Param (TypeBase Shape NoUniqueness)))
-> (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall 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 lore (m :: * -> *).
HasScope lore 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 -> [SubExp] -> BodyT SOACS
forall lore. Bindable lore => Stms lore -> [SubExp] -> Body lore
mkBody Stms SOACS
forall a. Monoid a => a
mempty ([SubExp] -> BodyT SOACS) -> [SubExp] -> BodyT SOACS
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)]
params
  Lambda
lam' <-
    [LParam (Lore InternaliseM)]
-> InternaliseM [SubExp]
-> InternaliseM (Lambda (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
[LParam (Lore m)] -> m [SubExp] -> m (Lambda (Lore m))
mkLambda [Param (TypeBase Shape NoUniqueness)]
[LParam (Lore InternaliseM)]
params (InternaliseM [SubExp]
 -> InternaliseM (Lambda (Lore InternaliseM)))
-> InternaliseM [SubExp]
-> InternaliseM (Lambda (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$
      ErrorMsg SubExp
-> SrcLoc
-> [TypeBase Shape NoUniqueness]
-> [SubExp]
-> InternaliseM [SubExp]
ensureResultShape
        ErrorMsg SubExp
"Row shape of value array does not match row shape of hist target"
        (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
img)
        [TypeBase Shape NoUniqueness]
rettype
        ([SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Body (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *). MonadBinder m => Body (Lore m) -> m [SubExp]
bodyBind Body (Lore 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 lore (m :: * -> *).
HasScope lore 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 lore (m :: * -> *).
HasScope lore 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 lore (m :: * -> *).
HasScope lore 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 <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"bucket_cmp" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
I.int64) SubExp
b_w SubExp
w_img
  Certificates
c <-
    [Char]
-> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certificates
assert
      [Char]
"bucket_cert"
      SubExp
cmp
      ErrorMsg SubExp
"length of index and value array does not match"
      SrcLoc
loc
  VName
buckets'' <-
    Certificates -> InternaliseM VName -> InternaliseM VName
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
c (InternaliseM VName -> InternaliseM VName)
-> InternaliseM VName -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$
      [Char] -> Exp (Lore InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m VName
letExp (VName -> [Char]
baseString VName
buckets') (Exp (Lore InternaliseM) -> InternaliseM VName)
-> Exp (Lore InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$
        BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
I.Reshape (ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp
reshapeOuter [SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimCoercion SubExp
w_img] Int
1 Shape
b_shape) VName
buckets'

  [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
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 lore. Op lore -> ExpT lore
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 lore.
SubExp -> [HistOp lore] -> Lambda lore -> [VName] -> SOAC lore
I.Hist SubExp
w_img [SubExp -> SubExp -> [VName] -> [SubExp] -> Lambda -> HistOp SOACS
forall lore.
SubExp
-> SubExp -> [VName] -> [SubExp] -> Lambda lore -> HistOp lore
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 :: [Char] -> StreamOrd -> Exp -> Exp -> InternaliseM [SubExp]
internaliseStreamMap [Char]
desc StreamOrd
o Exp
lam Exp
arr = do
  [VName]
arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"stream_input" Exp
arr
  Lambda
lam' <- InternaliseLambda -> Exp -> [SubExp] -> InternaliseM Lambda
internaliseStreamMapLambda InternaliseLambda
internaliseLambda Exp
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 lore (m :: * -> *).
HasScope lore m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
  let form :: StreamForm SOACS
form = StreamOrd -> Commutativity -> Lambda -> StreamForm SOACS
forall lore.
StreamOrd -> Commutativity -> Lambda lore -> StreamForm lore
I.Parallel StreamOrd
o Commutativity
Commutative ([LParam] -> BodyT SOACS -> [TypeBase Shape NoUniqueness] -> Lambda
forall lore.
[LParam lore]
-> BodyT lore -> [TypeBase Shape NoUniqueness] -> LambdaT lore
I.Lambda [] (Stms SOACS -> [SubExp] -> BodyT SOACS
forall lore. Bindable lore => Stms lore -> [SubExp] -> Body lore
mkBody Stms SOACS
forall a. Monoid a => a
mempty []) [])
  [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Op SOACS -> ExpT SOACS
forall lore. Op lore -> ExpT lore
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp
-> [VName] -> StreamForm SOACS -> [SubExp] -> Lambda -> SOAC SOACS
forall lore.
SubExp
-> [VName]
-> StreamForm lore
-> [SubExp]
-> Lambda lore
-> SOAC lore
I.Stream SubExp
w [VName]
arrs StreamForm SOACS
form [] Lambda
lam'

internaliseStreamRed ::
  String ->
  StreamOrd ->
  Commutativity ->
  E.Exp ->
  E.Exp ->
  E.Exp ->
  InternaliseM [SubExp]
internaliseStreamRed :: [Char]
-> StreamOrd
-> Commutativity
-> Exp
-> Exp
-> Exp
-> InternaliseM [SubExp]
internaliseStreamRed [Char]
desc StreamOrd
o Commutativity
comm Exp
lam0 Exp
lam Exp
arr = do
  [VName]
arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"stream_input" Exp
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 lore (m :: * -> *).
HasScope lore m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType) [VName]
arrs
  ([Param (TypeBase Shape NoUniqueness)]
lam_params, BodyT SOACS
lam_body) <-
    InternaliseLambda
-> Exp
-> [TypeBase Shape NoUniqueness]
-> InternaliseM ([LParam], BodyT SOACS)
internaliseStreamLambda InternaliseLambda
internaliseLambda Exp
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 (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [Param (TypeBase Shape NoUniqueness) -> VName
forall dec. Param dec -> VName
I.paramName Param (TypeBase Shape NoUniqueness)
chunk_param] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
    BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ 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 (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [Param (TypeBase Shape NoUniqueness) -> VName
forall dec. Param dec -> VName
I.paramName Param (TypeBase Shape NoUniqueness)
p] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
        PrimType -> [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] -> BasicOp) -> [SubExp] -> BasicOp
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
  [SubExp]
nes <- BodyT SOACS -> InternaliseM [SubExp]
forall (m :: * -> *). MonadBinder m => Body (Lore m) -> m [SubExp]
bodyBind (BodyT SOACS -> InternaliseM [SubExp])
-> InternaliseM (BodyT SOACS) -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BodyT SOACS -> InternaliseM (BodyT SOACS)
forall lore (m :: * -> *).
(Renameable lore, MonadFreshNames m) =>
Body lore -> m (Body lore)
renameBody BodyT SOACS
lam_body

  [TypeBase Shape NoUniqueness]
nes_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]
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 lore (m :: * -> *).
HasScope lore 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
-> Exp
-> [TypeBase Shape NoUniqueness]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM Lambda
internaliseFoldLambda InternaliseLambda
internaliseLambda Exp
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 ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
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 lore. LambdaT lore -> [LParam lore]
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 <- [Char] -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName ([Char] -> InternaliseM VName) -> [Char] -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$ VName -> [Char]
baseString (VName -> [Char]) -> VName -> [Char]
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 (Lore InternaliseM)]
-> InternaliseM [SubExp]
-> InternaliseM (Lambda (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
[LParam (Lore m)] -> m [SubExp] -> m (Lambda (Lore m))
mkLambda [Param (TypeBase Shape NoUniqueness)]
[LParam (Lore InternaliseM)]
lam_params' (InternaliseM [SubExp]
 -> InternaliseM (Lambda (Lore InternaliseM)))
-> InternaliseM [SubExp]
-> InternaliseM (Lambda (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ do
    [SubExp]
lam_res <- Body (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *). MonadBinder m => Body (Lore m) -> m [SubExp]
bodyBind Body (Lore 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"
        (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
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 lore. LambdaT lore -> [LParam lore]
I.lambdaParams Lambda
lam0')
        [SubExp]
lam_res
    ErrorMsg SubExp
-> SrcLoc
-> [TypeBase Shape NoUniqueness]
-> [SubExp]
-> InternaliseM [SubExp]
ensureResultShape
      ErrorMsg SubExp
"shape of result does not match shape of initial value"
      (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
lam0)
      [TypeBase Shape NoUniqueness]
nes_ts
      ([SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ( Lambda (Lore InternaliseM)
-> [InternaliseM (Exp (Lore InternaliseM))]
-> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
Lambda (Lore m) -> [m (Exp (Lore m))] -> m [SubExp]
eLambda Lambda (Lore InternaliseM)
Lambda
lam0' ([InternaliseM (ExpT SOACS)] -> InternaliseM [SubExp])
-> ([SubExp] -> [InternaliseM (ExpT SOACS)])
-> [SubExp]
-> InternaliseM [SubExp]
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 :: * -> *). MonadBinder m => SubExp -> m (Exp (Lore m))
eSubExp ([SubExp] -> InternaliseM [SubExp])
-> [SubExp] -> InternaliseM [SubExp]
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 lore.
StreamOrd -> Commutativity -> Lambda lore -> StreamForm lore
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 lore (m :: * -> *).
HasScope lore m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
  [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Op SOACS -> ExpT SOACS
forall lore. Op lore -> ExpT lore
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp
-> [VName] -> StreamForm SOACS -> [SubExp] -> Lambda -> SOAC SOACS
forall lore.
SubExp
-> [VName]
-> StreamForm lore
-> [SubExp]
-> Lambda lore
-> SOAC lore
I.Stream SubExp
w [VName]
arrs StreamForm SOACS
form [SubExp]
nes Lambda
lam'

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

  VName
acc_cert_v <- [Char] -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"acc_cert"
  [TypeBase Shape NoUniqueness]
dest_ts <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall lore (m :: * -> *).
HasScope lore 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 <- [Char]
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"acc_p" TypeBase Shape NoUniqueness
acc_t
  Lambda
withacc_lam <- [LParam (Lore InternaliseM)]
-> InternaliseM [SubExp]
-> InternaliseM (Lambda (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
[LParam (Lore m)] -> m [SubExp] -> m (Lambda (Lore 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 (Lore InternaliseM)
acc_p] (InternaliseM [SubExp]
 -> InternaliseM (Lambda (Lore InternaliseM)))
-> InternaliseM [SubExp]
-> InternaliseM (Lambda (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ do
    Lambda
lam' <-
      InternaliseLambda -> Exp -> [SubExp] -> InternaliseM Lambda
internaliseMapLambda InternaliseLambda
internaliseLambda Exp
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 lore (m :: * -> *).
HasScope lore m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
bs'
    [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
"acc_res" (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Op SOACS -> ExpT SOACS
forall lore. Op lore -> ExpT lore
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
forall lore. SubExp -> [VName] -> ScremaForm lore -> SOAC lore
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 lore. Lambda lore -> ScremaForm lore
I.mapSOAC Lambda
lam')

  Maybe (Lambda, [SubExp])
op' <-
    case Maybe (Exp, Exp)
op of
      Just (Exp
op_lam, Exp
ne) -> do
        [SubExp]
ne' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"hist_ne" Exp
ne
        [TypeBase Shape NoUniqueness]
ne_ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
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 Exp
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 <- [Char]
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"idx" (TypeBase Shape NoUniqueness
 -> InternaliseM (Param (TypeBase Shape NoUniqueness)))
-> TypeBase Shape NoUniqueness
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64
        let op_lam' :: Lambda
op_lam' = [LParam] -> BodyT SOACS -> [TypeBase Shape NoUniqueness] -> Lambda
forall lore.
[LParam lore]
-> BodyT lore -> [TypeBase Shape NoUniqueness] -> LambdaT lore
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 (Exp, Exp)
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 lore (m :: * -> *).
HasScope lore 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
$
    [Char] -> Exp (Lore InternaliseM) -> InternaliseM [VName]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [VName]
letTupExp [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM [VName])
-> Exp (Lore InternaliseM) -> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$ [(Shape, [VName], Maybe (Lambda, [SubExp]))]
-> Lambda -> ExpT SOACS
forall lore.
[(Shape, [VName], Maybe (Lambda lore, [SubExp]))]
-> Lambda lore -> ExpT lore
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 :: [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
desc Exp
e = do
  [SubExp]
vs <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
  case [SubExp]
vs of
    [SubExp
se] -> SubExp -> InternaliseM SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return SubExp
se
    [SubExp]
_ -> [Char] -> InternaliseM SubExp
forall a. HasCallStack => [Char] -> a
error [Char]
"Internalise.internaliseExp1: was passed not just a single subexpression"

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

internaliseExpToVars :: String -> E.Exp -> InternaliseM [I.VName]
internaliseExpToVars :: [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
desc Exp
e =
  (SubExp -> InternaliseM VName) -> [SubExp] -> InternaliseM [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM VName
asIdent ([SubExp] -> InternaliseM [VName])
-> InternaliseM [SubExp] -> InternaliseM [VName]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
  where
    asIdent :: SubExp -> InternaliseM VName
asIdent (I.Var VName
v) = VName -> InternaliseM VName
forall (m :: * -> *) a. Monad m => a -> m a
return VName
v
    asIdent SubExp
se = [Char] -> Exp (Lore InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m VName
letExp [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM VName)
-> Exp (Lore InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
se

internaliseOperation ::
  String ->
  E.Exp ->
  (I.VName -> InternaliseM I.BasicOp) ->
  InternaliseM [I.SubExp]
internaliseOperation :: [Char]
-> Exp -> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
internaliseOperation [Char]
s Exp
e VName -> InternaliseM BasicOp
op = do
  [VName]
vs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
s Exp
e
  [Char] -> [Exp (Lore InternaliseM)] -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> [Exp (Lore m)] -> m [SubExp]
letSubExps [Char]
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 lore. BasicOp -> ExpT lore
I.BasicOp (InternaliseM BasicOp -> InternaliseM (ExpT SOACS))
-> (VName -> InternaliseM BasicOp)
-> VName
-> InternaliseM (ExpT SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> InternaliseM BasicOp
op) [VName]
vs

certifyingNonzero ::
  SrcLoc ->
  IntType ->
  SubExp ->
  InternaliseM a ->
  InternaliseM a
certifyingNonzero :: forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
x InternaliseM a
m = do
  SubExp
zero <-
    [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"zero" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
        CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp (PrimType -> CmpOp
CmpEq (IntType -> PrimType
IntType IntType
t)) SubExp
x (IntType -> Integer -> SubExp
intConst IntType
t Integer
0)
  SubExp
nonzero <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"nonzero" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
UnOp UnOp
Not SubExp
zero
  Certificates
c <- [Char]
-> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certificates
assert [Char]
"nonzero_cert" SubExp
nonzero ErrorMsg SubExp
"division by zero" SrcLoc
loc
  Certificates -> InternaliseM a -> InternaliseM a
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
c InternaliseM a
m

certifyingNonnegative ::
  SrcLoc ->
  IntType ->
  SubExp ->
  InternaliseM a ->
  InternaliseM a
certifyingNonnegative :: forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonnegative SrcLoc
loc IntType
t SubExp
x InternaliseM a
m = do
  SubExp
nonnegative <-
    [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"nonnegative" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
        CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp (IntType -> CmpOp
CmpSle IntType
t) (IntType -> Integer -> SubExp
intConst IntType
t Integer
0) SubExp
x
  Certificates
c <- [Char]
-> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certificates
assert [Char]
"nonzero_cert" SubExp
nonnegative ErrorMsg SubExp
"negative exponent" SrcLoc
loc
  Certificates -> InternaliseM a -> InternaliseM a
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
c InternaliseM a
m

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

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

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

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

findFuncall ::
  E.AppExp ->
  InternaliseM
    ( E.QualName VName,
      [(E.Exp, Maybe VName)]
    )
findFuncall :: AppExp -> InternaliseM (QualName VName, [(Exp, Maybe VName)])
findFuncall (E.Apply Exp
f Exp
arg (Info (Diet
_, Maybe VName
argext)) SrcLoc
_)
  | E.AppExp AppExp
f_e Info AppRes
_ <- Exp
f = do
    (QualName VName
fname, [(Exp, Maybe VName)]
args) <- AppExp -> InternaliseM (QualName VName, [(Exp, Maybe VName)])
findFuncall AppExp
f_e
    (QualName VName, [(Exp, Maybe VName)])
-> InternaliseM (QualName VName, [(Exp, Maybe VName)])
forall (m :: * -> *) a. Monad m => a -> m a
return (QualName VName
fname, [(Exp, Maybe VName)]
args [(Exp, Maybe VName)]
-> [(Exp, Maybe VName)] -> [(Exp, Maybe VName)]
forall a. [a] -> [a] -> [a]
++ [(Exp
arg, Maybe VName
argext)])
  | E.Var QualName VName
fname Info PatternType
_ SrcLoc
_ <- Exp
f =
    (QualName VName, [(Exp, Maybe VName)])
-> InternaliseM (QualName VName, [(Exp, Maybe VName)])
forall (m :: * -> *) a. Monad m => a -> m a
return (QualName VName
fname, [(Exp
arg, Maybe VName
argext)])
findFuncall AppExp
e =
  [Char] -> InternaliseM (QualName VName, [(Exp, Maybe VName)])
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM (QualName VName, [(Exp, Maybe VName)]))
-> [Char] -> InternaliseM (QualName VName, [(Exp, Maybe VName)])
forall a b. (a -> b) -> a -> b
$ [Char]
"Invalid function expression in application: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ AppExp -> [Char]
forall a. Pretty a => a -> [Char]
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 [SubExp]
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 (ShapeBase ExtSize) 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 lore (m :: * -> *) a.
ExtendedScope lore m a -> Scope lore -> m a
extendedScope ((SubExp
 -> ExtendedScope SOACS InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp]
-> ExtendedScope SOACS InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse SubExp
-> ExtendedScope SOACS InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
res) Scope SOACS
stmsscope
  where
    stmsscope :: Scope SOACS
stmsscope = Stms SOACS -> Scope SOACS
forall lore a. Scoped lore a => a -> Scope lore
scopeOf Stms SOACS
stms

internaliseLambda :: InternaliseLambda
internaliseLambda :: InternaliseLambda
internaliseLambda (E.Parens Exp
e SrcLoc
_) [TypeBase Shape NoUniqueness]
rowtypes =
  InternaliseLambda
internaliseLambda Exp
e [TypeBase Shape NoUniqueness]
rowtypes
internaliseLambda (E.Lambda [PatternBase Info VName]
params Exp
body Maybe (TypeExp VName)
_ (Info (Aliasing
_, StructType
rettype)) SrcLoc
_) [TypeBase Shape NoUniqueness]
rowtypes =
  [PatternBase 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.
[PatternBase Info VName]
-> [TypeBase Shape NoUniqueness]
-> ([LParam] -> InternaliseM a)
-> InternaliseM a
bindingLambdaParams [PatternBase 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' <- [Char] -> Exp -> InternaliseM (BodyT SOACS)
internaliseBody [Char]
"lam" Exp
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 Exp
e [TypeBase Shape NoUniqueness]
_ = [Char]
-> InternaliseM
     ([Param (TypeBase Shape NoUniqueness)], BodyT SOACS,
      [TypeBase Shape NoUniqueness])
forall a. HasCallStack => [Char] -> a
error ([Char]
 -> InternaliseM
      ([Param (TypeBase Shape NoUniqueness)], BodyT SOACS,
       [TypeBase Shape NoUniqueness]))
-> [Char]
-> InternaliseM
     ([Param (TypeBase Shape NoUniqueness)], BodyT SOACS,
      [TypeBase Shape NoUniqueness])
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseLambda: unexpected expression:\n" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Exp -> [Char]
forall a. Pretty a => a -> [Char]
pretty Exp
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
-> [Exp] -> SrcLoc -> Maybe ([Char] -> InternaliseM [SubExp])
isOverloadedFunction QualName VName
qname [Exp]
args SrcLoc
loc = do
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ VName -> Int
baseTag (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qname) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag
  let handlers :: [[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])]
handlers =
        [ [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
forall {a}.
(Eq a, IsString a) =>
[Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleSign,
          [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
forall {f :: * -> *}.
Applicative f =>
[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM (f SubExp))
handleIntrinsicOps,
          [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
handleOps,
          [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
handleSOACs,
          [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
forall {a}.
(Eq a, IsString a) =>
[Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleAccs,
          [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
handleRest
        ]
  [Maybe ([Char] -> InternaliseM [SubExp])]
-> Maybe ([Char] -> InternaliseM [SubExp])
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, MonadPlus m) =>
t (m a) -> m a
msum [[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
h [Exp]
args ([Char] -> Maybe ([Char] -> InternaliseM [SubExp]))
-> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ VName -> [Char]
baseString (VName -> [Char]) -> VName -> [Char]
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qname | [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
h <- [[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])]
handlers]
  where
    handleSign :: [Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleSign [Exp
x] a
"sign_i8" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
I.Int8 Exp
x
    handleSign [Exp
x] a
"sign_i16" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
I.Int16 Exp
x
    handleSign [Exp
x] a
"sign_i32" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
I.Int32 Exp
x
    handleSign [Exp
x] a
"sign_i64" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
I.Int64 Exp
x
    handleSign [Exp
x] a
"unsign_i8" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
I.Int8 Exp
x
    handleSign [Exp
x] a
"unsign_i16" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
I.Int16 Exp
x
    handleSign [Exp
x] a
"unsign_i32" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
I.Int32 Exp
x
    handleSign [Exp
x] a
"unsign_i64" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
I.Int64 Exp
x
    handleSign [Exp]
_ a
_ = Maybe ([Char] -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

    handleIntrinsicOps :: [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM (f SubExp))
handleIntrinsicOps [Exp
x] [Char]
s
      | Just UnOp
unop <- (UnOp -> Bool) -> [UnOp] -> Maybe UnOp
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (([Char] -> [Char] -> Bool
forall a. Eq a => a -> a -> Bool
== [Char]
s) ([Char] -> Bool) -> (UnOp -> [Char]) -> UnOp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UnOp -> [Char]
forall a. Pretty a => a -> [Char]
pretty) [UnOp]
allUnOps = ([Char] -> InternaliseM (f SubExp))
-> Maybe ([Char] -> InternaliseM (f SubExp))
forall a. a -> Maybe a
Just (([Char] -> InternaliseM (f SubExp))
 -> Maybe ([Char] -> InternaliseM (f SubExp)))
-> ([Char] -> InternaliseM (f SubExp))
-> Maybe ([Char] -> InternaliseM (f SubExp))
forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
        SubExp
x' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"x" Exp
x
        (SubExp -> f SubExp)
-> InternaliseM SubExp -> InternaliseM (f SubExp)
forall (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
$ [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
unop SubExp
x'
    handleIntrinsicOps [TupLit [Exp
x, Exp
y] SrcLoc
_] [Char]
s
      | Just BinOp
bop <- (BinOp -> Bool) -> [BinOp] -> Maybe BinOp
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (([Char] -> [Char] -> Bool
forall a. Eq a => a -> a -> Bool
== [Char]
s) ([Char] -> Bool) -> (BinOp -> [Char]) -> BinOp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BinOp -> [Char]
forall a. Pretty a => a -> [Char]
pretty) [BinOp]
allBinOps = ([Char] -> InternaliseM (f SubExp))
-> Maybe ([Char] -> InternaliseM (f SubExp))
forall a. a -> Maybe a
Just (([Char] -> InternaliseM (f SubExp))
 -> Maybe ([Char] -> InternaliseM (f SubExp)))
-> ([Char] -> InternaliseM (f SubExp))
-> Maybe ([Char] -> InternaliseM (f SubExp))
forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
        SubExp
x' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"x" Exp
x
        SubExp
y' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"y" Exp
y
        (SubExp -> f SubExp)
-> InternaliseM SubExp -> InternaliseM (f SubExp)
forall (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
$ [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
bop SubExp
x' SubExp
y'
      | Just CmpOp
cmp <- (CmpOp -> Bool) -> [CmpOp] -> Maybe CmpOp
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (([Char] -> [Char] -> Bool
forall a. Eq a => a -> a -> Bool
== [Char]
s) ([Char] -> Bool) -> (CmpOp -> [Char]) -> CmpOp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CmpOp -> [Char]
forall a. Pretty a => a -> [Char]
pretty) [CmpOp]
allCmpOps = ([Char] -> InternaliseM (f SubExp))
-> Maybe ([Char] -> InternaliseM (f SubExp))
forall a. a -> Maybe a
Just (([Char] -> InternaliseM (f SubExp))
 -> Maybe ([Char] -> InternaliseM (f SubExp)))
-> ([Char] -> InternaliseM (f SubExp))
-> Maybe ([Char] -> InternaliseM (f SubExp))
forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
        SubExp
x' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"x" Exp
x
        SubExp
y' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"y" Exp
y
        (SubExp -> f SubExp)
-> InternaliseM SubExp -> InternaliseM (f SubExp)
forall (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
$ [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp CmpOp
cmp SubExp
x' SubExp
y'
    handleIntrinsicOps [Exp
x] [Char]
s
      | Just ConvOp
conv <- (ConvOp -> Bool) -> [ConvOp] -> Maybe ConvOp
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (([Char] -> [Char] -> Bool
forall a. Eq a => a -> a -> Bool
== [Char]
s) ([Char] -> Bool) -> (ConvOp -> [Char]) -> ConvOp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ConvOp -> [Char]
forall a. Pretty a => a -> [Char]
pretty) [ConvOp]
allConvOps = ([Char] -> InternaliseM (f SubExp))
-> Maybe ([Char] -> InternaliseM (f SubExp))
forall a. a -> Maybe a
Just (([Char] -> InternaliseM (f SubExp))
 -> Maybe ([Char] -> InternaliseM (f SubExp)))
-> ([Char] -> InternaliseM (f SubExp))
-> Maybe ([Char] -> InternaliseM (f SubExp))
forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
        SubExp
x' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"x" Exp
x
        (SubExp -> f SubExp)
-> InternaliseM SubExp -> InternaliseM (f SubExp)
forall (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
$ [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp ConvOp
conv SubExp
x'
    handleIntrinsicOps [Exp]
_ [Char]
_ = Maybe ([Char] -> InternaliseM (f SubExp))
forall a. Maybe a
Nothing

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

        doComparison :: [Char] -> SubExp -> SubExp -> m SubExp
doComparison [Char]
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 -> [Char] -> Exp (Lore m) -> m SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
desc (Exp (Lore m) -> m SubExp) -> Exp (Lore m) -> m SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
t) SubExp
x SubExp
y
            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) ->
                [Char] -> Exp (Lore m) -> m SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"dim_eq" (Exp (Lore m) -> m SubExp) -> Exp (Lore m) -> m SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
int64) SubExp
x_dim SubExp
y_dim
              SubExp
shapes_match <- [Char] -> Exp (Lore m) -> m SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"shapes_match" (Exp (Lore m) -> m SubExp) -> m (Exp (Lore m)) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [SubExp] -> m (Exp (Lore m))
forall (m :: * -> *). MonadBinder m => [SubExp] -> m (Exp (Lore m))
eAll [SubExp]
dims_match
              Body (Lore m)
compare_elems_body <- Binder (Lore m) (Body (Lore m)) -> m (Body (Lore m))
forall lore (m :: * -> *) somelore.
(Bindable lore, MonadFreshNames m, HasScope somelore m,
 SameScope somelore lore) =>
Binder lore (Body lore) -> m (Body lore)
runBodyBinder (Binder (Lore m) (Body (Lore m)) -> m (Body (Lore m)))
-> Binder (Lore m) (Body (Lore m)) -> m (Body (Lore m))
forall a b. (a -> b) -> a -> b
$ do
                -- Flatten both x and y.
                SubExp
x_num_elems <-
                  [Char]
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"x_num_elems"
                    (Exp (Lore m) -> BinderT (Lore m) (State VNameSource) SubExp)
-> BinderT (Lore m) (State VNameSource) (Exp (Lore m))
-> BinderT (Lore m) (State VNameSource) SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinOp
-> SubExp
-> [SubExp]
-> BinderT
     (Lore m)
     (State VNameSource)
     (Exp (Lore (BinderT (Lore m) (State VNameSource))))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> SubExp -> [SubExp] -> m (Exp (Lore m))
foldBinOp (IntType -> Overflow -> BinOp
I.Mul IntType
Int64 Overflow
I.OverflowUndef) (Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
1 :: Int64)) [SubExp]
x_dims
                VName
x' <- [Char]
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m VName
letExp [Char]
"x" (Exp (Lore (BinderT (Lore m) (State VNameSource)))
 -> BinderT (Lore m) (State VNameSource) VName)
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
x
                VName
y' <- [Char]
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m VName
letExp [Char]
"x" (Exp (Lore (BinderT (Lore m) (State VNameSource)))
 -> BinderT (Lore m) (State VNameSource) VName)
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
y
                VName
x_flat <- [Char]
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m VName
letExp [Char]
"x_flat" (Exp (Lore (BinderT (Lore m) (State VNameSource)))
 -> BinderT (Lore m) (State VNameSource) VName)
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
I.Reshape [SubExp -> DimChange SubExp
forall d. d -> DimChange d
I.DimNew SubExp
x_num_elems] VName
x'
                VName
y_flat <- [Char]
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m VName
letExp [Char]
"y_flat" (Exp (Lore (BinderT (Lore m) (State VNameSource)))
 -> BinderT (Lore m) (State VNameSource) VName)
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
I.Reshape [SubExp -> DimChange SubExp
forall d. d -> DimChange d
I.DimNew SubExp
x_num_elems] VName
y'

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

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

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

    handleSOACs :: [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
handleSOACs [TupLit [Exp
lam, Exp
arr] SrcLoc
_] [Char]
"map" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      [VName]
arr' <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"map_arr" Exp
arr
      Lambda
lam' <- InternaliseLambda -> Exp -> [SubExp] -> InternaliseM Lambda
internaliseMapLambda InternaliseLambda
internaliseLambda Exp
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 lore (m :: * -> *).
HasScope lore m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arr'
      [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
        Op SOACS -> ExpT SOACS
forall lore. Op lore -> ExpT lore
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
          SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
forall lore. SubExp -> [VName] -> ScremaForm lore -> SOAC lore
I.Screma SubExp
w [VName]
arr' (Lambda -> ScremaForm SOACS
forall lore. Lambda lore -> ScremaForm lore
I.mapSOAC Lambda
lam')
    handleSOACs [TupLit [Exp
k, Exp
lam, Exp
arr] SrcLoc
_] [Char]
"partition" = do
      Int
k' <- Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int32 -> Int) -> Maybe Int32 -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> Maybe Int32
forall {vn}. ExpBase Info vn -> Maybe Int32
fromInt32 Exp
k
      ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
_desc -> do
        [VName]
arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"partition_input" Exp
arr
        Lambda
lam' <- InternaliseLambda -> Int -> Exp -> [SubExp] -> InternaliseM Lambda
internalisePartitionLambda InternaliseLambda
internaliseLambda Int
k' Exp
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 [Exp
lam, Exp
ne, Exp
arr] SrcLoc
_] [Char]
"reduce" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char]
-> [Char]
-> (SubExp
    -> Lambda -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS))
-> (Exp, Exp, Exp, SrcLoc)
-> InternaliseM [SubExp]
internaliseScanOrReduce [Char]
desc [Char]
"reduce" SubExp
-> Lambda -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)
forall {f :: * -> *} {lore}.
(Bindable lore, MonadFreshNames f) =>
SubExp -> Lambda lore -> [SubExp] -> [VName] -> f (SOAC lore)
reduce (Exp
lam, Exp
ne, Exp
arr, SrcLoc
loc)
      where
        reduce :: SubExp -> Lambda lore -> [SubExp] -> [VName] -> f (SOAC lore)
reduce SubExp
w Lambda lore
red_lam [SubExp]
nes [VName]
arrs =
          SubExp -> [VName] -> ScremaForm lore -> SOAC lore
forall lore. SubExp -> [VName] -> ScremaForm lore -> SOAC lore
I.Screma SubExp
w [VName]
arrs
            (ScremaForm lore -> SOAC lore)
-> f (ScremaForm lore) -> f (SOAC lore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Reduce lore] -> f (ScremaForm lore)
forall lore (m :: * -> *).
(Bindable lore, MonadFreshNames m) =>
[Reduce lore] -> m (ScremaForm lore)
I.reduceSOAC [Commutativity -> Lambda lore -> [SubExp] -> Reduce lore
forall lore.
Commutativity -> Lambda lore -> [SubExp] -> Reduce lore
Reduce Commutativity
Noncommutative Lambda lore
red_lam [SubExp]
nes]
    handleSOACs [TupLit [Exp
lam, Exp
ne, Exp
arr] SrcLoc
_] [Char]
"reduce_comm" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char]
-> [Char]
-> (SubExp
    -> Lambda -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS))
-> (Exp, Exp, Exp, SrcLoc)
-> InternaliseM [SubExp]
internaliseScanOrReduce [Char]
desc [Char]
"reduce" SubExp
-> Lambda -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)
forall {f :: * -> *} {lore}.
(Bindable lore, MonadFreshNames f) =>
SubExp -> Lambda lore -> [SubExp] -> [VName] -> f (SOAC lore)
reduce (Exp
lam, Exp
ne, Exp
arr, SrcLoc
loc)
      where
        reduce :: SubExp -> Lambda lore -> [SubExp] -> [VName] -> f (SOAC lore)
reduce SubExp
w Lambda lore
red_lam [SubExp]
nes [VName]
arrs =
          SubExp -> [VName] -> ScremaForm lore -> SOAC lore
forall lore. SubExp -> [VName] -> ScremaForm lore -> SOAC lore
I.Screma SubExp
w [VName]
arrs
            (ScremaForm lore -> SOAC lore)
-> f (ScremaForm lore) -> f (SOAC lore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Reduce lore] -> f (ScremaForm lore)
forall lore (m :: * -> *).
(Bindable lore, MonadFreshNames m) =>
[Reduce lore] -> m (ScremaForm lore)
I.reduceSOAC [Commutativity -> Lambda lore -> [SubExp] -> Reduce lore
forall lore.
Commutativity -> Lambda lore -> [SubExp] -> Reduce lore
Reduce Commutativity
Commutative Lambda lore
red_lam [SubExp]
nes]
    handleSOACs [TupLit [Exp
lam, Exp
ne, Exp
arr] SrcLoc
_] [Char]
"scan" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char]
-> [Char]
-> (SubExp
    -> Lambda -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS))
-> (Exp, Exp, Exp, SrcLoc)
-> InternaliseM [SubExp]
internaliseScanOrReduce [Char]
desc [Char]
"scan" SubExp
-> Lambda -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)
forall {f :: * -> *} {lore}.
(Bindable lore, MonadFreshNames f) =>
SubExp -> Lambda lore -> [SubExp] -> [VName] -> f (SOAC lore)
reduce (Exp
lam, Exp
ne, Exp
arr, SrcLoc
loc)
      where
        reduce :: SubExp -> Lambda lore -> [SubExp] -> [VName] -> f (SOAC lore)
reduce SubExp
w Lambda lore
scan_lam [SubExp]
nes [VName]
arrs =
          SubExp -> [VName] -> ScremaForm lore -> SOAC lore
forall lore. SubExp -> [VName] -> ScremaForm lore -> SOAC lore
I.Screma SubExp
w [VName]
arrs (ScremaForm lore -> SOAC lore)
-> f (ScremaForm lore) -> f (SOAC lore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Scan lore] -> f (ScremaForm lore)
forall lore (m :: * -> *).
(Bindable lore, MonadFreshNames m) =>
[Scan lore] -> m (ScremaForm lore)
I.scanSOAC [Lambda lore -> [SubExp] -> Scan lore
forall lore. Lambda lore -> [SubExp] -> Scan lore
Scan Lambda lore
scan_lam [SubExp]
nes]
    handleSOACs [TupLit [Exp
op, Exp
f, Exp
arr] SrcLoc
_] [Char]
"reduce_stream" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char]
-> StreamOrd
-> Commutativity
-> Exp
-> Exp
-> Exp
-> InternaliseM [SubExp]
internaliseStreamRed [Char]
desc StreamOrd
InOrder Commutativity
Noncommutative Exp
op Exp
f Exp
arr
    handleSOACs [TupLit [Exp
op, Exp
f, Exp
arr] SrcLoc
_] [Char]
"reduce_stream_per" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char]
-> StreamOrd
-> Commutativity
-> Exp
-> Exp
-> Exp
-> InternaliseM [SubExp]
internaliseStreamRed [Char]
desc StreamOrd
Disorder Commutativity
Commutative Exp
op Exp
f Exp
arr
    handleSOACs [TupLit [Exp
f, Exp
arr] SrcLoc
_] [Char]
"map_stream" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char] -> StreamOrd -> Exp -> Exp -> InternaliseM [SubExp]
internaliseStreamMap [Char]
desc StreamOrd
InOrder Exp
f Exp
arr
    handleSOACs [TupLit [Exp
f, Exp
arr] SrcLoc
_] [Char]
"map_stream_per" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char] -> StreamOrd -> Exp -> Exp -> InternaliseM [SubExp]
internaliseStreamMap [Char]
desc StreamOrd
Disorder Exp
f Exp
arr
    handleSOACs [TupLit [Exp
rf, Exp
dest, Exp
op, Exp
ne, Exp
buckets, Exp
img] SrcLoc
_] [Char]
"hist" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char]
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> SrcLoc
-> InternaliseM [SubExp]
internaliseHist [Char]
desc Exp
rf Exp
dest Exp
op Exp
ne Exp
buckets Exp
img SrcLoc
loc
    handleSOACs [Exp]
_ [Char]
_ = Maybe ([Char] -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

    handleAccs :: [Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleAccs [TupLit [Exp
dest, Exp
f, Exp
bs] SrcLoc
_] a
"scatter_stream" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char]
-> Exp -> Maybe (Exp, Exp) -> Exp -> Exp -> InternaliseM [SubExp]
internaliseStreamAcc [Char]
desc Exp
dest Maybe (Exp, Exp)
forall a. Maybe a
Nothing Exp
f Exp
bs
    handleAccs [TupLit [Exp
dest, Exp
op, Exp
ne, Exp
f, Exp
bs] SrcLoc
_] a
"hist_stream" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char]
-> Exp -> Maybe (Exp, Exp) -> Exp -> Exp -> InternaliseM [SubExp]
internaliseStreamAcc [Char]
desc Exp
dest ((Exp, Exp) -> Maybe (Exp, Exp)
forall a. a -> Maybe a
Just (Exp
op, Exp
ne)) Exp
f Exp
bs
    handleAccs [TupLit [Exp
acc, Exp
i, Exp
v] SrcLoc
_] a
"acc_write" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      VName
acc' <- [VName] -> VName
forall a. [a] -> a
head ([VName] -> VName) -> InternaliseM [VName] -> InternaliseM VName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"acc" Exp
acc
      SubExp
i' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"acc_i" Exp
i
      [SubExp]
vs <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"acc_v" Exp
v
      (SubExp -> [SubExp])
-> InternaliseM SubExp -> InternaliseM [SubExp]
forall (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
$ [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
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 [Exp]
_ a
_ = Maybe ([Char] -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

    handleRest :: [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
handleRest [Exp
x] [Char]
"!" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ Exp -> [Char] -> InternaliseM [SubExp]
complementF Exp
x
    handleRest [Exp
x] [Char]
"opaque" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      (SubExp -> InternaliseM SubExp)
-> [SubExp] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
desc (ExpT SOACS -> InternaliseM SubExp)
-> (SubExp -> ExpT SOACS) -> SubExp -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS)
-> (SubExp -> BasicOp) -> SubExp -> ExpT SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
Opaque) ([SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"opaque_arg" Exp
x
    handleRest [E.TupLit [Exp
a, Exp
si, Exp
v] SrcLoc
_] [Char]
"scatter" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ Int -> Exp -> Exp -> Exp -> [Char] -> InternaliseM [SubExp]
scatterF Int
1 Exp
a Exp
si Exp
v
    handleRest [E.TupLit [Exp
a, Exp
si, Exp
v] SrcLoc
_] [Char]
"scatter_2d" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ Int -> Exp -> Exp -> Exp -> [Char] -> InternaliseM [SubExp]
scatterF Int
2 Exp
a Exp
si Exp
v
    handleRest [E.TupLit [Exp
a, Exp
si, Exp
v] SrcLoc
_] [Char]
"scatter_3d" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ Int -> Exp -> Exp -> Exp -> [Char] -> InternaliseM [SubExp]
scatterF Int
3 Exp
a Exp
si Exp
v
    handleRest [E.TupLit [Exp
n, Exp
m, Exp
arr] SrcLoc
_] [Char]
"unflatten" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      [VName]
arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"unflatten_arr" Exp
arr
      SubExp
n' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"n" Exp
n
      SubExp
m' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"m" Exp
m
      -- 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 lore (m :: * -> *).
HasScope lore m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
      SubExp
dim_ok <-
        [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"dim_ok"
          (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< CmpOp
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
CmpOp -> m (Exp (Lore m)) -> m (Exp (Lore m)) -> m (Exp (Lore m))
eCmpOp
            (PrimType -> CmpOp
I.CmpEq PrimType
I.int64)
            (BinOp
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> m (Exp (Lore m)) -> m (Exp (Lore m)) -> m (Exp (Lore m))
eBinOp (IntType -> Overflow -> BinOp
I.Mul IntType
Int64 Overflow
I.OverflowUndef) (SubExp -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => SubExp -> m (Exp (Lore m))
eSubExp SubExp
n') (SubExp -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => SubExp -> m (Exp (Lore m))
eSubExp SubExp
m'))
            (SubExp -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => SubExp -> m (Exp (Lore m))
eSubExp SubExp
old_dim)
      Certificates
dim_ok_cert <-
        [Char]
-> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certificates
assert
          [Char]
"dim_ok_cert"
          SubExp
dim_ok
          ErrorMsg SubExp
"new shape has different number of elements than old shape"
          SrcLoc
loc
      Certificates -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
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 lore (m :: * -> *).
HasScope lore m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
arr'
          [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
            BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
              ShapeChange SubExp -> VName -> BasicOp
I.Reshape (ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp
reshapeOuter [SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimNew SubExp
n', SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimNew SubExp
m'] Int
1 (Shape -> ShapeChange SubExp) -> Shape -> ShapeChange SubExp
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoUniqueness -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape TypeBase Shape NoUniqueness
arr_t) VName
arr'
    handleRest [Exp
arr] [Char]
"flatten" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      [VName]
arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"flatten_arr" Exp
arr
      [VName] -> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
arrs ((VName -> InternaliseM SubExp) -> InternaliseM [SubExp])
-> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
arr' -> do
        TypeBase Shape NoUniqueness
arr_t <- VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall lore (m :: * -> *).
HasScope lore m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
arr'
        let n :: SubExp
n = Int -> TypeBase Shape NoUniqueness -> SubExp
forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
0 TypeBase Shape NoUniqueness
arr_t
            m :: SubExp
m = Int -> TypeBase Shape NoUniqueness -> SubExp
forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
1 TypeBase Shape NoUniqueness
arr_t
        SubExp
k <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"flat_dim" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Mul IntType
Int64 Overflow
I.OverflowUndef) SubExp
n SubExp
m
        [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
            ShapeChange SubExp -> VName -> BasicOp
I.Reshape (ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp
reshapeOuter [SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimNew SubExp
k] Int
2 (Shape -> ShapeChange SubExp) -> Shape -> ShapeChange SubExp
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoUniqueness -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape TypeBase Shape NoUniqueness
arr_t) VName
arr'
    handleRest [TupLit [Exp
x, Exp
y] SrcLoc
_] [Char]
"concat" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      [VName]
xs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"concat_x" Exp
x
      [VName]
ys <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"concat_y" Exp
y
      SubExp
outer_size <- Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 ([TypeBase Shape NoUniqueness] -> SubExp)
-> InternaliseM [TypeBase Shape NoUniqueness]
-> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall lore (m :: * -> *).
HasScope lore m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
xs
      let sumdims :: SubExp -> SubExp -> m SubExp
sumdims SubExp
xsize SubExp
ysize =
            [Char] -> Exp (Lore m) -> m SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"conc_tmp" (Exp (Lore m) -> m SubExp) -> Exp (Lore m) -> m SubExp
forall a b. (a -> b) -> a -> b
$
              BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$
                BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Add IntType
I.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 :: * -> *}. MonadBinder 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 lore (m :: * -> *).
HasScope lore m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType) [[VName]
ys]

      let conc :: VName -> VName -> ExpT SOACS
conc VName
xarr VName
yarr =
            BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ Int -> VName -> [VName] -> SubExp -> BasicOp
I.Concat Int
0 VName
xarr [VName
yarr] SubExp
ressize
      [Char] -> [Exp (Lore InternaliseM)] -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> [Exp (Lore m)] -> m [SubExp]
letSubExps [Char]
desc ([Exp (Lore InternaliseM)] -> InternaliseM [SubExp])
-> [Exp (Lore 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 [Exp
offset, Exp
e] SrcLoc
_] [Char]
"rotate" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      SubExp
offset' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"rotation_offset" Exp
offset
      [Char]
-> Exp -> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
internaliseOperation [Char]
desc Exp
e ((VName -> InternaliseM BasicOp) -> InternaliseM [SubExp])
-> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
v -> do
        Int
r <- TypeBase Shape NoUniqueness -> Int
forall shape u. ArrayShape shape => TypeBase shape u -> Int
I.arrayRank (TypeBase Shape NoUniqueness -> Int)
-> InternaliseM (TypeBase Shape NoUniqueness) -> InternaliseM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall lore (m :: * -> *).
HasScope lore 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 [Exp
e] [Char]
"transpose" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char]
-> Exp -> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
internaliseOperation [Char]
desc Exp
e ((VName -> InternaliseM BasicOp) -> InternaliseM [SubExp])
-> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
v -> do
        Int
r <- TypeBase Shape NoUniqueness -> Int
forall shape u. ArrayShape shape => TypeBase shape u -> Int
I.arrayRank (TypeBase Shape NoUniqueness -> Int)
-> InternaliseM (TypeBase Shape NoUniqueness) -> InternaliseM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall lore (m :: * -> *).
HasScope lore 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 [Exp
x, Exp
y] SrcLoc
_] [Char]
"zip" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      (VName -> InternaliseM SubExp) -> [VName] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"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 lore. BasicOp -> ExpT lore
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
<$> [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars ([Char]
desc [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"_zip_x") Exp
x
                InternaliseM ([VName] -> [VName])
-> InternaliseM [VName] -> InternaliseM [VName]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars ([Char]
desc [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"_zip_y") Exp
y
            )
    handleRest [Exp
x] [Char]
"unzip" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ ([Char] -> Exp -> InternaliseM [SubExp])
-> Exp -> [Char] -> InternaliseM [SubExp]
forall a b c. (a -> b -> c) -> b -> a -> c
flip [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp Exp
x
    handleRest [Exp
x] [Char]
"trace" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ ([Char] -> Exp -> InternaliseM [SubExp])
-> Exp -> [Char] -> InternaliseM [SubExp]
forall a b c. (a -> b -> c) -> b -> a -> c
flip [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp Exp
x
    handleRest [Exp
x] [Char]
"break" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ ([Char] -> Exp -> InternaliseM [SubExp])
-> Exp -> [Char] -> InternaliseM [SubExp]
forall a b c. (a -> b -> c) -> b -> a -> c
flip [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp Exp
x
    handleRest [Exp]
_ [Char]
_ = Maybe ([Char] -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

    toSigned :: IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
int_to Exp
e [Char]
desc = do
      SubExp
e' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"trunc_arg" Exp
e
      case Exp -> PatternType
E.typeOf Exp
e of
        E.Scalar (E.Prim PrimType
E.Bool) ->
          [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
            SubExp
-> BodyT SOACS
-> BodyT SOACS
-> IfDec (BranchType SOACS)
-> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If
              SubExp
e'
              ([SubExp] -> BodyT SOACS
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
1])
              ([SubExp] -> BodyT SOACS
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
0])
              (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ [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)) ->
          [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (IntType -> IntType -> ConvOp
I.SExt IntType
int_from IntType
int_to) SubExp
e'
        E.Scalar (E.Prim (E.Unsigned IntType
int_from)) ->
          [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (IntType -> IntType -> ConvOp
I.ZExt IntType
int_from IntType
int_to) SubExp
e'
        E.Scalar (E.Prim (E.FloatType FloatType
float_from)) ->
          [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (FloatType -> IntType -> ConvOp
I.FPToSI FloatType
float_from IntType
int_to) SubExp
e'
        PatternType
_ -> [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise: non-numeric type in ToSigned"

    toUnsigned :: IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
int_to Exp
e [Char]
desc = do
      SubExp
e' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"trunc_arg" Exp
e
      case Exp -> PatternType
E.typeOf Exp
e of
        E.Scalar (E.Prim PrimType
E.Bool) ->
          [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
            SubExp
-> BodyT SOACS
-> BodyT SOACS
-> IfDec (BranchType SOACS)
-> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If
              SubExp
e'
              ([SubExp] -> BodyT SOACS
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
1])
              ([SubExp] -> BodyT SOACS
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
0])
              (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ [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)) ->
          [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (IntType -> IntType -> ConvOp
I.ZExt IntType
int_from IntType
int_to) SubExp
e'
        E.Scalar (E.Prim (E.Unsigned IntType
int_from)) ->
          [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (IntType -> IntType -> ConvOp
I.ZExt IntType
int_from IntType
int_to) SubExp
e'
        E.Scalar (E.Prim (E.FloatType FloatType
float_from)) ->
          [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (FloatType -> IntType -> ConvOp
I.FPToUI FloatType
float_from IntType
int_to) SubExp
e'
        PatternType
_ -> [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise.internaliseExp: non-numeric type in ToUnsigned"

    complementF :: Exp -> [Char] -> InternaliseM [SubExp]
complementF Exp
e [Char]
desc = do
      SubExp
e' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"complement_arg" Exp
e
      TypeBase Shape NoUniqueness
et <- SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType SubExp
e'
      case TypeBase Shape NoUniqueness
et of
        I.Prim (I.IntType IntType
t) ->
          [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (IntType -> UnOp
I.Complement IntType
t) SubExp
e'
        I.Prim PrimType
I.Bool ->
          [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
I.Not SubExp
e'
        TypeBase Shape NoUniqueness
_ ->
          [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise.internaliseExp: non-int/bool type in Complement"

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

      SubExp
si_w <- Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
I.arraysSize Int
0 ([TypeBase Shape NoUniqueness] -> SubExp)
-> InternaliseM [TypeBase Shape NoUniqueness]
-> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall lore (m :: * -> *).
HasScope lore 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 lore (m :: * -> *).
HasScope lore m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
svs

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

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

      [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 lore (m :: * -> *).
HasScope lore 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
_ -> [Char] -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"write_index") [TypeBase Shape NoUniqueness]
indexType
      [VName]
valueNames <- Int -> InternaliseM VName -> InternaliseM [VName]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM ([TypeBase Shape NoUniqueness] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase Shape NoUniqueness]
sv_ts) (InternaliseM VName -> InternaliseM [VName])
-> InternaliseM VName -> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$ [Char] -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"write_value"

      [TypeBase Shape NoUniqueness]
sa_ts <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall lore (m :: * -> *).
HasScope lore 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 lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param (TypeBase Shape NoUniqueness)] -> Scope SOACS
forall lore dec.
(LParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
scopeOfLParams [Param (TypeBase Shape NoUniqueness)]
bodyParams) (InternaliseM (BodyT SOACS) -> InternaliseM (BodyT SOACS))
-> (InternaliseM [SubExp] -> InternaliseM (BodyT SOACS))
-> InternaliseM [SubExp]
-> InternaliseM (BodyT SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InternaliseM [SubExp] -> InternaliseM (BodyT SOACS)
forall (m :: * -> *).
MonadBinder m =>
m [SubExp] -> m (Body (Lore m))
buildBody_ (InternaliseM [SubExp] -> InternaliseM (BodyT SOACS))
-> InternaliseM [SubExp] -> 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 ->
          [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"write_res" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
name
        ErrorMsg SubExp
-> SrcLoc
-> [TypeBase Shape NoUniqueness]
-> [SubExp]
-> InternaliseM [SubExp]
ensureResultShape
          ErrorMsg SubExp
"scatter value has wrong size"
          SrcLoc
loc
          [TypeBase Shape NoUniqueness]
bodyTypes
          [SubExp]
results

      let lam :: Lambda
lam =
            Lambda :: forall lore.
[LParam lore]
-> BodyT lore -> [TypeBase Shape NoUniqueness] -> LambdaT lore
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
      [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Op SOACS -> ExpT SOACS
forall lore. Op lore -> ExpT lore
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> Lambda -> [VName] -> [(Shape, Int, VName)] -> SOAC SOACS
forall lore.
SubExp
-> Lambda lore -> [VName] -> [(Shape, Int, VName)] -> SOAC lore
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

funcall ::
  String ->
  QualName VName ->
  [SubExp] ->
  SrcLoc ->
  InternaliseM ([SubExp], [I.ExtType])
funcall :: [Char]
-> QualName VName
-> [SubExp]
-> SrcLoc
-> InternaliseM ([SubExp], [ExtType])
funcall [Char]
desc (QualName [VName]
_ VName
fname) [SubExp]
args SrcLoc
loc = do
  ([VName]
shapes, [DeclType]
value_paramts, [Param DeclType]
fun_params, [(SubExp, TypeBase Shape NoUniqueness)] -> Maybe [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]
forall (m :: * -> *).
(HasScope SOACS m, Monad m) =>
[VName] -> [FParam] -> [TypeBase Shape NoUniqueness] -> m [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 ->
      [Char] -> InternaliseM ([SubExp], [ExtType])
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM ([SubExp], [ExtType]))
-> [Char] -> InternaliseM ([SubExp], [ExtType])
forall a b. (a -> b) -> a -> b
$
        [[Char]] -> [Char]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
          [ [Char]
"Cannot apply ",
            VName -> [Char]
forall a. Pretty a => a -> [Char]
pretty VName
fname,
            [Char]
" to ",
            Int -> [Char]
forall a. Show a => a -> [Char]
show ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
args'),
            [Char]
" arguments\n ",
            [SubExp] -> [Char]
forall a. Pretty a => a -> [Char]
pretty [SubExp]
args',
            [Char]
"\nof types\n ",
            [TypeBase Shape NoUniqueness] -> [Char]
forall a. Pretty a => a -> [Char]
pretty [TypeBase Shape NoUniqueness]
argts',
            [Char]
"\nFunction has ",
            Int -> [Char]
forall a. Show a => a -> [Char]
show ([Param DeclType] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Param DeclType]
fun_params),
            [Char]
" parameters\n ",
            [Param DeclType] -> [Char]
forall a. Pretty a => a -> [Char]
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. MonadBinder m => Attrs -> m a -> m a
attributing Attrs
attrs (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
          [Char] -> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Lore InternaliseM) -> InternaliseM [SubExp])
-> Exp (Lore InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
            Name
-> [(SubExp, Diet)]
-> [RetType SOACS]
-> (Safety, SrcLoc, [SrcLoc])
-> ExpT SOACS
forall lore.
Name
-> [(SubExp, Diet)]
-> [RetType lore]
-> (Safety, SrcLoc, [SrcLoc])
-> ExpT lore
I.Apply (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 PatternType
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
$ PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
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
$ (ExtSize -> SubExp -> Map VName SubExp)
-> [ExtSize] -> [SubExp] -> [Map VName SubExp]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith ExtSize -> SubExp -> Map VName SubExp
combine' (DeclExtType -> [ExtSize]
forall u. TypeBase (ShapeBase ExtSize) u -> [ExtSize]
arrayExtDims DeclExtType
t1) (TypeBase Shape NoUniqueness -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
arrayDims TypeBase Shape NoUniqueness
t2)
      combine' :: ExtSize -> 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' ExtSize
_ SubExp
_ = Map VName SubExp
forall a. Monoid a => a
mempty

  [(VName, SubExp)]
-> ((VName, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Map VName SubExp -> [(VName, SubExp)]
forall k a. Map k a -> [(k, a)]
M.toList (Map VName SubExp -> [(VName, SubExp)])
-> Map VName SubExp -> [(VName, SubExp)]
forall a b. (a -> b) -> a -> b
$ [Map VName SubExp] -> Map VName SubExp
forall a. Monoid a => [a] -> a
mconcat ([Map VName SubExp] -> Map VName SubExp)
-> [Map VName SubExp] -> Map VName SubExp
forall a b. (a -> b) -> a -> b
$ (DeclExtType -> 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 (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [VName
v] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se

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

-- Implement partitioning using maps, scans and writes.
partitionWithSOACS :: Int -> I.Lambda -> [I.VName] -> InternaliseM ([I.SubExp], [I.SubExp])
partitionWithSOACS :: Int -> Lambda -> [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 lore (m :: * -> *).
HasScope lore m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
  let w :: SubExp
w = Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [TypeBase Shape NoUniqueness]
arr_ts
  [VName]
classes_and_increments <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM [VName]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [VName]
letTupExp [Char]
"increments" (Exp (Lore InternaliseM) -> InternaliseM [VName])
-> Exp (Lore InternaliseM) -> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$ Op SOACS -> ExpT SOACS
forall lore. Op lore -> ExpT lore
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
forall lore. SubExp -> [VName] -> ScremaForm lore -> SOAC lore
I.Screma SubExp
w [VName]
arrs (Lambda -> ScremaForm SOACS
forall lore. Lambda lore -> ScremaForm lore
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]
_ -> [Char] -> InternaliseM (VName, [VName])
forall a. HasCallStack => [Char] -> a
error [Char]
"partitionWithSOACS"

  [Param (TypeBase Shape NoUniqueness)]
add_lam_x_params <-
    Int
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
-> InternaliseM [Param (TypeBase Shape NoUniqueness)]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
k (InternaliseM (Param (TypeBase Shape NoUniqueness))
 -> InternaliseM [Param (TypeBase Shape NoUniqueness)])
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
-> InternaliseM [Param (TypeBase Shape NoUniqueness)]
forall a b. (a -> b) -> a -> b
$ 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
<$> [Char] -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"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
<$> [Char] -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"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 <- Binder SOACS (BodyT SOACS) -> InternaliseM (BodyT SOACS)
forall lore (m :: * -> *) somelore.
(Bindable lore, MonadFreshNames m, HasScope somelore m,
 SameScope somelore lore) =>
Binder lore (Body lore) -> m (Body lore)
runBodyBinder (Binder SOACS (BodyT SOACS) -> InternaliseM (BodyT SOACS))
-> Binder SOACS (BodyT SOACS) -> InternaliseM (BodyT SOACS)
forall a b. (a -> b) -> a -> b
$
    Scope SOACS
-> Binder SOACS (BodyT SOACS) -> Binder SOACS (BodyT SOACS)
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param (TypeBase Shape NoUniqueness)] -> Scope SOACS
forall lore dec.
(LParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
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) (Binder SOACS (BodyT SOACS) -> Binder SOACS (BodyT SOACS))
-> Binder SOACS (BodyT SOACS) -> Binder SOACS (BodyT SOACS)
forall a b. (a -> b) -> a -> b
$
      ([SubExp] -> BodyT SOACS)
-> BinderT SOACS (State VNameSource) [SubExp]
-> Binder SOACS (BodyT SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [SubExp] -> BodyT SOACS
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody (BinderT SOACS (State VNameSource) [SubExp]
 -> Binder SOACS (BodyT SOACS))
-> BinderT SOACS (State VNameSource) [SubExp]
-> Binder 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))
    -> BinderT SOACS (State VNameSource) SubExp)
-> BinderT 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))
  -> BinderT SOACS (State VNameSource) SubExp)
 -> BinderT SOACS (State VNameSource) [SubExp])
-> ((Param (TypeBase Shape NoUniqueness),
     Param (TypeBase Shape NoUniqueness))
    -> BinderT SOACS (State VNameSource) SubExp)
-> BinderT SOACS (State VNameSource) [SubExp]
forall a b. (a -> b) -> a -> b
$ \(Param (TypeBase Shape NoUniqueness)
x, Param (TypeBase Shape NoUniqueness)
y) ->
          [Char]
-> Exp (Lore (BinderT SOACS (State VNameSource)))
-> BinderT SOACS (State VNameSource) SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"z" (Exp (Lore (BinderT SOACS (State VNameSource)))
 -> BinderT SOACS (State VNameSource) SubExp)
-> Exp (Lore (BinderT SOACS (State VNameSource)))
-> BinderT SOACS (State VNameSource) SubExp
forall a b. (a -> b) -> a -> b
$
            BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
              BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp
                (IntType -> Overflow -> BinOp
I.Add IntType
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 lore.
[LParam lore]
-> BodyT lore -> [TypeBase Shape NoUniqueness] -> LambdaT lore
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 lore (m :: * -> *).
(Bindable lore, MonadFreshNames m) =>
[Scan lore] -> m (ScremaForm lore)
I.scanSOAC [Lambda -> [SubExp] -> Scan SOACS
forall lore. Lambda lore -> [SubExp] -> Scan lore
I.Scan Lambda
add_lam [SubExp]
nes]
  [VName]
all_offsets <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM [VName]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [VName]
letTupExp [Char]
"offsets" (Exp (Lore InternaliseM) -> InternaliseM [VName])
-> Exp (Lore InternaliseM) -> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$ Op SOACS -> ExpT SOACS
forall lore. Op lore -> ExpT lore
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
forall lore. SubExp -> [VName] -> ScremaForm lore -> SOAC lore
I.Screma SubExp
w [VName]
increments ScremaForm SOACS
scan

  -- We have the offsets for each of the partitions, but we also need
  -- the total sizes, which are the last elements in the offests.  We
  -- just have to be careful in case the array is empty.
  SubExp
last_index <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"last_index" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
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 <- Binder SOACS (BodyT SOACS) -> InternaliseM (BodyT SOACS)
forall lore (m :: * -> *) somelore.
(Bindable lore, MonadFreshNames m, HasScope somelore m,
 SameScope somelore lore) =>
Binder lore (Body lore) -> m (Body lore)
runBodyBinder (Binder SOACS (BodyT SOACS) -> InternaliseM (BodyT SOACS))
-> Binder SOACS (BodyT SOACS) -> InternaliseM (BodyT SOACS)
forall a b. (a -> b) -> a -> b
$
    ([SubExp] -> BodyT SOACS)
-> BinderT SOACS (State VNameSource) [SubExp]
-> Binder SOACS (BodyT SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [SubExp] -> BodyT SOACS
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody (BinderT SOACS (State VNameSource) [SubExp]
 -> Binder SOACS (BodyT SOACS))
-> BinderT SOACS (State VNameSource) [SubExp]
-> Binder SOACS (BodyT SOACS)
forall a b. (a -> b) -> a -> b
$
      [VName]
-> (VName -> BinderT SOACS (State VNameSource) SubExp)
-> BinderT 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 -> BinderT SOACS (State VNameSource) SubExp)
 -> BinderT SOACS (State VNameSource) [SubExp])
-> (VName -> BinderT SOACS (State VNameSource) SubExp)
-> BinderT SOACS (State VNameSource) [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
offset_array ->
        [Char]
-> Exp (Lore (BinderT SOACS (State VNameSource)))
-> BinderT SOACS (State VNameSource) SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"last_offset" (Exp (Lore (BinderT SOACS (State VNameSource)))
 -> BinderT SOACS (State VNameSource) SubExp)
-> Exp (Lore (BinderT SOACS (State VNameSource)))
-> BinderT SOACS (State VNameSource) SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ VName -> [DimIndex SubExp] -> BasicOp
I.Index VName
offset_array [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
I.DimFix SubExp
last_index]
  let empty_body :: BodyT SOACS
empty_body = [SubExp] -> BodyT SOACS
forall lore. Bindable lore => [SubExp] -> Body lore
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 <- [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"is_empty" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
CmpEq PrimType
int64) SubExp
w (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)
  [VName]
sizes <-
    [Char] -> Exp (Lore InternaliseM) -> InternaliseM [VName]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [VName]
letTupExp [Char]
"partition_size" (Exp (Lore InternaliseM) -> InternaliseM [VName])
-> Exp (Lore InternaliseM) -> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$
      SubExp
-> BodyT SOACS
-> BodyT SOACS
-> IfDec (BranchType SOACS)
-> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
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 ->
    [Char] -> Exp (Lore InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m VName
letExp [Char]
"partition_dest" (Exp (Lore InternaliseM) -> InternaliseM VName)
-> Exp (Lore InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ PrimType -> [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
<$> [Char] -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"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
<$> [Char] -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"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
<$> [Char] -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"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 (Lore InternaliseM))
forall (m :: * -> *) a.
MonadBinder m =>
m a -> m (a, Stms (Lore m))
collectStms (InternaliseM SubExp
 -> InternaliseM (SubExp, Stms (Lore InternaliseM)))
-> InternaliseM SubExp
-> InternaliseM (SubExp, Stms (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$
        [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 lore.
[LParam lore]
-> BodyT lore -> [TypeBase Shape NoUniqueness] -> LambdaT lore
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 -> [SubExp] -> BodyT SOACS
forall lore. Bindable lore => Stms lore -> [SubExp] -> Body lore
mkBody Stms SOACS
offset_stms ([SubExp] -> BodyT SOACS) -> [SubExp] -> BodyT SOACS
forall a b. (a -> b) -> a -> b
$
              Int -> SubExp -> [SubExp]
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
offset
                [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ (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
I.paramName) [Param (TypeBase Shape NoUniqueness)]
value_params
        }
  [VName]
results <-
    [Char] -> Exp (Lore InternaliseM) -> InternaliseM [VName]
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m [VName]
letTupExp [Char]
"partition_res" (Exp (Lore InternaliseM) -> InternaliseM [VName])
-> Exp (Lore InternaliseM) -> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$
      Op SOACS -> ExpT SOACS
forall lore. Op lore -> ExpT lore
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
        SubExp -> Lambda -> [VName] -> [(Shape, Int, VName)] -> SOAC SOACS
forall lore.
SubExp
-> Lambda lore -> [VName] -> [(Shape, Int, VName)] -> SOAC lore
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' <-
    [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"partition_sizes" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
        [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 <-
        [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"is_this_one" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
CmpEq PrimType
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 <-
        [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"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 (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> SubExp -> [SubExp] -> m (Exp (Lore 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)
      [Char] -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
[Char] -> Exp (Lore m) -> m SubExp
letSubExp [Char]
"total_res" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
        SubExp
-> BodyT SOACS
-> BodyT SOACS
-> IfDec (BranchType SOACS)
-> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If
          SubExp
is_this_one
          ([SubExp] -> BodyT SOACS
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody [SubExp
this_one])
          ([SubExp] -> BodyT SOACS
forall lore. Bindable lore => [SubExp] -> Body lore
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 [[Char] -> ErrorMsgPart SubExp
forall a. [Char] -> ErrorMsgPart a
ErrorString ([Char] -> ErrorMsgPart SubExp) -> [Char] -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ QualName VName -> [Char]
forall a. Pretty a => a -> [Char]
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) =
      ([Char] -> ErrorMsgPart SubExp
forall a. [Char] -> ErrorMsgPart a
ErrorString (a -> [Char]
forall a. Pretty a => a -> [Char]
pretty a
k [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
": ") 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
$ SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt64 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
$ [Char] -> ErrorMsgPart SubExp
forall a. [Char] -> ErrorMsgPart a
ErrorString ([Char] -> ErrorMsgPart SubExp) -> [Char] -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ Int -> [Char]
forall a. Pretty a => a -> [Char]
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 :: forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg = [ErrorMsgPart a] -> ErrorMsg a
forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg ([ErrorMsgPart a] -> ErrorMsg a)
-> ([ErrorMsgPart a] -> [ErrorMsgPart a])
-> [ErrorMsgPart a]
-> ErrorMsg a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ErrorMsgPart a] -> [ErrorMsgPart a]
forall {a}. [ErrorMsgPart a] -> [ErrorMsgPart a]
compact
  where
    compact :: [ErrorMsgPart a] -> [ErrorMsgPart a]
compact [] = []
    compact (ErrorString [Char]
x : ErrorString [Char]
y : [ErrorMsgPart a]
parts) =
      [ErrorMsgPart a] -> [ErrorMsgPart a]
compact ([Char] -> ErrorMsgPart a
forall a. [Char] -> ErrorMsgPart a
ErrorString ([Char]
x [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
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