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

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

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

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

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

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

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

    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 (f :: * -> *) a. Applicative f => a -> f a
pure (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."]

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

    let all_params :: [Param DeclType]
all_params = [Param DeclType]
[FParam SOACS]
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 SOACS]]
params'

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

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

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

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

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

    OpaqueTypes -> InternaliseM ()
addOpaques OpaqueTypes
opaques

    (Body SOACS
entry_body, [TypeBase ExtShape Uniqueness]
ctx_ts) <- InternaliseM (Result, [TypeBase ExtShape Uniqueness])
-> InternaliseM
     (Body (Rep InternaliseM), [TypeBase ExtShape Uniqueness])
forall (m :: * -> *) a.
MonadBuilder m =>
m (Result, a) -> m (Body (Rep m), a)
buildBody (InternaliseM (Result, [TypeBase ExtShape Uniqueness])
 -> InternaliseM
      (Body (Rep InternaliseM), [TypeBase ExtShape Uniqueness]))
-> InternaliseM (Result, [TypeBase ExtShape Uniqueness])
-> InternaliseM
     (Body (Rep InternaliseM), [TypeBase ExtShape Uniqueness])
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 (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
ses
        Maybe [SubExp]
Nothing ->
          ([SubExp], [ExtType]) -> [SubExp]
forall a b. (a, b) -> a
fst (([SubExp], [ExtType]) -> [SubExp])
-> InternaliseM ([SubExp], [ExtType]) -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String
-> QualName VName
-> [SubExp]
-> SrcLoc
-> InternaliseM ([SubExp], [ExtType])
funcall String
"entry_result" (VName -> QualName VName
forall v. v -> QualName v
E.qualName VName
ofname) [SubExp]
args SrcLoc
loc
      [SubExp]
ctx <-
        [TypeBase ExtShape Uniqueness] -> [[SubExp]] -> [SubExp]
forall u a. [TypeBase ExtShape u] -> [[a]] -> [a]
extractShapeContext ([TypeBase ExtShape Uniqueness] -> [TypeBase ExtShape Uniqueness]
forall {u}. [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts ([TypeBase ExtShape Uniqueness] -> [TypeBase ExtShape Uniqueness])
-> [TypeBase ExtShape Uniqueness] -> [TypeBase ExtShape Uniqueness]
forall a b. (a -> b) -> a -> b
$ [[TypeBase ExtShape Uniqueness]] -> [TypeBase ExtShape Uniqueness]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[TypeBase ExtShape Uniqueness]]
entry_rettype)
          ([[SubExp]] -> [SubExp])
-> InternaliseM [[SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> InternaliseM [SubExp])
-> [SubExp] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((TypeBase Shape NoUniqueness -> [SubExp])
-> InternaliseM (TypeBase Shape NoUniqueness)
-> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase Shape NoUniqueness -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
I.arrayDims (InternaliseM (TypeBase Shape NoUniqueness)
 -> InternaliseM [SubExp])
-> (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> SubExp
-> InternaliseM [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType) [SubExp]
vals
      (Result, [TypeBase ExtShape Uniqueness])
-> InternaliseM (Result, [TypeBase ExtShape Uniqueness])
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([SubExp] -> Result
subExpsRes ([SubExp] -> Result) -> [SubExp] -> Result
forall a b. (a -> b) -> a -> b
$ [SubExp]
ctx [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
vals, (SubExp -> TypeBase ExtShape Uniqueness)
-> [SubExp] -> [TypeBase ExtShape Uniqueness]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase ExtShape Uniqueness
-> SubExp -> TypeBase ExtShape Uniqueness
forall a b. a -> b -> a
const (PrimType -> TypeBase ExtShape Uniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)) [SubExp]
ctx)

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

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

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

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

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

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

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

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

  let one :: SubExp
one = IntType -> Integer -> SubExp
intConst IntType
it Integer
1
      negone :: SubExp
negone = IntType -> Integer -> SubExp
intConst IntType
it (-Integer
1)
      default_step :: SubExp
default_step = case Inclusiveness 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 <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"subtracted_step" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$
            BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
it Overflow
I.OverflowWrap) SubExp
second' SubExp
start'
      SubExp
step_zero <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"step_zero" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (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 (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
subtracted_step, SubExp
step_zero)
    Maybe SubExp
Nothing ->
      (SubExp, SubExp) -> InternaliseM (SubExp, SubExp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
default_step, Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
False)

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

  SubExp
bounds_invalid_downwards <-
    String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"bounds_invalid_downwards" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$
        CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp CmpOp
le_op SubExp
start' SubExp
end'
  SubExp
bounds_invalid_upwards <-
    String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"bounds_invalid_upwards" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$
        CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp 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 <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"step_wrong_dir" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
one
      SubExp
distance <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"distance" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$
            BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
it Overflow
I.OverflowWrap) SubExp
start' SubExp
end'
      SubExp
distance_i64 <- IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
distance
      (SubExp, SubExp, SubExp) -> InternaliseM (SubExp, SubExp, SubExp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
distance_i64, SubExp
step_wrong_dir, SubExp
bounds_invalid_downwards)
    UpToExclusive {} -> do
      SubExp
step_wrong_dir <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"step_wrong_dir" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
negone
      SubExp
distance <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"distance" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
it Overflow
I.OverflowWrap) SubExp
end' SubExp
start'
      SubExp
distance_i64 <- IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
distance
      (SubExp, SubExp, SubExp) -> InternaliseM (SubExp, SubExp, SubExp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
distance_i64, SubExp
step_wrong_dir, SubExp
bounds_invalid_upwards)
    ToInclusive {} -> do
      SubExp
downwards <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"downwards" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
negone
      SubExp
distance_downwards_exclusive <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"distance_downwards_exclusive" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$
            BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
it Overflow
I.OverflowWrap) SubExp
start' SubExp
end'
      SubExp
distance_upwards_exclusive <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"distance_upwards_exclusive" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$
            BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
it Overflow
I.OverflowWrap) SubExp
end' SubExp
start'

      SubExp
bounds_invalid <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"bounds_invalid"
          (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
            (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
downwards)
            ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
bounds_invalid_downwards])
            ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
bounds_invalid_upwards])
      SubExp
distance_exclusive <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"distance_exclusive"
          (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
            (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
downwards)
            ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
distance_downwards_exclusive])
            ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
distance_upwards_exclusive])
      SubExp
distance_exclusive_i64 <- IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
distance_exclusive
      SubExp
distance <-
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"distance" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp 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 (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
distance, Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
False, SubExp
bounds_invalid)

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

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

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

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

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

      -- Some functions are magical (overloaded) and we handle that here.
      case () of
        -- Overloaded functions never take array arguments (except
        -- equality, but those cannot be existential), so we can safely
        -- ignore the existential dimensions.
        ()
          | Just String -> InternaliseM [SubExp]
internalise <- QualName VName
-> [Exp] -> SrcLoc -> Maybe (String -> 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 ->
              String -> InternaliseM [SubExp]
internalise String
desc
          | VName -> Int
baseTag (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag,
            Just (PrimType
rettype, [PrimType]
_) <- Name
-> Map Name (PrimType, [PrimType]) -> Maybe (PrimType, [PrimType])
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
fname Map Name (PrimType, [PrimType])
I.builtInFunctions -> do
              let tag :: [a] -> [(a, Diet)]
tag [a]
ses = [(a
se, Diet
I.Observe) | a
se <- [a]
ses]
              [[SubExp]]
args' <- [[SubExp]] -> [[SubExp]]
forall a. [a] -> [a]
reverse ([[SubExp]] -> [[SubExp]])
-> InternaliseM [[SubExp]] -> InternaliseM [[SubExp]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((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 (String -> (Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg String
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'
              String -> Exp SOACS -> InternaliseM [SubExp]
letValExp' String
desc (Exp SOACS -> InternaliseM [SubExp])
-> Exp SOACS -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Name
-> [(SubExp, Diet)]
-> [RetType SOACS]
-> (Safety, SrcLoc, [SrcLoc])
-> Exp SOACS
forall rep.
Name
-> [(SubExp, Diet)]
-> [RetType rep]
-> (Safety, SrcLoc, [SrcLoc])
-> Exp rep
I.Apply Name
fname [(SubExp, Diet)]
args'' [PrimType -> TypeBase ExtShape Uniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
rettype] (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 (String -> (Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg String
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
<$> String
-> QualName VName
-> [SubExp]
-> SrcLoc
-> InternaliseM ([SubExp], [ExtType])
funcall String
desc QualName VName
qfname [SubExp]
args' SrcLoc
loc
internaliseAppExp String
desc AppRes
_ (E.LetPat [SizeBinder VName]
sizes PatBase Info VName
pat Exp
e Exp
body SrcLoc
_) =
  String
-> [SizeBinder VName]
-> PatBase Info VName
-> Exp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
forall a.
String
-> [SizeBinder VName]
-> PatBase Info VName
-> Exp
-> InternaliseM a
-> InternaliseM a
internalisePat String
desc [SizeBinder VName]
sizes PatBase Info VName
pat Exp
e (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String -> Exp -> InternaliseM [SubExp]
internaliseExp String
desc Exp
body
internaliseAppExp String
_ AppRes
_ (E.LetFun VName
ofname ([TypeParamBase VName], [PatBase Info VName],
 Maybe (TypeExp VName), Info StructRetType, Exp)
_ Exp
_ SrcLoc
_) =
  String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error (String -> InternaliseM [SubExp])
-> String -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String
"Unexpected LetFun " String -> String -> String
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Pretty a => a -> String
pretty VName
ofname
internaliseAppExp String
desc AppRes
_ (E.DoLoop [VName]
sparams PatBase Info VName
mergepat Exp
mergeexp LoopFormBase Info VName
form Exp
loopbody SrcLoc
loc) = do
  [SubExp]
ses <- String -> Exp -> InternaliseM [SubExp]
internaliseExp String
"loop_init" Exp
mergeexp
  ((Body SOACS
loopbody', (LoopForm SOACS
form', [Param DeclType]
shapepat, [Param DeclType]
mergepat', [SubExp]
mergeinit')), Stms SOACS
initstms) <-
    InternaliseM
  (Body SOACS,
   (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     ((Body SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])),
      Stms (Rep InternaliseM))
forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms (InternaliseM
   (Body SOACS,
    (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
 -> InternaliseM
      ((Body SOACS,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])),
       Stms (Rep InternaliseM)))
-> InternaliseM
     (Body SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     ((Body SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])),
      Stms (Rep InternaliseM))
forall a b. (a -> b) -> a -> b
$ [SubExp]
-> LoopFormBase Info VName
-> InternaliseM
     (Body SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
handleForm [SubExp]
ses LoopFormBase Info VName
form

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

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

  -- Ensure that the initial loop values match the shapes of the loop
  -- parameters.  XXX: Ideally they should already match (by the
  -- source language type rules), but some of our transformations
  -- (esp. defunctionalisation) strips out some size information.  For
  -- a type-correct source program, these reshapes should simplify
  -- away.
  let args :: [SubExp]
args = [SubExp]
ctxinit [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
mergeinit'
  [SubExp]
args' <-
    ErrorMsg SubExp
-> SrcLoc
-> [VName]
-> [TypeBase Shape NoUniqueness]
-> [SubExp]
-> InternaliseM [SubExp]
forall u.
Typed (TypeBase Shape u) =>
ErrorMsg SubExp
-> SrcLoc
-> [VName]
-> [TypeBase Shape u]
-> [SubExp]
-> InternaliseM [SubExp]
ensureArgShapes
      ErrorMsg SubExp
"initial loop values have right shape"
      SrcLoc
loc
      ((Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
shapepat)
      ((Param DeclType -> TypeBase Shape NoUniqueness)
-> [Param DeclType] -> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> TypeBase Shape NoUniqueness
forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
paramType ([Param DeclType] -> [TypeBase Shape NoUniqueness])
-> [Param DeclType] -> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> a -> b
$ [Param DeclType]
shapepat [Param DeclType] -> [Param DeclType] -> [Param DeclType]
forall a. [a] -> [a] -> [a]
++ [Param DeclType]
mergepat')
      [SubExp]
args

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

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

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

    forLoop :: [Param DeclType]
-> [Param DeclType]
-> [SubExp]
-> LoopForm SOACS
-> InternaliseM
     (Body SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forLoop [Param DeclType]
mergepat' [Param DeclType]
shapepat [SubExp]
mergeinit LoopForm SOACS
form' =
      InternaliseM
  (Result,
   (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Body SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a. InternaliseM (Result, a) -> InternaliseM (Body SOACS, a)
bodyFromStms (InternaliseM
   (Result,
    (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
 -> InternaliseM
      (Body SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> (InternaliseM
      (Result,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
    -> InternaliseM
         (Result,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Body SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LoopForm SOACS
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf LoopForm SOACS
form' (InternaliseM
   (Result,
    (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
 -> InternaliseM
      (Body SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Body SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ do
        [SubExp]
ses <- String -> Exp -> InternaliseM [SubExp]
internaliseExp String
"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 SOACS]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM [SubExp]
argShapes ((Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
shapepat) [Param DeclType]
[FParam SOACS]
mergepat' [TypeBase Shape NoUniqueness]
sets
        (Result,
 (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall (f :: * -> *) a. Applicative f => a -> f a
pure
          ( [SubExp] -> Result
subExpsRes ([SubExp] -> Result) -> [SubExp] -> Result
forall a b. (a -> b) -> a -> b
$ [SubExp]
shapeargs [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
ses,
            ( LoopForm SOACS
form',
              [Param DeclType]
shapepat,
              [Param DeclType]
mergepat',
              [SubExp]
mergeinit
            )
          )

    handleForm :: [SubExp]
-> LoopFormBase Info VName
-> InternaliseM
     (Body SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
handleForm [SubExp]
mergeinit (E.ForIn PatBase Info VName
x Exp
arr) = do
      [VName]
arr' <- String -> Exp -> InternaliseM [VName]
internaliseExpToVars String
"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 rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arr'
      let w :: SubExp
w = Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [TypeBase Shape NoUniqueness]
arr_ts

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

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

      [TypeBase Shape NoUniqueness]
ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
mergeinit
      [TypeParamBase VName]
-> PatBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([FParam SOACS]
    -> [FParam SOACS]
    -> InternaliseM
         (Body SOACS,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a.
[TypeParamBase VName]
-> PatBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([FParam SOACS] -> [FParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' PatBase Info VName
mergepat [TypeBase Shape NoUniqueness]
ts (([FParam SOACS]
  -> [FParam SOACS]
  -> InternaliseM
       (Body SOACS,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
 -> InternaliseM
      (Body SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> ([FParam SOACS]
    -> [FParam SOACS]
    -> InternaliseM
         (Body SOACS,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$
        \[FParam SOACS]
shapepat [FParam SOACS]
mergepat' ->
          [Param DeclType]
-> [Param DeclType]
-> [SubExp]
-> LoopForm SOACS
-> InternaliseM
     (Body SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forLoop [Param DeclType]
[FParam SOACS]
mergepat' [Param DeclType]
[FParam SOACS]
shapepat [SubExp]
mergeinit (LoopForm SOACS
 -> InternaliseM
      (Body SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> LoopForm SOACS
-> InternaliseM
     (Body SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$
            VName
-> IntType -> SubExp -> [(LParam SOACS, VName)] -> LoopForm SOACS
forall rep.
VName -> IntType -> SubExp -> [(LParam rep, VName)] -> LoopForm rep
I.ForLoop (IdentBase Info VName -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
E.identName IdentBase Info VName
i) IntType
it SubExp
num_iterations' []
    handleForm [SubExp]
mergeinit (E.While 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]
-> PatBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([FParam SOACS]
    -> [FParam SOACS]
    -> InternaliseM
         (Body SOACS,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a.
[TypeParamBase VName]
-> PatBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([FParam SOACS] -> [FParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' PatBase Info VName
mergepat [TypeBase Shape NoUniqueness]
ts (([FParam SOACS]
  -> [FParam SOACS]
  -> InternaliseM
       (Body SOACS,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
 -> InternaliseM
      (Body SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> ([FParam SOACS]
    -> [FParam SOACS]
    -> InternaliseM
         (Body SOACS,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ \[FParam SOACS]
shapepat [FParam SOACS]
mergepat' -> do
        [TypeBase Shape NoUniqueness]
mergeinit_ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
mergeinit
        -- We need to insert 'cond' twice - once for the initial
        -- condition (do we enter the loop at all?), and once with the
        -- result values of the loop (do we continue into the next
        -- iteration?).  This is safe, as the type rules for the
        -- external language guarantees that 'cond' does not consume
        -- anything.
        [SubExp]
shapeinit <- [VName]
-> [FParam SOACS]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM [SubExp]
argShapes ((Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
[FParam SOACS]
shapepat) [FParam SOACS]
mergepat' [TypeBase Shape NoUniqueness]
mergeinit_ts

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

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

        InternaliseM
  (Result,
   (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Body SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a. InternaliseM (Result, a) -> InternaliseM (Body SOACS, a)
bodyFromStms (InternaliseM
   (Result,
    (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
 -> InternaliseM
      (Body SOACS,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Body SOACS,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ do
          [SubExp]
ses <- String -> Exp -> InternaliseM [SubExp]
internaliseExp String
"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 <- String -> DeclType -> InternaliseM (Param DeclType)
forall (m :: * -> *) dec.
MonadFreshNames m =>
String -> dec -> m (Param dec)
newParam String
"loop_while" (DeclType -> InternaliseM (Param DeclType))
-> DeclType -> InternaliseM (Param DeclType)
forall a b. (a -> b) -> a -> b
$ PrimType -> DeclType
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.Bool
          [SubExp]
shapeargs <- [VName]
-> [FParam SOACS]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM [SubExp]
argShapes ((Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
[FParam SOACS]
shapepat) [FParam SOACS]
mergepat' [TypeBase Shape NoUniqueness]
sets

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

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

internaliseExp :: String -> E.Exp -> InternaliseM [I.SubExp]
internaliseExp :: String -> Exp -> InternaliseM [SubExp]
internaliseExp String
desc (E.Parens Exp
e SrcLoc
_) =
  String -> Exp -> InternaliseM [SubExp]
internaliseExp String
desc Exp
e
internaliseExp String
desc (E.Hole (Info PatType
t) SrcLoc
loc) = do
  let msg :: String
msg = Doc -> String
forall a. Pretty a => a -> String
pretty (Doc -> String) -> Doc -> String
forall a b. (a -> b) -> a -> b
$ Doc
"Reached hole of type: " Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
align (PatType -> Doc
forall a. Pretty a => a -> Doc
ppr PatType
t)
      ts :: [TypeBase ExtShape Uniqueness]
ts = TypeBase Size () -> [TypeBase ExtShape Uniqueness]
internaliseType (PatType -> TypeBase Size ()
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatType
t)
  Certs
c <- String -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert String
"hole_c" (Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
False) ([ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [String -> ErrorMsgPart SubExp
forall a. String -> ErrorMsgPart a
ErrorString String
msg]) SrcLoc
loc
  case (TypeBase ExtShape Uniqueness -> Maybe DeclType)
-> [TypeBase ExtShape Uniqueness] -> Maybe [DeclType]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TypeBase ExtShape Uniqueness -> Maybe DeclType
forall u. TypeBase ExtShape u -> Maybe (TypeBase Shape u)
hasStaticShape [TypeBase ExtShape Uniqueness]
ts of
    Maybe [DeclType]
Nothing ->
      String -> InternaliseM [SubExp]
forall a. HasCallStack => String -> a
error (String -> InternaliseM [SubExp])
-> String -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ String
"Hole at " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> SrcLoc -> String
forall a. Located a => a -> String
locStr SrcLoc
loc String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" has existential type:\n" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> [TypeBase ExtShape Uniqueness] -> String
forall a. Show a => a -> String
show [TypeBase ExtShape Uniqueness]
ts
    Just [DeclType]
ts' ->
      -- Make sure we always generate a binding, even for primitives.
      Certs -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ (DeclType -> InternaliseM SubExp)
-> [DeclType] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((VName -> SubExp) -> InternaliseM VName -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap VName -> SubExp
I.Var (InternaliseM VName -> InternaliseM SubExp)
-> (Exp SOACS -> InternaliseM VName)
-> Exp SOACS
-> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Exp (Rep InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m VName
letExp String
desc (Exp SOACS -> InternaliseM SubExp)
-> (DeclType -> InternaliseM (Exp SOACS))
-> DeclType
-> InternaliseM SubExp
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< TypeBase Shape NoUniqueness -> InternaliseM (Exp SOACS)
forall (m :: * -> *).
MonadBuilder m =>
TypeBase Shape NoUniqueness -> m (Exp (Rep m))
eBlank (TypeBase Shape NoUniqueness -> InternaliseM (Exp SOACS))
-> (DeclType -> TypeBase Shape NoUniqueness)
-> DeclType
-> InternaliseM (Exp SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DeclType -> TypeBase Shape NoUniqueness
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl) [DeclType]
ts'
internaliseExp String
desc (E.QualParens (QualName VName, SrcLoc)
_ Exp
e SrcLoc
_) =
  String -> Exp -> InternaliseM [SubExp]
internaliseExp String
desc Exp
e
internaliseExp String
desc (E.StringLit [Word8]
vs SrcLoc
_) =
  (SubExp -> [SubExp])
-> InternaliseM SubExp -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> [SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InternaliseM SubExp -> InternaliseM [SubExp])
-> (Exp SOACS -> InternaliseM SubExp)
-> Exp SOACS
-> InternaliseM [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
desc (Exp SOACS -> InternaliseM [SubExp])
-> Exp SOACS -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$
      [SubExp] -> TypeBase Shape NoUniqueness -> BasicOp
I.ArrayLit ((Word8 -> SubExp) -> [Word8] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map Word8 -> SubExp
forall v. IsValue v => v -> SubExp
constant [Word8]
vs) (TypeBase Shape NoUniqueness -> BasicOp)
-> TypeBase Shape NoUniqueness -> BasicOp
forall a b. (a -> b) -> a -> b
$
        PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int8
internaliseExp String
_ (E.Var (E.QualName [VName]
_ VName
name) Info PatType
_ SrcLoc
_) = do
  Maybe [SubExp]
subst <- VName -> InternaliseM (Maybe [SubExp])
lookupSubst VName
name
  case Maybe [SubExp]
subst of
    Just [SubExp]
substs -> [SubExp] -> InternaliseM [SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
substs
    Maybe [SubExp]
Nothing -> [SubExp] -> InternaliseM [SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [VName -> SubExp
I.Var VName
name]
internaliseExp String
desc (E.AppExp AppExp
e (Info AppRes
appres)) = do
  [SubExp]
ses <- String -> AppRes -> AppExp -> InternaliseM [SubExp]
internaliseAppExp String
desc AppRes
appres AppExp
e
  AppRes -> [SubExp] -> InternaliseM ()
bindExtSizes AppRes
appres [SubExp]
ses
  [SubExp] -> InternaliseM [SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
ses
internaliseExp String
_ (E.TupLit [] SrcLoc
_) =
  [SubExp] -> InternaliseM [SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
forall v. IsValue v => v -> SubExp
constant PrimValue
UnitValue]
internaliseExp String
_ (E.RecordLit [] SrcLoc
_) =
  [SubExp] -> InternaliseM [SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
forall v. IsValue v => v -> SubExp
constant PrimValue
UnitValue]
internaliseExp String
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 (String -> Exp -> InternaliseM [SubExp]
internaliseExp String
desc) [Exp]
es
internaliseExp String
desc (E.RecordLit [FieldBase Info VName]
orig_fields SrcLoc
_) =
  ((Name, [SubExp]) -> [SubExp]) -> [(Name, [SubExp])] -> [SubExp]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Name, [SubExp]) -> [SubExp]
forall a b. (a, b) -> b
snd ([(Name, [SubExp])] -> [SubExp])
-> ([Map Name [SubExp]] -> [(Name, [SubExp])])
-> [Map Name [SubExp]]
-> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name [SubExp] -> [(Name, [SubExp])]
forall a. Map Name a -> [(Name, a)]
sortFields (Map Name [SubExp] -> [(Name, [SubExp])])
-> ([Map Name [SubExp]] -> Map Name [SubExp])
-> [Map Name [SubExp]]
-> [(Name, [SubExp])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Map Name [SubExp]] -> Map Name [SubExp]
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
M.unions ([Map Name [SubExp]] -> [SubExp])
-> InternaliseM [Map Name [SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (FieldBase Info VName -> InternaliseM (Map Name [SubExp]))
-> [FieldBase Info VName] -> InternaliseM [Map Name [SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM FieldBase Info VName -> InternaliseM (Map Name [SubExp])
internaliseField [FieldBase Info VName]
orig_fields
  where
    internaliseField :: FieldBase Info VName -> InternaliseM (Map Name [SubExp])
internaliseField (E.RecordFieldExplicit Name
name 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
<$> String -> Exp -> InternaliseM [SubExp]
internaliseExp String
desc Exp
e
    internaliseField (E.RecordFieldImplicit VName
name Info PatType
t SrcLoc
loc) =
      FieldBase Info VName -> InternaliseM (Map Name [SubExp])
internaliseField (FieldBase Info VName -> InternaliseM (Map Name [SubExp]))
-> FieldBase Info VName -> InternaliseM (Map Name [SubExp])
forall a b. (a -> b) -> a -> b
$
        Name -> 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 PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
E.Var (VName -> QualName VName
forall v. v -> QualName v
E.qualName VName
name) Info PatType
t SrcLoc
loc)
          SrcLoc
loc
internaliseExp String
desc (E.ArrayLit [Exp]
es (Info PatType
arr_t) SrcLoc
loc)
  -- If this is a multidimensional array literal of primitives, we
  -- treat it specially by flattening it out followed by a reshape.
  -- This cuts down on the amount of statements that are produced, and
  -- thus allows us to efficiently handle huge array literals - a
  -- corner case, but an important one.
  | Just (([Int]
eshape, [Exp]
e') : [([Int], [Exp])]
es') <- (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 PatType
basetype <- Int -> PatType -> Maybe PatType
forall dim as. Int -> TypeBase dim as -> Maybe (TypeBase dim as)
E.peelArray ([Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
eshape) PatType
arr_t = do
      let flat_lit :: Exp
flat_lit = [Exp] -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
[ExpBase f vn] -> f PatType -> 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') (PatType -> Info PatType
forall a. a -> Info a
Info PatType
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 <- String -> Exp -> InternaliseM [VName]
internaliseExpToVars String
"flat_literal" Exp
flat_lit
      [VName] -> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
flat_arrs ((VName -> InternaliseM SubExp) -> InternaliseM [SubExp])
-> (VName -> InternaliseM SubExp) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
flat_arr -> do
        TypeBase Shape NoUniqueness
flat_arr_t <- VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
flat_arr
        let new_shape' :: Shape
new_shape' =
              Shape -> Int -> Shape -> Shape
reshapeOuter
                ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
I.Shape ([SubExp] -> Shape) -> [SubExp] -> Shape
forall a b. (a -> b) -> a -> b
$ (Int -> SubExp) -> [Int] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (IntType -> Integer -> SubExp
intConst IntType
Int64 (Integer -> SubExp) -> (Int -> Integer) -> Int -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Integer
forall a. Integral a => a -> Integer
toInteger) [Int]
new_shape)
                Int
1
                (Shape -> Shape) -> Shape -> Shape
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoUniqueness -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape TypeBase Shape NoUniqueness
flat_arr_t
        String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
desc (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ ReshapeKind -> Shape -> VName -> BasicOp
I.Reshape ReshapeKind
I.ReshapeArbitrary Shape
new_shape' VName
flat_arr
  | Bool
otherwise = do
      [[SubExp]]
es' <- (Exp -> InternaliseM [SubExp]) -> [Exp] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> Exp -> InternaliseM [SubExp]
internaliseExp String
"arr_elem") [Exp]
es
      let arr_t_ext :: [TypeBase ExtShape Uniqueness]
arr_t_ext = TypeBase Size () -> [TypeBase ExtShape Uniqueness]
internaliseType (TypeBase Size () -> [TypeBase ExtShape Uniqueness])
-> TypeBase Size () -> [TypeBase ExtShape Uniqueness]
forall a b. (a -> b) -> a -> b
$ PatType -> TypeBase Size ()
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatType
arr_t

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  SubExp
slice_ok <-
    String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"slice_ok"
      (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
        (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
backwards)
        ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
backwards_ok])
        ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
forwards_ok])

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

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

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

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

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

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

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

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

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

internaliseStreamMap ::
  String ->
  StreamOrd ->
  E.Exp ->
  E.Exp ->
  InternaliseM [SubExp]
internaliseStreamMap :: String -> StreamOrd -> Exp -> Exp -> InternaliseM [SubExp]
internaliseStreamMap String
desc StreamOrd
o Exp
lam Exp
arr = do
  [VName]
arrs <- String -> Exp -> InternaliseM [VName]
internaliseExpToVars String
"stream_input" Exp
arr
  Lambda SOACS
lam' <- InternaliseLambda -> Exp -> [SubExp] -> InternaliseM (Lambda SOACS)
internaliseStreamMapLambda InternaliseLambda
internaliseLambda Exp
lam ([SubExp] -> InternaliseM (Lambda SOACS))
-> [SubExp] -> InternaliseM (Lambda SOACS)
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var [VName]
arrs
  SubExp
w <- Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 ([TypeBase Shape NoUniqueness] -> SubExp)
-> InternaliseM [TypeBase Shape NoUniqueness]
-> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
  let form :: StreamForm SOACS
form = StreamOrd -> Commutativity -> Lambda SOACS -> StreamForm SOACS
forall rep.
StreamOrd -> Commutativity -> Lambda rep -> StreamForm rep
I.Parallel StreamOrd
o Commutativity
Commutative ([LParam SOACS]
-> Body SOACS -> [TypeBase Shape NoUniqueness] -> Lambda SOACS
forall rep.
[LParam rep]
-> Body rep -> [TypeBase Shape NoUniqueness] -> Lambda rep
I.Lambda [] (Stms SOACS -> Result -> Body SOACS
forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody Stms SOACS
forall a. Monoid a => a
mempty []) [])
  String -> Exp SOACS -> InternaliseM [SubExp]
letValExp' String
desc (Exp SOACS -> InternaliseM [SubExp])
-> Exp SOACS -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Op SOACS -> Exp SOACS
forall rep. Op rep -> Exp rep
I.Op (Op SOACS -> Exp SOACS) -> Op SOACS -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ SubExp
-> [VName]
-> StreamForm SOACS
-> [SubExp]
-> Lambda SOACS
-> SOAC SOACS
forall rep.
SubExp
-> [VName] -> StreamForm rep -> [SubExp] -> Lambda rep -> SOAC rep
I.Stream SubExp
w [VName]
arrs StreamForm SOACS
form [] Lambda SOACS
lam'

internaliseStreamRed ::
  String ->
  StreamOrd ->
  Commutativity ->
  E.Exp ->
  E.Exp ->
  E.Exp ->
  InternaliseM [SubExp]
internaliseStreamRed :: String
-> StreamOrd
-> Commutativity
-> Exp
-> Exp
-> Exp
-> InternaliseM [SubExp]
internaliseStreamRed String
desc StreamOrd
o Commutativity
comm Exp
lam0 Exp
lam Exp
arr = do
  [VName]
arrs <- String -> Exp -> InternaliseM [VName]
internaliseExpToVars String
"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 u. TypeBase Shape u -> TypeBase Shape u
I.rowType (InternaliseM (TypeBase Shape NoUniqueness)
 -> InternaliseM (TypeBase Shape NoUniqueness))
-> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> VName
-> InternaliseM (TypeBase Shape NoUniqueness)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType) [VName]
arrs
  ([Param (TypeBase Shape NoUniqueness)]
lam_params, Body SOACS
lam_body) <-
    InternaliseLambda
-> Exp
-> [TypeBase Shape NoUniqueness]
-> InternaliseM ([LParam SOACS], Body 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 (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param (TypeBase Shape NoUniqueness) -> VName
forall dec. Param dec -> VName
I.paramName Param (TypeBase Shape NoUniqueness)
chunk_param] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
    BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$
      SubExp -> BasicOp
I.SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$
        Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)
  [Param (TypeBase Shape NoUniqueness)]
-> (Param (TypeBase Shape NoUniqueness) -> InternaliseM ())
-> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Param (TypeBase Shape NoUniqueness)]
lam_val_params ((Param (TypeBase Shape NoUniqueness) -> InternaliseM ())
 -> InternaliseM ())
-> (Param (TypeBase Shape NoUniqueness) -> InternaliseM ())
-> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \Param (TypeBase Shape NoUniqueness)
p ->
    [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param (TypeBase Shape NoUniqueness) -> VName
forall dec. Param dec -> VName
I.paramName Param (TypeBase Shape NoUniqueness)
p] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
      BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS)
-> ([SubExp] -> BasicOp) -> [SubExp] -> Exp SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimType -> [SubExp] -> BasicOp
I.Scratch (TypeBase Shape NoUniqueness -> PrimType
forall shape u. TypeBase shape u -> PrimType
I.elemType (TypeBase Shape NoUniqueness -> PrimType)
-> TypeBase Shape NoUniqueness -> PrimType
forall a b. (a -> b) -> a -> b
$ Param (TypeBase Shape NoUniqueness) -> TypeBase Shape NoUniqueness
forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
I.paramType Param (TypeBase Shape NoUniqueness)
p) ([SubExp] -> Exp SOACS) -> [SubExp] -> Exp SOACS
forall a b. (a -> b) -> a -> b
$
        TypeBase Shape NoUniqueness -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
I.arrayDims (TypeBase Shape NoUniqueness -> [SubExp])
-> TypeBase Shape NoUniqueness -> [SubExp]
forall a b. (a -> b) -> a -> b
$
          Param (TypeBase Shape NoUniqueness) -> TypeBase Shape NoUniqueness
forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
I.paramType Param (TypeBase Shape NoUniqueness)
p
  Result
nes <- Body SOACS -> InternaliseM Result
forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind (Body SOACS -> InternaliseM Result)
-> InternaliseM (Body SOACS) -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Body SOACS -> InternaliseM (Body SOACS)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Body rep -> m (Body rep)
renameBody Body SOACS
lam_body

  [TypeBase Shape NoUniqueness]
nes_ts <- (SubExpRes -> InternaliseM (TypeBase Shape NoUniqueness))
-> Result -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExpRes -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExpRes -> m (TypeBase Shape NoUniqueness)
I.subExpResType Result
nes
  SubExp
outsz <- Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 ([TypeBase Shape NoUniqueness] -> SubExp)
-> InternaliseM [TypeBase Shape NoUniqueness]
-> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
  let acc_arr_tps :: [TypeBase Shape NoUniqueness]
acc_arr_tps = [TypeBase Shape NoUniqueness
-> Shape -> NoUniqueness -> TypeBase Shape NoUniqueness
forall shape u_unused u.
ArrayShape shape =>
TypeBase shape u_unused -> shape -> u -> TypeBase shape u
I.arrayOf TypeBase Shape NoUniqueness
t ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
I.Shape [SubExp
outsz]) NoUniqueness
NoUniqueness | TypeBase Shape NoUniqueness
t <- [TypeBase Shape NoUniqueness]
nes_ts]
  Lambda SOACS
lam0' <- InternaliseLambda
-> Exp
-> [TypeBase Shape NoUniqueness]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM (Lambda SOACS)
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 (Result -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Result
nes) ([Param (TypeBase Shape NoUniqueness)]
 -> [Param (TypeBase Shape NoUniqueness)])
-> [Param (TypeBase Shape NoUniqueness)]
-> [Param (TypeBase Shape NoUniqueness)]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> [LParam SOACS]
forall rep. Lambda rep -> [LParam rep]
I.lambdaParams Lambda SOACS
lam0'
  [Param (TypeBase Shape NoUniqueness)]
lam_acc_params <- [Param (TypeBase Shape NoUniqueness)]
-> (Param (TypeBase Shape NoUniqueness)
    -> InternaliseM (Param (TypeBase Shape NoUniqueness)))
-> InternaliseM [Param (TypeBase Shape NoUniqueness)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Param (TypeBase Shape NoUniqueness)]
lam0_acc_params ((Param (TypeBase Shape NoUniqueness)
  -> InternaliseM (Param (TypeBase Shape NoUniqueness)))
 -> InternaliseM [Param (TypeBase Shape NoUniqueness)])
-> (Param (TypeBase Shape NoUniqueness)
    -> InternaliseM (Param (TypeBase Shape NoUniqueness)))
-> InternaliseM [Param (TypeBase Shape NoUniqueness)]
forall a b. (a -> b) -> a -> b
$ \Param (TypeBase Shape NoUniqueness)
p -> do
    VName
name <- String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName (String -> InternaliseM VName) -> String -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$ VName -> String
baseString (VName -> String) -> VName -> String
forall a b. (a -> b) -> a -> b
$ Param (TypeBase Shape NoUniqueness) -> VName
forall dec. Param dec -> VName
I.paramName Param (TypeBase Shape NoUniqueness)
p
    Param (TypeBase Shape NoUniqueness)
-> InternaliseM (Param (TypeBase Shape NoUniqueness))
forall (f :: * -> *) a. Applicative f => a -> f a
pure 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 SOACS
lam' <- [LParam (Rep InternaliseM)]
-> InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
mkLambda [Param (TypeBase Shape NoUniqueness)]
[LParam (Rep InternaliseM)]
lam_params' (InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM)))
-> InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM))
forall a b. (a -> b) -> a -> b
$ do
    Result
lam_res <- Body (Rep InternaliseM) -> InternaliseM Result
forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind Body (Rep InternaliseM)
Body 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 SOACS -> [LParam SOACS]
forall rep. Lambda rep -> [LParam rep]
I.lambdaParams Lambda SOACS
lam0')
        ((SubExpRes -> SubExp) -> Result -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp Result
lam_res)
    ErrorMsg SubExp
-> SrcLoc
-> [TypeBase Shape NoUniqueness]
-> Result
-> InternaliseM Result
ensureResultShape
      ErrorMsg SubExp
"shape of result does not match shape of initial value"
      (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
lam0)
      [TypeBase Shape NoUniqueness]
nes_ts
      (Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ( Lambda (Rep InternaliseM)
-> [InternaliseM (Exp (Rep InternaliseM))] -> InternaliseM Result
forall (m :: * -> *).
MonadBuilder m =>
Lambda (Rep m) -> [m (Exp (Rep m))] -> m Result
eLambda Lambda (Rep InternaliseM)
Lambda SOACS
lam0' ([InternaliseM (Exp SOACS)] -> InternaliseM Result)
-> ([SubExp] -> [InternaliseM (Exp SOACS)])
-> [SubExp]
-> InternaliseM Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SubExp -> InternaliseM (Exp SOACS))
-> [SubExp] -> [InternaliseM (Exp SOACS)]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> InternaliseM (Exp SOACS)
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp ([SubExp] -> InternaliseM Result)
-> [SubExp] -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
              (Param (TypeBase Shape NoUniqueness) -> SubExp)
-> [Param (TypeBase Shape NoUniqueness)] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
I.Var (VName -> SubExp)
-> (Param (TypeBase Shape NoUniqueness) -> VName)
-> Param (TypeBase Shape NoUniqueness)
-> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param (TypeBase Shape NoUniqueness) -> VName
forall dec. Param dec -> VName
paramName) [Param (TypeBase Shape NoUniqueness)]
lam_acc_params [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
lam_res'
          )

  let form :: StreamForm SOACS
form = StreamOrd -> Commutativity -> Lambda SOACS -> StreamForm SOACS
forall rep.
StreamOrd -> Commutativity -> Lambda rep -> StreamForm rep
I.Parallel StreamOrd
o Commutativity
comm Lambda SOACS
lam0'
  SubExp
w <- Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 ([TypeBase Shape NoUniqueness] -> SubExp)
-> InternaliseM [TypeBase Shape NoUniqueness]
-> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
  String -> Exp SOACS -> InternaliseM [SubExp]
letValExp' String
desc (Exp SOACS -> InternaliseM [SubExp])
-> Exp SOACS -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Op SOACS -> Exp SOACS
forall rep. Op rep -> Exp rep
I.Op (Op SOACS -> Exp SOACS) -> Op SOACS -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ SubExp
-> [VName]
-> StreamForm SOACS
-> [SubExp]
-> Lambda SOACS
-> SOAC SOACS
forall rep.
SubExp
-> [VName] -> StreamForm rep -> [SubExp] -> Lambda rep -> SOAC rep
I.Stream SubExp
w [VName]
arrs StreamForm SOACS
form ((SubExpRes -> SubExp) -> Result -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp Result
nes) Lambda SOACS
lam'

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

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

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

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

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

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

internaliseExpToVars :: String -> E.Exp -> InternaliseM [I.VName]
internaliseExpToVars :: String -> Exp -> InternaliseM [VName]
internaliseExpToVars String
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
=<< String -> Exp -> InternaliseM [SubExp]
internaliseExp String
desc Exp
e
  where
    asIdent :: SubExp -> InternaliseM VName
asIdent (I.Var VName
v) = VName -> InternaliseM VName
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v
    asIdent SubExp
se = String -> Exp (Rep InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m VName
letExp String
desc (Exp (Rep InternaliseM) -> InternaliseM VName)
-> Exp (Rep InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
se

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

certifyingNonzero ::
  SrcLoc ->
  IntType ->
  SubExp ->
  InternaliseM a ->
  InternaliseM a
certifyingNonzero :: forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
x InternaliseM a
m = do
  SubExp
zero <-
    String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"zero" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$
        CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp (PrimType -> CmpOp
CmpEq (IntType -> PrimType
IntType IntType
t)) SubExp
x (IntType -> Integer -> SubExp
intConst IntType
t Integer
0)
  SubExp
nonzero <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"nonzero" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
UnOp UnOp
I.Not SubExp
zero
  Certs
c <- String -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert String
"nonzero_cert" SubExp
nonzero ErrorMsg SubExp
"division by zero" SrcLoc
loc
  Certs -> InternaliseM a -> InternaliseM a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c InternaliseM a
m

certifyingNonnegative ::
  SrcLoc ->
  IntType ->
  SubExp ->
  InternaliseM a ->
  InternaliseM a
certifyingNonnegative :: forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonnegative SrcLoc
loc IntType
t SubExp
x InternaliseM a
m = do
  SubExp
nonnegative <-
    String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"nonnegative" (Exp SOACS -> InternaliseM SubExp)
-> (BasicOp -> Exp SOACS) -> BasicOp -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> InternaliseM SubExp) -> BasicOp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp (IntType -> CmpOp
CmpSle IntType
t) (IntType -> Integer -> SubExp
intConst IntType
t Integer
0) SubExp
x
  Certs
c <- String -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert String
"nonzero_cert" SubExp
nonnegative ErrorMsg SubExp
"negative exponent" SrcLoc
loc
  Certs -> InternaliseM a -> InternaliseM a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c InternaliseM a
m

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

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

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

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

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

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

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

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

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

-- | 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 (String -> 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] -> String -> Maybe (String -> InternaliseM [SubExp])]
handlers =
        [ [Exp] -> String -> Maybe (String -> InternaliseM [SubExp])
forall {a}.
(Eq a, IsString a) =>
[Exp] -> a -> Maybe (String -> InternaliseM [SubExp])
handleSign,
          [Exp] -> String -> Maybe (String -> InternaliseM [SubExp])
forall {f :: * -> *}.
Applicative f =>
[Exp] -> String -> Maybe (String -> InternaliseM (f SubExp))
handleIntrinsicOps,
          [Exp] -> String -> Maybe (String -> InternaliseM [SubExp])
handleOps,
          [Exp] -> String -> Maybe (String -> InternaliseM [SubExp])
handleSOACs,
          [Exp] -> String -> Maybe (String -> InternaliseM [SubExp])
forall {a}.
(Eq a, IsString a) =>
[Exp] -> a -> Maybe (String -> InternaliseM [SubExp])
handleAccs,
          [Exp] -> String -> Maybe (String -> InternaliseM [SubExp])
forall {a}.
(IsString a, Eq a) =>
[Exp] -> a -> Maybe (String -> InternaliseM [SubExp])
handleAD,
          [Exp] -> String -> Maybe (String -> InternaliseM [SubExp])
handleRest
        ]
  [Maybe (String -> InternaliseM [SubExp])]
-> Maybe (String -> InternaliseM [SubExp])
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, MonadPlus m) =>
t (m a) -> m a
msum [[Exp] -> String -> Maybe (String -> InternaliseM [SubExp])
h [Exp]
args (String -> Maybe (String -> InternaliseM [SubExp]))
-> String -> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ VName -> String
baseString (VName -> String) -> VName -> String
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qname | [Exp] -> String -> Maybe (String -> InternaliseM [SubExp])
h <- [[Exp] -> String -> Maybe (String -> InternaliseM [SubExp])]
handlers]
  where
    handleSign :: [Exp] -> a -> Maybe (String -> InternaliseM [SubExp])
handleSign [Exp
x] a
"sign_i8" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> String -> InternaliseM [SubExp]
toSigned IntType
I.Int8 Exp
x
    handleSign [Exp
x] a
"sign_i16" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> String -> InternaliseM [SubExp]
toSigned IntType
I.Int16 Exp
x
    handleSign [Exp
x] a
"sign_i32" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> String -> InternaliseM [SubExp]
toSigned IntType
I.Int32 Exp
x
    handleSign [Exp
x] a
"sign_i64" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> String -> InternaliseM [SubExp]
toSigned IntType
I.Int64 Exp
x
    handleSign [Exp
x] a
"unsign_i8" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> String -> InternaliseM [SubExp]
toUnsigned IntType
I.Int8 Exp
x
    handleSign [Exp
x] a
"unsign_i16" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> String -> InternaliseM [SubExp]
toUnsigned IntType
I.Int16 Exp
x
    handleSign [Exp
x] a
"unsign_i32" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> String -> InternaliseM [SubExp]
toUnsigned IntType
I.Int32 Exp
x
    handleSign [Exp
x] a
"unsign_i64" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> String -> InternaliseM [SubExp]
toUnsigned IntType
I.Int64 Exp
x
    handleSign [Exp]
_ a
_ = Maybe (String -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

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

    -- Short-circuiting operators are magical.
    handleOps :: [Exp] -> String -> Maybe (String -> InternaliseM [SubExp])
handleOps [Exp
x, Exp
y] String
"&&" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String -> Exp -> InternaliseM [SubExp]
internaliseExp String
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
$ PatType -> [VName] -> AppRes
AppRes (ScalarTypeBase Size (Set Alias) -> PatType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar (ScalarTypeBase Size (Set Alias) -> PatType)
-> ScalarTypeBase Size (Set Alias) -> PatType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Size (Set Alias)
forall dim as. PrimType -> ScalarTypeBase dim as
E.Prim PrimType
E.Bool) [])
    handleOps [Exp
x, Exp
y] String
"||" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String -> Exp -> InternaliseM [SubExp]
internaliseExp String
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
$ PatType -> [VName] -> AppRes
AppRes (ScalarTypeBase Size (Set Alias) -> PatType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar (ScalarTypeBase Size (Set Alias) -> PatType)
-> ScalarTypeBase Size (Set Alias) -> PatType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Size (Set Alias)
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] String
op
      | Just String -> SubExp -> InternaliseM [SubExp]
cmp_f <- String -> Maybe (String -> SubExp -> InternaliseM [SubExp])
forall {a} {m :: * -> *}.
(IsString a, MonadBuilder m, Eq a) =>
a -> Maybe (String -> SubExp -> m [SubExp])
isEqlOp String
op = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc -> do
          [SubExp]
xe' <- String -> Exp -> InternaliseM [SubExp]
internaliseExp String
"x" Exp
xe
          [SubExp]
ye' <- String -> Exp -> InternaliseM [SubExp]
internaliseExp String
"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 (String -> SubExp -> SubExp -> InternaliseM SubExp
forall {m :: * -> *}.
(MonadBuilder m, Buildable (Rep m), BuilderOps (Rep m),
 Op (Rep m) ~ SOAC (Rep m)) =>
String -> SubExp -> SubExp -> m SubExp
doComparison String
desc) [SubExp]
xe' [SubExp]
ye'
          String -> SubExp -> InternaliseM [SubExp]
cmp_f String
desc (SubExp -> InternaliseM [SubExp])
-> InternaliseM SubExp -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"eq" (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [SubExp] -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
eAll [SubExp]
rs
      where
        isEqlOp :: a -> Maybe (String -> SubExp -> m [SubExp])
isEqlOp a
"!=" = (String -> SubExp -> m [SubExp])
-> Maybe (String -> SubExp -> m [SubExp])
forall a. a -> Maybe a
Just ((String -> SubExp -> m [SubExp])
 -> Maybe (String -> SubExp -> m [SubExp]))
-> (String -> SubExp -> m [SubExp])
-> Maybe (String -> SubExp -> m [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc SubExp
eq ->
          String -> Exp (Rep m) -> m [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [SubExp]
letTupExp' String
desc (Exp (Rep m) -> m [SubExp]) -> Exp (Rep m) -> m [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
I.Not SubExp
eq
        isEqlOp a
"==" = (String -> SubExp -> m [SubExp])
-> Maybe (String -> SubExp -> m [SubExp])
forall a. a -> Maybe a
Just ((String -> SubExp -> m [SubExp])
 -> Maybe (String -> SubExp -> m [SubExp]))
-> (String -> SubExp -> m [SubExp])
-> Maybe (String -> SubExp -> m [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
_ SubExp
eq ->
          [SubExp] -> m [SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp
eq]
        isEqlOp a
_ = Maybe (String -> SubExp -> m [SubExp])
forall a. Maybe a
Nothing

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

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

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

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

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

    handleAccs :: [Exp] -> a -> Maybe (String -> InternaliseM [SubExp])
handleAccs [TupLit [Exp
dest, Exp
f, Exp
bs] SrcLoc
_] a
"scatter_stream" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> Exp -> Maybe (Exp, Exp) -> Exp -> Exp -> InternaliseM [SubExp]
internaliseStreamAcc String
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" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> Exp -> Maybe (Exp, Exp) -> Exp -> Exp -> InternaliseM [SubExp]
internaliseStreamAcc String
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" = (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((String -> InternaliseM [SubExp])
 -> Maybe (String -> InternaliseM [SubExp]))
-> (String -> InternaliseM [SubExp])
-> Maybe (String -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \String
desc -> do
      VName
acc' <- [VName] -> VName
forall a. [a] -> a
head ([VName] -> VName) -> InternaliseM [VName] -> InternaliseM VName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> Exp -> InternaliseM [VName]
internaliseExpToVars String
"acc" Exp
acc
      SubExp
i' <- String -> Exp -> InternaliseM SubExp
internaliseExp1 String
"acc_i" Exp
i
      [SubExp]
vs <- String -> Exp -> InternaliseM [SubExp]
internaliseExp String
"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
$ String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
desc (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ VName -> [SubExp] -> [SubExp] -> BasicOp
UpdateAcc VName
acc' [SubExp
i'] [SubExp]
vs
    handleAccs [Exp]
_ a
_ = Maybe (String -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

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

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

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

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

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

    scatterF :: Int -> Exp -> Exp -> Exp -> String -> InternaliseM [SubExp]
scatterF Int
dim Exp
a Exp
si Exp
v String
desc = do
      [VName]
si' <- String -> Exp -> InternaliseM [VName]
internaliseExpToVars String
"write_arg_i" Exp
si
      [VName]
svs <- String -> Exp -> InternaliseM [VName]
internaliseExpToVars String
"write_arg_v" Exp
v
      [VName]
sas <- String -> Exp -> InternaliseM [VName]
internaliseExpToVars String
"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 rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
si'
      [TypeBase Shape NoUniqueness]
sv_ts <- (VName -> InternaliseM (TypeBase Shape NoUniqueness))
-> [VName] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
svs

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  [(VName, SubExp)]
-> ((VName, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Map VName SubExp -> [(VName, SubExp)]
forall k a. Map k a -> [(k, a)]
M.toList (Map VName SubExp -> [(VName, SubExp)])
-> Map VName SubExp -> [(VName, SubExp)]
forall a b. (a -> b) -> a -> b
$ [Map VName SubExp] -> Map VName SubExp
forall a. Monoid a => [a] -> a
mconcat ([Map VName SubExp] -> Map VName SubExp)
-> [Map VName SubExp] -> Map VName SubExp
forall a b. (a -> b) -> a -> b
$ (TypeBase ExtShape Uniqueness
 -> TypeBase Shape NoUniqueness -> Map VName SubExp)
-> [TypeBase ExtShape Uniqueness]
-> [TypeBase Shape NoUniqueness]
-> [Map VName SubExp]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeBase ExtShape Uniqueness
-> TypeBase Shape NoUniqueness -> Map VName SubExp
combine [TypeBase ExtShape Uniqueness]
ts [TypeBase Shape NoUniqueness]
ses_ts) (((VName, SubExp) -> InternaliseM ()) -> InternaliseM ())
-> ((VName, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \(VName
v, SubExp
se) ->
    [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
v] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp 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 (f :: * -> *) a. Applicative f => a -> f a
pure (Safety -> InternaliseM Safety) -> Safety -> InternaliseM Safety
forall a b. (a -> b) -> a -> b
$ if Bool
check then Safety
I.Safe else Safety
I.Unsafe

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

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

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

  -- We have the offsets for each of the partitions, but we also need
  -- the total sizes, which are the last elements in the offests.  We
  -- just have to be careful in case the array is empty.
  SubExp
last_index <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"last_index" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp 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)
  let nonempty_body :: InternaliseM (Body SOACS)
nonempty_body = Builder SOACS (Body SOACS) -> InternaliseM (Body SOACS)
forall rep (m :: * -> *) somerep.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
 SameScope somerep rep) =>
Builder rep (Body rep) -> m (Body rep)
runBodyBuilder (Builder SOACS (Body SOACS) -> InternaliseM (Body SOACS))
-> Builder SOACS (Body SOACS) -> InternaliseM (Body SOACS)
forall a b. (a -> b) -> a -> b
$
        ([SubExp] -> Body SOACS)
-> BuilderT SOACS (State VNameSource) [SubExp]
-> Builder SOACS (Body SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [SubExp] -> Body SOACS
forall rep. Buildable rep => [SubExp] -> Body rep
resultBody (BuilderT SOACS (State VNameSource) [SubExp]
 -> Builder SOACS (Body SOACS))
-> BuilderT SOACS (State VNameSource) [SubExp]
-> Builder SOACS (Body SOACS)
forall a b. (a -> b) -> a -> b
$
          [VName]
-> (VName -> BuilderT SOACS (State VNameSource) SubExp)
-> BuilderT SOACS (State VNameSource) [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
all_offsets ((VName -> BuilderT SOACS (State VNameSource) SubExp)
 -> BuilderT SOACS (State VNameSource) [SubExp])
-> (VName -> BuilderT SOACS (State VNameSource) SubExp)
-> BuilderT SOACS (State VNameSource) [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
offset_array ->
            String
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"last_offset" (Exp (Rep (BuilderT SOACS (State VNameSource)))
 -> BuilderT SOACS (State VNameSource) SubExp)
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp
I.Index VName
offset_array (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ [DimIndex SubExp] -> Slice SubExp
forall d. [DimIndex d] -> Slice d
Slice [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
I.DimFix SubExp
last_index]
      empty_body :: InternaliseM (Body (Rep InternaliseM))
empty_body = [SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM ([SubExp] -> InternaliseM (Body (Rep InternaliseM)))
-> [SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall a b. (a -> b) -> a -> b
$ Int -> SubExp -> [SubExp]
forall a. Int -> a -> [a]
replicate Int
k (SubExp -> [SubExp]) -> SubExp -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)
  SubExp
is_empty <- String -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"is_empty" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
CmpEq PrimType
int64) SubExp
w (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)
  [VName]
sizes <-
    String -> Exp (Rep InternaliseM) -> InternaliseM [VName]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [VName]
letTupExp String
"partition_size" (Exp SOACS -> InternaliseM [VName])
-> InternaliseM (Exp SOACS) -> InternaliseM [VName]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
is_empty) InternaliseM (Body (Rep InternaliseM))
empty_body InternaliseM (Body (Rep InternaliseM))
InternaliseM (Body SOACS)
nonempty_body

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

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

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 (f :: * -> *) a. Applicative f => a -> f a
pure [String -> ErrorMsgPart SubExp
forall a. String -> ErrorMsgPart a
ErrorString (String -> ErrorMsgPart SubExp) -> String -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ QualName VName -> String
forall a. Pretty a => a -> String
pretty QualName VName
qn]
typeExpForError (E.TEUnique TypeExp VName
te SrcLoc
_) =
  (ErrorMsgPart SubExp
"*" ErrorMsgPart SubExp
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. a -> [a] -> [a]
:) ([ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp])
-> InternaliseM [ErrorMsgPart SubExp]
-> InternaliseM [ErrorMsgPart SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
te
typeExpForError (E.TEDim [VName]
dims TypeExp VName
te SrcLoc
_) =
  (String -> ErrorMsgPart SubExp
forall a. String -> ErrorMsgPart a
ErrorString (String
"?" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
dims' String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
".") ErrorMsgPart SubExp
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. a -> [a] -> [a]
:) ([ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp])
-> InternaliseM [ErrorMsgPart SubExp]
-> InternaliseM [ErrorMsgPart SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
te
  where
    dims' :: String
dims' = [String] -> String
forall a. Monoid a => [a] -> a
mconcat ((VName -> String) -> [VName] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map VName -> String
forall a. Pretty a => a -> String
onDim [VName]
dims)
    onDim :: a -> String
onDim a
d = String
"[" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> a -> String
forall a. Pretty a => a -> String
pretty a
d String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"]"
typeExpForError (E.TEArray SizeExp VName
d TypeExp VName
te SrcLoc
_) = do
  ErrorMsgPart SubExp
d' <- SizeExp VName -> InternaliseM (ErrorMsgPart SubExp)
dimExpForError SizeExp VName
d
  [ErrorMsgPart SubExp]
te' <- TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
te
  [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp
"[", ErrorMsgPart SubExp
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 (f :: * -> *) a. Applicative f => a -> f a
pure ([ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp
"("] [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
-> [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
", "] [[ErrorMsgPart SubExp]]
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 (f :: * -> *) a. Applicative f => a -> f a
pure ([ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp
"{"] [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
-> [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
", "] [[ErrorMsgPart SubExp]]
fields' [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"}"]
  where
    onField :: (a, TypeExp VName) -> InternaliseM [ErrorMsgPart SubExp]
onField (a
k, TypeExp VName
te) =
      (String -> ErrorMsgPart SubExp
forall a. String -> ErrorMsgPart a
ErrorString (a -> String
forall a. Pretty a => a -> String
pretty a
k String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
": ") ErrorMsgPart SubExp
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. a -> [a] -> [a]
:) ([ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp])
-> InternaliseM [ErrorMsgPart SubExp]
-> InternaliseM [ErrorMsgPart SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
te
typeExpForError (E.TEArrow Maybe VName
_ TypeExp VName
t1 TypeExp VName
t2 SrcLoc
_) = do
  [ErrorMsgPart SubExp]
t1' <- TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
t1
  [ErrorMsgPart SubExp]
t2' <- TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
t2
  [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp]
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 SizeExp 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
<$> SizeExp VName -> InternaliseM (ErrorMsgPart SubExp)
dimExpForError SizeExp VName
d
  [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp]
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 (f :: * -> *) a. Applicative f => a -> f a
pure ([ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp]
-> [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
" | "] [[ErrorMsgPart SubExp]]
cs'
  where
    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 (f :: * -> *) a. Applicative f => a -> f a
pure ([ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp])
-> [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp]
-> [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
" "] [[ErrorMsgPart SubExp]]
c'

dimExpForError :: E.SizeExp VName -> InternaliseM (ErrorMsgPart SubExp)
dimExpForError :: SizeExp VName -> InternaliseM (ErrorMsgPart SubExp)
dimExpForError (SizeExpNamed QualName VName
d SrcLoc
_) = do
  Maybe [SubExp]
substs <- VName -> InternaliseM (Maybe [SubExp])
lookupSubst (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 (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
v
    Maybe [SubExp]
_ -> 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
$ 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 (f :: * -> *) a. Applicative f => a -> f a
pure (ErrorMsgPart SubExp -> InternaliseM (ErrorMsgPart SubExp))
-> ErrorMsgPart SubExp -> InternaliseM (ErrorMsgPart SubExp)
forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
d'
dimExpForError (SizeExpConst Int
d SrcLoc
_) =
  ErrorMsgPart SubExp -> InternaliseM (ErrorMsgPart SubExp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ErrorMsgPart SubExp -> InternaliseM (ErrorMsgPart SubExp))
-> ErrorMsgPart SubExp -> InternaliseM (ErrorMsgPart SubExp)
forall a b. (a -> b) -> a -> b
$ String -> ErrorMsgPart SubExp
forall a. String -> ErrorMsgPart a
ErrorString (String -> ErrorMsgPart SubExp) -> String -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ Int -> String
forall a. Pretty a => a -> String
pretty Int
d
dimExpForError SizeExp VName
SizeExpAny = ErrorMsgPart SubExp -> InternaliseM (ErrorMsgPart SubExp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ErrorMsgPart SubExp
""

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