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

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

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

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

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

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

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

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

internaliseValBind :: E.ValBind -> InternaliseM ()
internaliseValBind :: ValBind -> InternaliseM ()
internaliseValBind fb :: ValBind
fb@(E.ValBind Maybe (Info EntryPoint)
entry VName
fname Maybe (TypeExp VName)
retdecl (Info (StructType
rettype, [VName]
_)) [TypeParamBase VName]
tparams [PatternBase Info VName]
params ExpBase Info VName
body Maybe DocComment
_ [AttrInfo]
attrs SrcLoc
loc) = do
  InternaliseM () -> InternaliseM ()
forall a. InternaliseM a -> InternaliseM a
localConstsScope (InternaliseM () -> InternaliseM ())
-> InternaliseM () -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
    [TypeParamBase VName]
-> [PatternBase Info VName]
-> ([FParam] -> [[FParam]] -> InternaliseM ())
-> InternaliseM ()
forall a.
[TypeParamBase VName]
-> [PatternBase Info VName]
-> ([FParam] -> [[FParam]] -> InternaliseM a)
-> InternaliseM a
bindingParams [TypeParamBase VName]
tparams [PatternBase Info VName]
params (([FParam] -> [[FParam]] -> InternaliseM ()) -> InternaliseM ())
-> ([FParam] -> [[FParam]] -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \[FParam]
shapeparams [[FParam]]
params' -> do
      let shapenames :: [VName]
shapenames = (Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
[FParam]
shapeparams

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

      (([TypeBase ExtShape Uniqueness]
rettype', Result
body_res), Stms SOACS
body_stms) <- InternaliseM ([TypeBase ExtShape Uniqueness], Result)
-> InternaliseM
     (([TypeBase ExtShape Uniqueness], Result),
      Stms (Lore InternaliseM))
forall (m :: * -> *) a.
MonadBinder m =>
m a -> m (a, Stms (Lore m))
collectStms (InternaliseM ([TypeBase ExtShape Uniqueness], Result)
 -> InternaliseM
      (([TypeBase ExtShape Uniqueness], Result),
       Stms (Lore InternaliseM)))
-> InternaliseM ([TypeBase ExtShape Uniqueness], Result)
-> InternaliseM
     (([TypeBase ExtShape Uniqueness], Result),
      Stms (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ do
        Result
body_res <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp (VName -> String
baseString VName
fname String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"_res") ExpBase Info VName
body
        [TypeBase ExtShape Uniqueness]
rettype_bad <- StructType -> InternaliseM [TypeBase ExtShape Uniqueness]
internaliseReturnType StructType
rettype
        let rettype' :: [TypeBase ExtShape Uniqueness]
rettype' = [TypeBase ExtShape Uniqueness] -> [TypeBase ExtShape Uniqueness]
forall u. [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts [TypeBase ExtShape Uniqueness]
rettype_bad
        ([TypeBase ExtShape Uniqueness], Result)
-> InternaliseM ([TypeBase ExtShape Uniqueness], Result)
forall (m :: * -> *) a. Monad m => a -> m a
return ([TypeBase ExtShape Uniqueness]
rettype', Result
body_res)
      Body
body' <-
        ErrorMsg SubExp -> SrcLoc -> [ExtType] -> Body -> InternaliseM Body
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') (Body -> InternaliseM Body) -> Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$
          Stms SOACS -> Result -> Body
forall lore. Bindable lore => Stms lore -> Result -> Body lore
mkBody Stms SOACS
body_stms Result
body_res

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

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

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

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

generateEntryPoint :: E.EntryPoint -> E.ValBind -> InternaliseM ()
generateEntryPoint :: EntryPoint -> ValBind -> InternaliseM ()
generateEntryPoint (E.EntryPoint [EntryType]
e_paramts EntryType
e_rettype) ValBind
vb = InternaliseM () -> InternaliseM ()
forall a. InternaliseM a -> InternaliseM a
localConstsScope (InternaliseM () -> InternaliseM ())
-> InternaliseM () -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ do
  let (E.ValBind Maybe (Info EntryPoint)
_ VName
ofname Maybe (TypeExp VName)
_ (Info (StructType
rettype, [VName]
_)) [TypeParamBase VName]
tparams [PatternBase Info VName]
params ExpBase Info VName
_ Maybe DocComment
_ [AttrInfo]
attrs SrcLoc
loc) = ValBind
vb
  [TypeParamBase VName]
-> [PatternBase Info VName]
-> ([FParam] -> [[FParam]] -> InternaliseM ())
-> InternaliseM ()
forall a.
[TypeParamBase VName]
-> [PatternBase Info VName]
-> ([FParam] -> [[FParam]] -> InternaliseM a)
-> InternaliseM a
bindingParams [TypeParamBase VName]
tparams [PatternBase Info VName]
params (([FParam] -> [[FParam]] -> InternaliseM ()) -> InternaliseM ())
-> ([FParam] -> [[FParam]] -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \[FParam]
shapeparams [[FParam]]
params' -> do
    [[TypeBase ExtShape Uniqueness]]
entry_rettype <- StructType -> InternaliseM [[TypeBase ExtShape Uniqueness]]
internaliseEntryReturnType StructType
rettype
    let entry' :: EntryPoint
entry' = [(EntryType, [FParam])]
-> (EntryType, [[TypeBase ExtShape Uniqueness]]) -> EntryPoint
entryPoint ([EntryType]
-> [[Param DeclType]] -> [(EntryType, [Param DeclType])]
forall a b. [a] -> [b] -> [(a, b)]
zip [EntryType]
e_paramts [[Param DeclType]]
[[FParam]]
params') (EntryType
e_rettype, [[TypeBase ExtShape Uniqueness]]
entry_rettype)
        args :: Result
args = (Param DeclType -> SubExp) -> [Param DeclType] -> Result
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] -> Result) -> [Param DeclType] -> Result
forall a b. (a -> b) -> a -> b
$ [[Param DeclType]] -> [Param DeclType]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Param DeclType]]
[[FParam]]
params'

    Body
entry_body <- InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
m (Body (Lore m)) -> m (Body (Lore m))
insertStmsM (InternaliseM (Body (Lore InternaliseM))
 -> InternaliseM (Body (Lore InternaliseM)))
-> InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ do
      -- Special case the (rare) situation where the entry point is
      -- not a function.
      Maybe Result
maybe_const <- VName -> InternaliseM (Maybe Result)
lookupConst VName
ofname
      Result
vals <- case Maybe Result
maybe_const of
        Just Result
ses ->
          Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
ses
        Maybe Result
Nothing ->
          (Result, [ExtType]) -> Result
forall a b. (a, b) -> a
fst ((Result, [ExtType]) -> Result)
-> InternaliseM (Result, [ExtType]) -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String
-> QualName VName
-> Result
-> SrcLoc
-> InternaliseM (Result, [ExtType])
funcall String
"entry_result" (VName -> QualName VName
forall v. v -> QualName v
E.qualName VName
ofname) Result
args SrcLoc
loc
      Result
ctx <-
        [TypeBase ExtShape Uniqueness] -> [Result] -> Result
forall u a. [TypeBase ExtShape u] -> [[a]] -> [a]
extractShapeContext ([[TypeBase ExtShape Uniqueness]] -> [TypeBase ExtShape Uniqueness]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[TypeBase ExtShape Uniqueness]]
entry_rettype)
          ([Result] -> Result)
-> InternaliseM [Result] -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> InternaliseM Result) -> Result -> InternaliseM [Result]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Type -> Result) -> InternaliseM Type -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Type -> Result
forall u. TypeBase Shape u -> Result
I.arrayDims (InternaliseM Type -> InternaliseM Result)
-> (SubExp -> InternaliseM Type) -> SubExp -> InternaliseM Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType) Result
vals
      Result -> InternaliseM (Body (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => Result -> m (Body (Lore m))
resultBodyM (Result
ctx Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ Result
vals)

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

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

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

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

internaliseBody :: String -> E.Exp -> InternaliseM Body
internaliseBody :: String -> ExpBase Info VName -> InternaliseM Body
internaliseBody String
desc ExpBase Info VName
e =
  InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
m (Body (Lore m)) -> m (Body (Lore m))
insertStmsM (InternaliseM (Body (Lore InternaliseM))
 -> InternaliseM (Body (Lore InternaliseM)))
-> InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody (Result -> Body) -> InternaliseM Result -> InternaliseM Body
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> ExpBase Info VName -> InternaliseM Result
internaliseExp (String
desc String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"_res") ExpBase Info VName
e

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

internaliseExp :: String -> E.Exp -> InternaliseM [I.SubExp]
internaliseExp :: String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc (E.Parens ExpBase Info VName
e SrcLoc
_) =
  String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc ExpBase Info VName
e
internaliseExp String
desc (E.QualParens (QualName VName, SrcLoc)
_ ExpBase Info VName
e SrcLoc
_) =
  String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc ExpBase Info VName
e
internaliseExp String
desc (E.StringLit [Word8]
vs SrcLoc
_) =
  (SubExp -> Result) -> InternaliseM SubExp -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> Result
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InternaliseM SubExp -> InternaliseM Result)
-> InternaliseM SubExp -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
    String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
desc (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ Result -> Type -> BasicOp
I.ArrayLit ((Word8 -> SubExp) -> [Word8] -> Result
forall a b. (a -> b) -> [a] -> [b]
map Word8 -> SubExp
forall v. IsValue v => v -> SubExp
constant [Word8]
vs) (Type -> BasicOp) -> Type -> BasicOp
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int8
internaliseExp String
_ (E.Var (E.QualName [VName]
_ VName
name) Info PatternType
_ SrcLoc
_) = do
  Maybe Result
subst <- VName -> InternaliseM (Maybe Result)
lookupSubst VName
name
  case Maybe Result
subst of
    Just Result
substs -> Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
substs
    Maybe Result
Nothing -> Result -> InternaliseM Result
forall (f :: * -> *) a. Applicative f => a -> f a
pure [VName -> SubExp
I.Var VName
name]
internaliseExp String
desc (E.Index ExpBase Info VName
e [DimIndexBase Info VName]
idxs (Info PatternType
ret, Info [VName]
retext) SrcLoc
loc) = do
  [VName]
vs <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"indexed" ExpBase Info VName
e
  Result
dims <- case [VName]
vs of
    [] -> Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return [] -- Will this happen?
    VName
v : [VName]
_ -> Type -> Result
forall u. TypeBase Shape u -> Result
I.arrayDims (Type -> Result) -> InternaliseM Type -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
v
  ([DimIndex SubExp]
idxs', Certificates
cs) <- SrcLoc
-> Result
-> [DimIndexBase Info VName]
-> InternaliseM ([DimIndex SubExp], Certificates)
internaliseSlice SrcLoc
loc Result
dims [DimIndexBase Info VName]
idxs
  let index :: VName -> InternaliseM (ExpT SOACS)
index VName
v = do
        Type
v_t <- VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
v
        ExpT SOACS -> InternaliseM (ExpT SOACS)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpT SOACS -> InternaliseM (ExpT SOACS))
-> ExpT SOACS -> InternaliseM (ExpT SOACS)
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ VName -> [DimIndex SubExp] -> BasicOp
I.Index VName
v ([DimIndex SubExp] -> BasicOp) -> [DimIndex SubExp] -> BasicOp
forall a b. (a -> b) -> a -> b
$ Type -> [DimIndex SubExp] -> [DimIndex SubExp]
fullSlice Type
v_t [DimIndex SubExp]
idxs'
  Result
ses <- Certificates -> InternaliseM Result -> InternaliseM Result
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
cs (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ String -> [Exp (Lore InternaliseM)] -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
String -> [Exp (Lore m)] -> m Result
letSubExps String
desc ([ExpT SOACS] -> InternaliseM Result)
-> InternaliseM [ExpT SOACS] -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (VName -> InternaliseM (ExpT SOACS))
-> [VName] -> InternaliseM [ExpT SOACS]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM (ExpT SOACS)
index [VName]
vs
  StructType -> [VName] -> Result -> InternaliseM ()
bindExtSizes (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret) [VName]
retext Result
ses
  Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
ses

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

    [Type]
rowtypes <-
      case (TypeBase ExtShape Uniqueness -> Maybe Type)
-> [TypeBase ExtShape Uniqueness] -> Maybe [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Type -> Type) -> Maybe Type -> Maybe Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
rowType (Maybe Type -> Maybe Type)
-> (TypeBase ExtShape Uniqueness -> Maybe Type)
-> TypeBase ExtShape Uniqueness
-> Maybe Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExtType -> Maybe Type
forall u. TypeBase ExtShape u -> Maybe (TypeBase Shape u)
hasStaticShape (ExtType -> Maybe Type)
-> (TypeBase ExtShape Uniqueness -> ExtType)
-> TypeBase ExtShape Uniqueness
-> Maybe Type
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 [Type]
ts -> [Type] -> InternaliseM [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Type]
ts
        Maybe [Type]
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 [Result]
es' of
            [] -> String -> InternaliseM [Type]
forall a. HasCallStack => String -> a
error (String -> InternaliseM [Type]) -> String -> InternaliseM [Type]
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp ArrayLit: existential type: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PatternType -> String
forall a. Pretty a => a -> String
pretty PatternType
arr_t
            Result
e' : [Result]
_ -> (SubExp -> InternaliseM Type) -> Result -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType Result
e'

    let arraylit :: Result -> Type -> InternaliseM (ExpT SOACS)
arraylit Result
ks Type
rt = do
          Result
ks' <-
            (SubExp -> InternaliseM SubExp) -> Result -> InternaliseM Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM
              ( ErrorMsg SubExp
-> SrcLoc -> Type -> String -> SubExp -> InternaliseM SubExp
ensureShape
                  ErrorMsg SubExp
"shape of element differs from shape of first element"
                  SrcLoc
loc
                  Type
rt
                  String
"elem_reshaped"
              )
              Result
ks
          ExpT SOACS -> InternaliseM (ExpT SOACS)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpT SOACS -> InternaliseM (ExpT SOACS))
-> ExpT SOACS -> InternaliseM (ExpT SOACS)
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ Result -> Type -> BasicOp
I.ArrayLit Result
ks' Type
rt

    String -> [Exp (Lore InternaliseM)] -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
String -> [Exp (Lore m)] -> m Result
letSubExps String
desc
      ([ExpT SOACS] -> InternaliseM Result)
-> InternaliseM [ExpT SOACS] -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< if [Result] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Result]
es'
        then (Type -> InternaliseM (ExpT SOACS))
-> [Type] -> InternaliseM [ExpT SOACS]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Result -> Type -> InternaliseM (ExpT SOACS)
arraylit []) [Type]
rowtypes
        else (Result -> Type -> InternaliseM (ExpT SOACS))
-> [Result] -> [Type] -> InternaliseM [ExpT SOACS]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM Result -> Type -> InternaliseM (ExpT SOACS)
arraylit ([Result] -> [Result]
forall a. [[a]] -> [[a]]
transpose [Result]
es') [Type]
rowtypes
  where
    isArrayLiteral :: E.Exp -> Maybe ([Int], [E.Exp])
    isArrayLiteral :: ExpBase Info VName -> Maybe ([Int], [ExpBase Info VName])
isArrayLiteral (E.ArrayLit [ExpBase Info VName]
inner_es Info PatternType
_ SrcLoc
_) = do
      ([Int]
eshape, [ExpBase Info VName]
e) : [([Int], [ExpBase Info VName])]
inner_es' <- (ExpBase Info VName -> Maybe ([Int], [ExpBase Info VName]))
-> [ExpBase Info VName] -> Maybe [([Int], [ExpBase Info VName])]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ExpBase Info VName -> Maybe ([Int], [ExpBase Info VName])
isArrayLiteral [ExpBase Info VName]
inner_es
      Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ (([Int], [ExpBase Info VName]) -> Bool)
-> [([Int], [ExpBase Info VName])] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (([Int]
eshape [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
==) ([Int] -> Bool)
-> (([Int], [ExpBase Info VName]) -> [Int])
-> ([Int], [ExpBase Info VName])
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Int], [ExpBase Info VName]) -> [Int]
forall a b. (a, b) -> a
fst) [([Int], [ExpBase Info VName])]
inner_es'
      ([Int], [ExpBase Info VName])
-> Maybe ([Int], [ExpBase Info VName])
forall (m :: * -> *) a. Monad m => a -> m a
return ([ExpBase Info VName] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ExpBase Info VName]
inner_es Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
eshape, [ExpBase Info VName]
e [ExpBase Info VName]
-> [ExpBase Info VName] -> [ExpBase Info VName]
forall a. [a] -> [a] -> [a]
++ (([Int], [ExpBase Info VName]) -> [ExpBase Info VName])
-> [([Int], [ExpBase Info VName])] -> [ExpBase Info VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Int], [ExpBase Info VName]) -> [ExpBase Info VName]
forall a b. (a, b) -> b
snd [([Int], [ExpBase Info VName])]
inner_es')
    isArrayLiteral ExpBase Info VName
e =
      ([Int], [ExpBase Info VName])
-> Maybe ([Int], [ExpBase Info VName])
forall a. a -> Maybe a
Just ([], [ExpBase Info VName
e])
internaliseExp String
desc (E.Range ExpBase Info VName
start Maybe (ExpBase Info VName)
maybe_second Inclusiveness (ExpBase Info VName)
end (Info PatternType
ret, Info [VName]
retext) SrcLoc
loc) = do
  SubExp
start' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"range_start" ExpBase Info VName
start
  SubExp
end' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"range_end" (ExpBase Info VName -> InternaliseM SubExp)
-> ExpBase Info VName -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ case Inclusiveness (ExpBase Info VName)
end of
    DownToExclusive ExpBase Info VName
e -> ExpBase Info VName
e
    ToInclusive ExpBase Info VName
e -> ExpBase Info VName
e
    UpToExclusive ExpBase Info VName
e -> ExpBase Info VName
e
  Maybe SubExp
maybe_second' <-
    (ExpBase Info VName -> InternaliseM SubExp)
-> Maybe (ExpBase Info VName) -> InternaliseM (Maybe SubExp)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"range_second") Maybe (ExpBase Info VName)
maybe_second

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

  (IntType
it, CmpOp
le_op, CmpOp
lt_op) <-
    case ExpBase Info VName -> PatternType
E.typeOf ExpBase Info VName
start of
      E.Scalar (E.Prim (E.Signed IntType
it)) -> (IntType, CmpOp, CmpOp) -> InternaliseM (IntType, CmpOp, CmpOp)
forall (m :: * -> *) a. Monad m => a -> m a
return (IntType
it, IntType -> CmpOp
CmpSle IntType
it, IntType -> CmpOp
CmpSlt IntType
it)
      E.Scalar (E.Prim (E.Unsigned IntType
it)) -> (IntType, CmpOp, CmpOp) -> InternaliseM (IntType, CmpOp, CmpOp)
forall (m :: * -> *) a. Monad m => a -> m a
return (IntType
it, IntType -> CmpOp
CmpUle IntType
it, IntType -> CmpOp
CmpUlt IntType
it)
      PatternType
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]
++ PatternType -> String
forall a. Pretty a => a -> String
pretty PatternType
start_t

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

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

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

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

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

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

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

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

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

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

  SubExp
se <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
desc (BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> SubExp -> SubExp -> IntType -> BasicOp
I.Iota SubExp
num_elems SubExp
start' SubExp
step IntType
it)
  StructType -> [VName] -> Result -> InternaliseM ()
bindExtSizes (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret) [VName]
retext [SubExp
se]
  Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return [SubExp
se]
internaliseExp String
desc (E.Ascript ExpBase Info VName
e TypeDeclBase Info VName
_ SrcLoc
_) =
  String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc ExpBase Info VName
e
internaliseExp String
desc (E.Coerce ExpBase Info VName
e (TypeDecl TypeExp VName
dt (Info StructType
et)) (Info PatternType
ret, Info [VName]
retext) SrcLoc
loc) = do
  Result
ses <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc ExpBase Info VName
e
  [TypeBase ExtShape Uniqueness]
ts <- StructType -> InternaliseM [TypeBase ExtShape Uniqueness]
internaliseReturnType StructType
et
  [ErrorMsgPart SubExp]
dt' <- TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
dt
  StructType -> [VName] -> Result -> InternaliseM ()
bindExtSizes (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret) [VName]
retext Result
ses
  [(SubExp, TypeBase ExtShape Uniqueness)]
-> ((SubExp, TypeBase ExtShape Uniqueness) -> InternaliseM SubExp)
-> InternaliseM Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (Result
-> [TypeBase ExtShape Uniqueness]
-> [(SubExp, TypeBase ExtShape Uniqueness)]
forall a b. [a] -> [b] -> [(a, b)]
zip Result
ses [TypeBase ExtShape Uniqueness]
ts) (((SubExp, TypeBase ExtShape Uniqueness) -> InternaliseM SubExp)
 -> InternaliseM Result)
-> ((SubExp, TypeBase ExtShape Uniqueness) -> InternaliseM SubExp)
-> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ \(SubExp
e', TypeBase ExtShape Uniqueness
t') -> do
    Result
dims <- Type -> Result
forall u. TypeBase Shape u -> Result
arrayDims (Type -> Result) -> InternaliseM Type -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
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) -> Result -> [ErrorMsgPart SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt64 Result
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'
internaliseExp String
desc (E.Negate ExpBase Info VName
e SrcLoc
_) = do
  SubExp
e' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"negate_arg" ExpBase Info VName
e
  Type
et <- SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
e'
  case Type
et of
    I.Prim (I.IntType IntType
t) ->
      String -> Exp (Lore InternaliseM) -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
t Overflow
I.OverflowWrap) (IntType -> Integer -> SubExp
I.intConst IntType
t Integer
0) SubExp
e'
    I.Prim (I.FloatType FloatType
t) ->
      String -> Exp (Lore InternaliseM) -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (FloatType -> BinOp
I.FSub FloatType
t) (FloatType -> Double -> SubExp
I.floatConst FloatType
t Double
0) SubExp
e'
    Type
_ -> String -> InternaliseM Result
forall a. HasCallStack => String -> a
error String
"Futhark.Internalise.internaliseExp: non-numeric type in Negate"
internaliseExp String
desc e :: ExpBase Info VName
e@E.Apply {} = do
  (QualName VName
qfname, [(ExpBase Info VName, Maybe VName)]
args, StructType
ret, [VName]
retext) <- ExpBase Info VName
-> InternaliseM
     (QualName VName, [(ExpBase Info VName, Maybe VName)], StructType,
      [VName])
findFuncall ExpBase Info VName
e
  -- Argument evaluation is outermost-in so that any existential sizes
  -- created by function applications can be brought into scope.
  let fname :: Name
fname = String -> Name
nameFromString (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ Name -> String
forall a. Pretty a => a -> String
pretty (Name -> String) -> Name -> String
forall a b. (a -> b) -> a -> b
$ VName -> Name
baseName (VName -> Name) -> VName -> Name
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname
      loc :: SrcLoc
loc = ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
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.
  Result
ses <-
    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 Result
internalise <- QualName VName
-> [ExpBase Info VName]
-> SrcLoc
-> Maybe (String -> InternaliseM Result)
isOverloadedFunction QualName VName
qfname (((ExpBase Info VName, Maybe VName) -> ExpBase Info VName)
-> [(ExpBase Info VName, Maybe VName)] -> [ExpBase Info VName]
forall a b. (a -> b) -> [a] -> [b]
map (ExpBase Info VName, Maybe VName) -> ExpBase Info VName
forall a b. (a, b) -> a
fst [(ExpBase Info VName, Maybe VName)]
args) SrcLoc
loc ->
          String -> InternaliseM Result
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]
          [Result]
args' <- [Result] -> [Result]
forall a. [a] -> [a]
reverse ([Result] -> [Result])
-> InternaliseM [Result] -> InternaliseM [Result]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((ExpBase Info VName, Maybe VName) -> InternaliseM Result)
-> [(ExpBase Info VName, Maybe VName)] -> InternaliseM [Result]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> (ExpBase Info VName, Maybe VName) -> InternaliseM Result
internaliseArg String
arg_desc) ([(ExpBase Info VName, Maybe VName)]
-> [(ExpBase Info VName, Maybe VName)]
forall a. [a] -> [a]
reverse [(ExpBase Info VName, Maybe VName)]
args)
          let args'' :: [(SubExp, Diet)]
args'' = (Result -> [(SubExp, Diet)]) -> [Result] -> [(SubExp, Diet)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Result -> [(SubExp, Diet)]
forall a. [a] -> [(a, Diet)]
tag [Result]
args'
          String -> Exp (Lore InternaliseM) -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
            Name
-> [(SubExp, Diet)]
-> [RetType SOACS]
-> (Safety, SrcLoc, [SrcLoc])
-> ExpT SOACS
forall lore.
Name
-> [(SubExp, Diet)]
-> [RetType lore]
-> (Safety, SrcLoc, [SrcLoc])
-> ExpT lore
I.Apply
              Name
fname
              [(SubExp, Diet)]
args''
              [PrimType -> TypeBase ExtShape Uniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
rettype]
              (Safety
Safe, SrcLoc
loc, [])
        | Bool
otherwise -> do
          Result
args' <- [Result] -> Result
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([Result] -> Result)
-> ([Result] -> [Result]) -> [Result] -> Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Result] -> [Result]
forall a. [a] -> [a]
reverse ([Result] -> Result)
-> InternaliseM [Result] -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((ExpBase Info VName, Maybe VName) -> InternaliseM Result)
-> [(ExpBase Info VName, Maybe VName)] -> InternaliseM [Result]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> (ExpBase Info VName, Maybe VName) -> InternaliseM Result
internaliseArg String
arg_desc) ([(ExpBase Info VName, Maybe VName)]
-> [(ExpBase Info VName, Maybe VName)]
forall a. [a] -> [a]
reverse [(ExpBase Info VName, Maybe VName)]
args)
          (Result, [ExtType]) -> Result
forall a b. (a, b) -> a
fst ((Result, [ExtType]) -> Result)
-> InternaliseM (Result, [ExtType]) -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String
-> QualName VName
-> Result
-> SrcLoc
-> InternaliseM (Result, [ExtType])
funcall String
desc QualName VName
qfname Result
args' SrcLoc
loc

  StructType -> [VName] -> Result -> InternaliseM ()
bindExtSizes StructType
ret [VName]
retext Result
ses
  Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
ses
internaliseExp String
desc (E.LetPat PatternBase Info VName
pat ExpBase Info VName
e ExpBase Info VName
body (Info PatternType
ret, Info [VName]
retext) SrcLoc
_) = do
  Result
ses <- String
-> PatternBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM Result)
-> InternaliseM Result
forall a.
String
-> PatternBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat String
desc PatternBase Info VName
pat ExpBase Info VName
e ExpBase Info VName
body (String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc)
  StructType -> [VName] -> Result -> InternaliseM ()
bindExtSizes (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret) [VName]
retext Result
ses
  Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
ses
internaliseExp String
_ (E.LetFun VName
ofname ([TypeParamBase VName], [PatternBase Info VName],
 Maybe (TypeExp VName), Info StructType, ExpBase Info VName)
_ ExpBase Info VName
_ Info PatternType
_ SrcLoc
_) =
  String -> InternaliseM Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
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
internaliseExp String
desc (E.DoLoop [VName]
sparams PatternBase Info VName
mergepat ExpBase Info VName
mergeexp LoopFormBase Info VName
form ExpBase Info VName
loopbody (Info (PatternType
ret, [VName]
retext)) SrcLoc
loc) = do
  Result
ses <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
"loop_init" ExpBase Info VName
mergeexp
  ((Body
loopbody', (LoopForm SOACS
form', [Param DeclType]
shapepat, [Param DeclType]
mergepat', Result
mergeinit')), Stms SOACS
initstms) <-
    InternaliseM
  (Body,
   (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
-> InternaliseM
     ((Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)),
      Stms (Lore InternaliseM))
forall (m :: * -> *) a.
MonadBinder m =>
m a -> m (a, Stms (Lore m))
collectStms (InternaliseM
   (Body,
    (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
 -> InternaliseM
      ((Body,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)),
       Stms (Lore InternaliseM)))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
-> InternaliseM
     ((Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)),
      Stms (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ Result
-> LoopFormBase Info VName
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
handleForm Result
ses LoopFormBase Info VName
form

  Stms (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *). MonadBinder m => Stms (Lore m) -> m ()
addStms Stms (Lore InternaliseM)
Stms SOACS
initstms
  [Type]
mergeinit_ts' <- (SubExp -> InternaliseM Type) -> Result -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType Result
mergeinit'

  Result
ctxinit <- [VName] -> [FParam] -> [Type] -> InternaliseM Result
forall (m :: * -> *).
(HasScope SOACS m, Monad m) =>
[VName] -> [FParam] -> [Type] -> m Result
argShapes ((Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
shapepat) [Param DeclType]
[FParam]
mergepat' [Type]
mergeinit_ts'

  let ctxmerge :: [(Param DeclType, SubExp)]
ctxmerge = [Param DeclType] -> Result -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
shapepat Result
ctxinit
      valmerge :: [(Param DeclType, SubExp)]
valmerge = [Param DeclType] -> Result -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
mergepat' Result
mergeinit'
      dropCond :: [VName] -> [VName]
dropCond = case LoopFormBase Info VName
form of
        E.While {} -> Int -> [VName] -> [VName]
forall a. Int -> [a] -> [a]
drop Int
1
        LoopFormBase Info VName
_ -> [VName] -> [VName]
forall a. a -> a
id

  -- Ensure that the result of the loop matches the shapes of the
  -- merge parameters.  XXX: Ideally they should already match (by
  -- the source language type rules), but some of our
  -- transformations (esp. defunctionalisation) strips out some size
  -- information.  For a type-correct source program, these reshapes
  -- should simplify away.
  let merge :: [(Param DeclType, SubExp)]
merge = [(Param DeclType, SubExp)]
ctxmerge [(Param DeclType, SubExp)]
-> [(Param DeclType, SubExp)] -> [(Param DeclType, SubExp)]
forall a. [a] -> [a] -> [a]
++ [(Param DeclType, SubExp)]
valmerge
      merge_ts :: [Type]
merge_ts = ((Param DeclType, SubExp) -> Type)
-> [(Param DeclType, SubExp)] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Param DeclType -> Type
forall dec. Typed dec => Param dec -> Type
I.paramType (Param DeclType -> Type)
-> ((Param DeclType, SubExp) -> Param DeclType)
-> (Param DeclType, SubExp)
-> Type
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
loopbody'' <-
    Scope SOACS -> InternaliseM Body -> InternaliseM Body
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param DeclType] -> Scope SOACS
forall lore dec.
(FParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
scopeOfFParams ([Param DeclType] -> Scope SOACS)
-> [Param DeclType] -> Scope SOACS
forall a b. (a -> b) -> a -> b
$ ((Param DeclType, SubExp) -> Param DeclType)
-> [(Param DeclType, SubExp)] -> [Param DeclType]
forall a b. (a -> b) -> [a] -> [b]
map (Param DeclType, SubExp) -> Param DeclType
forall a b. (a, b) -> a
fst [(Param DeclType, SubExp)]
merge) (InternaliseM Body -> InternaliseM Body)
-> InternaliseM Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$
      LoopForm SOACS -> InternaliseM Body -> InternaliseM Body
forall lore a (m :: * -> *) b.
(Scoped lore a, LocalScope lore m) =>
a -> m b -> m b
inScopeOf LoopForm SOACS
form' (InternaliseM Body -> InternaliseM Body)
-> InternaliseM Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$
        InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
m (Body (Lore m)) -> m (Body (Lore m))
insertStmsM (InternaliseM (Body (Lore InternaliseM))
 -> InternaliseM (Body (Lore InternaliseM)))
-> InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$
          Result -> InternaliseM Body
forall (m :: * -> *). MonadBinder m => Result -> m (Body (Lore m))
resultBodyM
            (Result -> InternaliseM Body)
-> InternaliseM Result -> InternaliseM Body
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ErrorMsg SubExp
-> SrcLoc -> [VName] -> [Type] -> Result -> InternaliseM Result
forall u.
Typed (TypeBase Shape u) =>
ErrorMsg SubExp
-> SrcLoc
-> [VName]
-> [TypeBase Shape u]
-> Result
-> InternaliseM Result
ensureArgShapes
              ErrorMsg SubExp
"shape of loop result does not match shapes in loop parameter"
              SrcLoc
loc
              (((Param DeclType, SubExp) -> VName)
-> [(Param DeclType, SubExp)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName (Param DeclType -> VName)
-> ((Param DeclType, SubExp) -> Param DeclType)
-> (Param DeclType, SubExp)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param DeclType, SubExp) -> Param DeclType
forall a b. (a, b) -> a
fst) [(Param DeclType, SubExp)]
ctxmerge)
              [Type]
merge_ts
            (Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Body (Lore InternaliseM) -> InternaliseM Result
forall (m :: * -> *). MonadBinder m => Body (Lore m) -> m Result
bodyBind Body (Lore InternaliseM)
Body
loopbody'

  Attrs
attrs <- (InternaliseEnv -> Attrs) -> InternaliseM Attrs
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks InternaliseEnv -> Attrs
envAttrs
  Result
loop_res <-
    (VName -> SubExp) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var ([VName] -> Result) -> ([VName] -> [VName]) -> [VName] -> Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [VName] -> [VName]
dropCond
      ([VName] -> Result) -> InternaliseM [VName] -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Attrs -> InternaliseM [VName] -> InternaliseM [VName]
forall (m :: * -> *) a. MonadBinder m => Attrs -> m a -> m a
attributing
        Attrs
attrs
        (String -> Exp (Lore InternaliseM) -> InternaliseM [VName]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [VName]
letTupExp String
desc ([(FParam, SubExp)]
-> [(FParam, SubExp)] -> LoopForm SOACS -> Body -> ExpT SOACS
forall lore.
[(FParam lore, SubExp)]
-> [(FParam lore, SubExp)]
-> LoopForm lore
-> BodyT lore
-> ExpT lore
I.DoLoop [(Param DeclType, SubExp)]
[(FParam, SubExp)]
ctxmerge [(Param DeclType, SubExp)]
[(FParam, SubExp)]
valmerge LoopForm SOACS
form' Body
loopbody''))
  StructType -> [VName] -> Result -> InternaliseM ()
bindExtSizes (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret) [VName]
retext Result
loop_res
  Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
loop_res
  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]
-> Result
-> LoopForm SOACS
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forLoop [Param DeclType]
mergepat' [Param DeclType]
shapepat Result
mergeinit LoopForm SOACS
form' =
      InternaliseM
  (Result,
   (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a. InternaliseM (Result, a) -> InternaliseM (Body, a)
bodyFromStms (InternaliseM
   (Result,
    (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a b. (a -> b) -> a -> b
$
        LoopForm SOACS
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall lore a (m :: * -> *) b.
(Scoped lore a, LocalScope lore m) =>
a -> m b -> m b
inScopeOf LoopForm SOACS
form' (InternaliseM
   (Result,
    (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
 -> InternaliseM
      (Result,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a b. (a -> b) -> a -> b
$ do
          Result
ses <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
"loopres" ExpBase Info VName
loopbody
          [Type]
sets <- (SubExp -> InternaliseM Type) -> Result -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType Result
ses
          Result
shapeargs <- [VName] -> [FParam] -> [Type] -> InternaliseM Result
forall (m :: * -> *).
(HasScope SOACS m, Monad m) =>
[VName] -> [FParam] -> [Type] -> m Result
argShapes ((Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
shapepat) [Param DeclType]
[FParam]
mergepat' [Type]
sets
          (Result,
 (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall (m :: * -> *) a. Monad m => a -> m a
return
            ( Result
shapeargs Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ Result
ses,
              ( LoopForm SOACS
form',
                [Param DeclType]
shapepat,
                [Param DeclType]
mergepat',
                Result
mergeinit
              )
            )

    handleForm :: Result
-> LoopFormBase Info VName
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
handleForm Result
mergeinit (E.ForIn PatternBase Info VName
x ExpBase Info VName
arr) = do
      [VName]
arr' <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"for_in_arr" ExpBase Info VName
arr
      [Type]
arr_ts <- (VName -> InternaliseM Type) -> [VName] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType [VName]
arr'
      let w :: SubExp
w = Int -> [Type] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [Type]
arr_ts

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

      [TypeParamBase VName]
-> PatternBase Info VName
-> ([FParam]
    -> [FParam]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a.
[TypeParamBase VName]
-> PatternBase Info VName
-> ([FParam] -> [FParam] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' PatternBase Info VName
mergepat (([FParam]
  -> [FParam]
  -> InternaliseM
       (Body,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> ([FParam]
    -> [FParam]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a b. (a -> b) -> a -> b
$
        \[FParam]
shapepat [FParam]
mergepat' ->
          [PatternBase Info VName]
-> [Type]
-> ([LParam]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a.
[PatternBase Info VName]
-> [Type] -> ([LParam] -> InternaliseM a) -> InternaliseM a
bindingLambdaParams [PatternBase Info VName
x] ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
rowType [Type]
arr_ts) (([LParam]
  -> InternaliseM
       (Body,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> ([LParam]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a b. (a -> b) -> a -> b
$ \[LParam]
x_params -> do
            let loopvars :: [(Param Type, VName)]
loopvars = [Param Type] -> [VName] -> [(Param Type, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param Type]
[LParam]
x_params [VName]
arr'
            [Param DeclType]
-> [Param DeclType]
-> Result
-> LoopForm SOACS
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forLoop [Param DeclType]
[FParam]
mergepat' [Param DeclType]
[FParam]
shapepat Result
mergeinit (LoopForm SOACS
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> LoopForm SOACS
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a b. (a -> b) -> a -> b
$
              VName -> IntType -> SubExp -> [(LParam, VName)] -> LoopForm SOACS
forall lore.
VName
-> IntType -> SubExp -> [(LParam lore, VName)] -> LoopForm lore
I.ForLoop VName
i IntType
Int64 SubExp
w [(Param Type, VName)]
[(LParam, VName)]
loopvars
    handleForm Result
mergeinit (E.For IdentBase Info VName
i ExpBase Info VName
num_iterations) = do
      SubExp
num_iterations' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"upper_bound" ExpBase Info VName
num_iterations
      Type
num_iterations_t <- SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
I.subExpType SubExp
num_iterations'
      IntType
it <- case Type
num_iterations_t of
        I.Prim (IntType IntType
it) -> IntType -> InternaliseM IntType
forall (m :: * -> *) a. Monad m => a -> m a
return IntType
it
        Type
_ -> String -> InternaliseM IntType
forall a. HasCallStack => String -> a
error String
"internaliseExp DoLoop: invalid type"

      [TypeParamBase VName]
-> PatternBase Info VName
-> ([FParam]
    -> [FParam]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a.
[TypeParamBase VName]
-> PatternBase Info VName
-> ([FParam] -> [FParam] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' PatternBase Info VName
mergepat (([FParam]
  -> [FParam]
  -> InternaliseM
       (Body,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> ([FParam]
    -> [FParam]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a b. (a -> b) -> a -> b
$
        \[FParam]
shapepat [FParam]
mergepat' ->
          [Param DeclType]
-> [Param DeclType]
-> Result
-> LoopForm SOACS
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forLoop [Param DeclType]
[FParam]
mergepat' [Param DeclType]
[FParam]
shapepat Result
mergeinit (LoopForm SOACS
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> LoopForm SOACS
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a b. (a -> b) -> a -> b
$
            VName -> IntType -> SubExp -> [(LParam, VName)] -> LoopForm SOACS
forall lore.
VName
-> IntType -> SubExp -> [(LParam lore, VName)] -> LoopForm lore
I.ForLoop (IdentBase Info VName -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
E.identName IdentBase Info VName
i) IntType
it SubExp
num_iterations' []
    handleForm Result
mergeinit (E.While ExpBase Info VName
cond) =
      [TypeParamBase VName]
-> PatternBase Info VName
-> ([FParam]
    -> [FParam]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a.
[TypeParamBase VName]
-> PatternBase Info VName
-> ([FParam] -> [FParam] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' PatternBase Info VName
mergepat (([FParam]
  -> [FParam]
  -> InternaliseM
       (Body,
        (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> ([FParam]
    -> [FParam]
    -> InternaliseM
         (Body,
          (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a b. (a -> b) -> a -> b
$ \[FParam]
shapepat [FParam]
mergepat' -> do
        [Type]
mergeinit_ts <- (SubExp -> InternaliseM Type) -> Result -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType Result
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.
        Result
shapeinit <- [VName] -> [FParam] -> [Type] -> InternaliseM Result
forall (m :: * -> *).
(HasScope SOACS m, Monad m) =>
[VName] -> [FParam] -> [Type] -> m Result
argShapes ((Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
[FParam]
shapepat) [FParam]
mergepat' [Type]
mergeinit_ts

        (SubExp
loop_initial_cond, Stms SOACS
init_loop_cond_bnds) <- InternaliseM SubExp
-> InternaliseM (SubExp, Stms (Lore InternaliseM))
forall (m :: * -> *) a.
MonadBinder m =>
m a -> m (a, Stms (Lore m))
collectStms (InternaliseM SubExp
 -> InternaliseM (SubExp, Stms (Lore InternaliseM)))
-> InternaliseM SubExp
-> InternaliseM (SubExp, Stms (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ do
          [(Param DeclType, SubExp)]
-> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param DeclType] -> Result -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
[FParam]
shapepat Result
shapeinit) (((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ())
-> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \(Param DeclType
p, SubExp
se) ->
            [VName] -> Exp (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
          [(Param DeclType, SubExp)]
-> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param DeclType] -> Result -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
[FParam]
mergepat' Result
mergeinit) (((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ())
-> ((Param DeclType, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \(Param DeclType
p, SubExp
se) ->
            Bool -> InternaliseM () -> InternaliseM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (SubExp
se SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== VName -> SubExp
I.Var (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p)) (InternaliseM () -> InternaliseM ())
-> InternaliseM () -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
              [VName] -> Exp (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
                BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
                  case SubExp
se of
                    I.Var VName
v
                      | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Type -> Bool
forall shape u. TypeBase shape u -> Bool
primType (Type -> Bool) -> Type -> Bool
forall a b. (a -> b) -> a -> b
$ Param DeclType -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param DeclType
p ->
                        ShapeChange SubExp -> VName -> BasicOp
Reshape ((SubExp -> DimChange SubExp) -> Result -> ShapeChange SubExp
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimCoercion (Result -> ShapeChange SubExp) -> Result -> ShapeChange SubExp
forall a b. (a -> b) -> a -> b
$ Type -> Result
forall u. TypeBase Shape u -> Result
arrayDims (Type -> Result) -> Type -> Result
forall a b. (a -> b) -> a -> b
$ Param DeclType -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param DeclType
p) VName
v
                    SubExp
_ -> SubExp -> BasicOp
SubExp SubExp
se
          String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"loop_cond" ExpBase Info VName
cond

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

        InternaliseM
  (Result,
   (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a. InternaliseM (Result, a) -> InternaliseM (Body, a)
bodyFromStms (InternaliseM
   (Result,
    (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
 -> InternaliseM
      (Body,
       (LoopForm SOACS, [Param DeclType], [Param DeclType], Result)))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
-> InternaliseM
     (Body,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall a b. (a -> b) -> a -> b
$ do
          Result
ses <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
"loopres" ExpBase Info VName
loopbody
          [Type]
sets <- (SubExp -> InternaliseM Type) -> Result -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType Result
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
          Result
shapeargs <- [VName] -> [FParam] -> [Type] -> InternaliseM Result
forall (m :: * -> *).
(HasScope SOACS m, Monad m) =>
[VName] -> [FParam] -> [Type] -> m Result
argShapes ((Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
[FParam]
shapepat) [FParam]
mergepat' [Type]
sets

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

          (Result,
 (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
-> InternaliseM
     (Result,
      (LoopForm SOACS, [Param DeclType], [Param DeclType], Result))
forall (m :: * -> *) a. Monad m => a -> m a
return
            ( Result
shapeargs Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ Result
loop_end_cond Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ Result
ses,
              ( VName -> LoopForm SOACS
forall lore. VName -> LoopForm lore
I.WhileLoop (VName -> LoopForm SOACS) -> VName -> LoopForm SOACS
forall a b. (a -> b) -> a -> b
$ Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName Param DeclType
loop_while,
                [Param DeclType]
[FParam]
shapepat,
                Param DeclType
loop_while Param DeclType -> [Param DeclType] -> [Param DeclType]
forall a. a -> [a] -> [a]
: [Param DeclType]
[FParam]
mergepat',
                SubExp
loop_initial_cond SubExp -> Result -> Result
forall a. a -> [a] -> [a]
: Result
mergeinit
              )
            )
internaliseExp String
desc (E.LetWith IdentBase Info VName
name IdentBase Info VName
src [DimIndexBase Info VName]
idxs ExpBase Info VName
ve ExpBase Info VName
body Info PatternType
t SrcLoc
loc) = do
  let pat :: PatternBase Info VName
pat = VName -> Info PatternType -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase f vn
E.Id (IdentBase Info VName -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
E.identName IdentBase Info VName
name) (IdentBase Info VName -> Info PatternType
forall (f :: * -> *) vn. IdentBase f vn -> f PatternType
E.identType IdentBase Info VName
name) SrcLoc
loc
      src_t :: Info PatternType
src_t = PatternType -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
E.fromStruct (PatternType -> PatternType)
-> Info PatternType -> Info PatternType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IdentBase Info VName -> Info PatternType
forall (f :: * -> *) vn. IdentBase f vn -> f PatternType
E.identType IdentBase Info VName
src
      e :: ExpBase Info VName
e = ExpBase Info VName
-> [DimIndexBase Info VName]
-> ExpBase Info VName
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> [DimIndexBase f vn] -> ExpBase f vn -> SrcLoc -> ExpBase f vn
E.Update (QualName VName -> Info PatternType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
E.Var (VName -> QualName VName
forall v. v -> QualName v
E.qualName (VName -> QualName VName) -> VName -> QualName VName
forall a b. (a -> b) -> a -> b
$ IdentBase Info VName -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
E.identName IdentBase Info VName
src) Info PatternType
src_t SrcLoc
loc) [DimIndexBase Info VName]
idxs ExpBase Info VName
ve SrcLoc
loc
  String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc (ExpBase Info VName -> InternaliseM Result)
-> ExpBase Info VName -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ PatternBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
PatternBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
E.LetPat PatternBase Info VName
pat ExpBase Info VName
e ExpBase Info VName
body (Info PatternType
t, [VName] -> Info [VName]
forall a. a -> Info a
Info []) SrcLoc
loc
internaliseExp String
desc (E.Update ExpBase Info VName
src [DimIndexBase Info VName]
slice ExpBase Info VName
ve SrcLoc
loc) = do
  Result
ves <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
"lw_val" ExpBase Info VName
ve
  [VName]
srcs <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"src" ExpBase Info VName
src
  Result
dims <- case [VName]
srcs of
    [] -> Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return [] -- Will this happen?
    VName
v : [VName]
_ -> Type -> Result
forall u. TypeBase Shape u -> Result
I.arrayDims (Type -> Result) -> InternaliseM Type -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
v
  ([DimIndex SubExp]
idxs', Certificates
cs) <- SrcLoc
-> Result
-> [DimIndexBase Info VName]
-> InternaliseM ([DimIndex SubExp], Certificates)
internaliseSlice SrcLoc
loc Result
dims [DimIndexBase Info VName]
slice

  let comb :: VName -> SubExp -> InternaliseM VName
comb VName
sname SubExp
ve' = do
        Type
sname_t <- VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
sname
        let full_slice :: [DimIndex SubExp]
full_slice = Type -> [DimIndex SubExp] -> [DimIndex SubExp]
fullSlice Type
sname_t [DimIndex SubExp]
idxs'
            rowtype :: Type
rowtype = Type
sname_t Type -> Result -> Type
forall oldshape u.
TypeBase oldshape u -> Result -> TypeBase Shape u
`setArrayDims` [DimIndex SubExp] -> Result
forall d. Slice d -> [d]
sliceDims [DimIndex SubExp]
full_slice
        SubExp
ve'' <-
          ErrorMsg SubExp
-> SrcLoc -> Type -> String -> SubExp -> InternaliseM SubExp
ensureShape
            ErrorMsg SubExp
"shape of value does not match shape of source array"
            SrcLoc
loc
            Type
rowtype
            String
"lw_val_correct_shape"
            SubExp
ve'
        String
-> VName
-> [DimIndex SubExp]
-> Exp (Lore InternaliseM)
-> InternaliseM VName
forall (m :: * -> *).
MonadBinder m =>
String -> VName -> [DimIndex SubExp] -> Exp (Lore m) -> m VName
letInPlace String
desc VName
sname [DimIndex SubExp]
full_slice (Exp (Lore InternaliseM) -> InternaliseM VName)
-> Exp (Lore InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
ve''
  Certificates -> InternaliseM Result -> InternaliseM Result
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
cs (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var ([VName] -> Result) -> InternaliseM [VName] -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> SubExp -> InternaliseM VName)
-> [VName] -> Result -> 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 Result
ves
internaliseExp String
desc (E.RecordUpdate ExpBase Info VName
src [Name]
fields ExpBase Info VName
ve Info PatternType
_ SrcLoc
_) = do
  Result
src' <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc ExpBase Info VName
src
  Result
ve' <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc ExpBase Info VName
ve
  StructType -> [Name] -> Result -> Result -> InternaliseM Result
forall a. StructType -> [Name] -> [a] -> [a] -> InternaliseM [a]
replace (ExpBase Info VName -> PatternType
E.typeOf ExpBase Info VName
src PatternType -> () -> StructType
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` ()) [Name]
fields Result
ve' Result
src'
  where
    replace :: StructType -> [Name] -> [a] -> [a] -> InternaliseM [a]
replace (E.Scalar (E.Record Map Name StructType
m)) (Name
f : [Name]
fs) [a]
ve' [a]
src'
      | Just StructType
t <- Name -> Map Name StructType -> Maybe StructType
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f Map Name StructType
m = do
        Int
i <-
          ([Int] -> Int) -> InternaliseM [Int] -> InternaliseM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (InternaliseM [Int] -> InternaliseM Int)
-> InternaliseM [Int] -> InternaliseM Int
forall a b. (a -> b) -> a -> b
$
            ((Name, StructType) -> InternaliseM Int)
-> [(Name, StructType)] -> InternaliseM [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (StructType -> InternaliseM Int
internalisedTypeSize (StructType -> InternaliseM Int)
-> ((Name, StructType) -> StructType)
-> (Name, StructType)
-> InternaliseM Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, StructType) -> StructType
forall a b. (a, b) -> b
snd) ([(Name, StructType)] -> InternaliseM [Int])
-> [(Name, StructType)] -> InternaliseM [Int]
forall a b. (a -> b) -> a -> b
$
              ((Name, StructType) -> Bool)
-> [(Name, StructType)] -> [(Name, StructType)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
/= Name
f) (Name -> Bool)
-> ((Name, StructType) -> Name) -> (Name, StructType) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, StructType) -> Name
forall a b. (a, b) -> a
fst) ([(Name, StructType)] -> [(Name, StructType)])
-> [(Name, StructType)] -> [(Name, StructType)]
forall a b. (a -> b) -> a -> b
$ Map Name StructType -> [(Name, StructType)]
forall a. Map Name a -> [(Name, a)]
sortFields Map Name StructType
m
        Int
k <- StructType -> InternaliseM Int
internalisedTypeSize StructType
t
        let ([a]
bef, [a]
to_update, [a]
aft) = Int -> Int -> [a] -> ([a], [a], [a])
forall a. Int -> Int -> [a] -> ([a], [a], [a])
splitAt3 Int
i Int
k [a]
src'
        [a]
src'' <- StructType -> [Name] -> [a] -> [a] -> InternaliseM [a]
replace StructType
t [Name]
fs [a]
ve' [a]
to_update
        [a] -> InternaliseM [a]
forall (m :: * -> *) a. Monad m => a -> m a
return ([a] -> InternaliseM [a]) -> [a] -> InternaliseM [a]
forall a b. (a -> b) -> a -> b
$ [a]
bef [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
src'' [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
aft
    replace StructType
_ [Name]
_ [a]
ve' [a]
_ = [a] -> InternaliseM [a]
forall (m :: * -> *) a. Monad m => a -> m a
return [a]
ve'
internaliseExp String
desc (E.Attr AttrInfo
attr ExpBase Info VName
e SrcLoc
_) =
  (InternaliseEnv -> InternaliseEnv)
-> InternaliseM Result -> InternaliseM Result
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local InternaliseEnv -> InternaliseEnv
f (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc ExpBase Info VName
e
  where
    attrs :: Attrs
attrs = Attr -> Attrs
oneAttr (Attr -> Attrs) -> Attr -> Attrs
forall a b. (a -> b) -> a -> b
$ AttrInfo -> Attr
internaliseAttr AttrInfo
attr
    f :: InternaliseEnv -> InternaliseEnv
f InternaliseEnv
env
      | Attr
"unsafe" Attr -> Attrs -> Bool
`inAttrs` Attrs
attrs,
        Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ InternaliseEnv -> Bool
envSafe InternaliseEnv
env =
        InternaliseEnv
env {envDoBoundsChecks :: Bool
envDoBoundsChecks = Bool
False}
      | Bool
otherwise =
        InternaliseEnv
env {envAttrs :: Attrs
envAttrs = InternaliseEnv -> Attrs
envAttrs InternaliseEnv
env Attrs -> Attrs -> Attrs
forall a. Semigroup a => a -> a -> a
<> Attrs
attrs}
internaliseExp String
desc (E.Assert ExpBase Info VName
e1 ExpBase Info VName
e2 (Info String
check) SrcLoc
loc) = do
  SubExp
e1' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"assert_cond" ExpBase Info VName
e1
  Certificates
c <- String
-> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certificates
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.
  Certificates -> InternaliseM Result -> InternaliseM Result
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying Certificates
c (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ (SubExp -> InternaliseM SubExp) -> Result -> InternaliseM Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM SubExp
forall (m :: * -> *). MonadBinder m => SubExp -> m SubExp
rebind (Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc ExpBase Info VName
e2
  where
    rebind :: SubExp -> m SubExp
rebind SubExp
v = do
      VName
v' <- String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"assert_res"
      [VName] -> Exp (Lore m) -> m ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [VName
v'] (Exp (Lore m) -> m ()) -> Exp (Lore m) -> m ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
v
      SubExp -> m SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> m SubExp) -> SubExp -> m SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
v'
internaliseExp String
_ (E.Constr Name
c [ExpBase Info VName]
es (Info (E.Scalar (E.Sum Map Name [PatternType]
fs))) SrcLoc
_) = do
  ([TypeBase ExtShape Uniqueness]
ts, Map Name (Int, [Int])
constr_map) <- Map Name [StructType]
-> InternaliseM
     ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
internaliseSumType (Map Name [StructType]
 -> InternaliseM
      ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int])))
-> Map Name [StructType]
-> InternaliseM
     ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
forall a b. (a -> b) -> a -> b
$ ([PatternType] -> [StructType])
-> Map Name [PatternType] -> Map Name [StructType]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((PatternType -> StructType) -> [PatternType] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct) Map Name [PatternType]
fs
  Result
es' <- [Result] -> Result
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([Result] -> Result)
-> InternaliseM [Result] -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ExpBase Info VName -> InternaliseM Result)
-> [ExpBase Info VName] -> InternaliseM [Result]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
"payload") [ExpBase Info VName]
es

  let noExt :: p -> m SubExp
noExt p
_ = SubExp -> m SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> m SubExp) -> SubExp -> m SubExp
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0
  [Type]
ts' <- (Int -> InternaliseM SubExp) -> [ExtType] -> InternaliseM [Type]
forall (m :: * -> *) u.
Monad m =>
(Int -> m SubExp) -> [TypeBase ExtShape u] -> m [TypeBase Shape u]
instantiateShapes Int -> InternaliseM SubExp
forall (m :: * -> *) p. Monad m => p -> m SubExp
noExt ([ExtType] -> InternaliseM [Type])
-> [ExtType] -> InternaliseM [Type]
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 -> Result -> Result
forall a. a -> [a] -> [a]
:) (Result -> Result) -> InternaliseM Result -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [Type] -> [(Int, SubExp)] -> InternaliseM Result
forall (f :: * -> *) a.
(Num a, MonadBinder f, Eq a) =>
a -> [Type] -> [(a, SubExp)] -> f Result
clauses Int
0 [Type]
ts' ([Int] -> Result -> [(Int, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
js Result
es')
    Maybe (Int, [Int])
Nothing ->
      String -> InternaliseM Result
forall a. HasCallStack => String -> a
error String
"internaliseExp Constr: missing constructor"
  where
    clauses :: a -> [Type] -> [(a, SubExp)] -> f Result
clauses a
j (Type
t : [Type]
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 -> Result -> Result
forall a. a -> [a] -> [a]
:) (Result -> Result) -> f Result -> f Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> [Type] -> [(a, SubExp)] -> f Result
clauses (a
j a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) [Type]
ts [(a, SubExp)]
js_to_es
      | Bool
otherwise = do
        SubExp
blank <- String -> Exp (Lore f) -> f SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"zero" (Exp (Lore f) -> f SubExp) -> f (Exp (Lore f)) -> f SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Type -> f (Exp (Lore f))
forall (m :: * -> *). MonadBinder m => Type -> m (Exp (Lore m))
eBlank Type
t
        (SubExp
blank SubExp -> Result -> Result
forall a. a -> [a] -> [a]
:) (Result -> Result) -> f Result -> f Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> [Type] -> [(a, SubExp)] -> f Result
clauses (a
j a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) [Type]
ts [(a, SubExp)]
js_to_es
    clauses a
_ [] [(a, SubExp)]
_ =
      Result -> f Result
forall (m :: * -> *) a. Monad m => a -> m a
return []
internaliseExp String
_ (E.Constr Name
_ [ExpBase Info VName]
_ (Info PatternType
t) SrcLoc
loc) =
  String -> InternaliseM Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp: constructor with type " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PatternType -> String
forall a. Pretty a => a -> String
pretty PatternType
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
internaliseExp String
desc (E.Match ExpBase Info VName
e NonEmpty (CaseBase Info VName)
cs (Info PatternType
ret, Info [VName]
retext) SrcLoc
_) = do
  Result
ses <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp (String
desc String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_scrutinee") ExpBase Info VName
e
  Result
res <-
    case NonEmpty (CaseBase Info VName)
-> (CaseBase Info VName, Maybe (NonEmpty (CaseBase Info VName)))
forall a. NonEmpty a -> (a, Maybe (NonEmpty a))
NE.uncons NonEmpty (CaseBase Info VName)
cs of
      (CasePat PatternBase Info VName
pCase ExpBase Info VName
eCase SrcLoc
_, Maybe (NonEmpty (CaseBase Info VName))
Nothing) -> do
        (SubExp
_, Result
pertinent) <- PatternBase Info VName -> Result -> InternaliseM (SubExp, Result)
generateCond PatternBase Info VName
pCase Result
ses
        PatternBase Info VName
-> Result
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM Result)
-> InternaliseM Result
forall a.
PatternBase Info VName
-> Result
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat' PatternBase Info VName
pCase Result
pertinent ExpBase Info VName
eCase (String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc)
      (CaseBase Info VName
c, Just NonEmpty (CaseBase Info VName)
cs') -> do
        let CasePat PatternBase Info VName
pLast ExpBase Info VName
eLast SrcLoc
_ = NonEmpty (CaseBase Info VName) -> CaseBase Info VName
forall a. NonEmpty a -> a
NE.last NonEmpty (CaseBase Info VName)
cs'
        Body
bFalse <- do
          (SubExp
_, Result
pertinent) <- PatternBase Info VName -> Result -> InternaliseM (SubExp, Result)
generateCond PatternBase Info VName
pLast Result
ses
          Body
eLast' <- PatternBase Info VName
-> Result
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM Body)
-> InternaliseM Body
forall a.
PatternBase Info VName
-> Result
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat' PatternBase Info VName
pLast Result
pertinent ExpBase Info VName
eLast (String -> ExpBase Info VName -> InternaliseM Body
internaliseBody String
desc)
          (Body -> CaseBase Info VName -> InternaliseM Body)
-> Body -> [CaseBase Info VName] -> InternaliseM Body
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\Body
bf CaseBase Info VName
c' -> [InternaliseM (Exp (Lore InternaliseM))]
-> InternaliseM (Body (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
[m (Exp (Lore m))] -> m (Body (Lore m))
eBody ([InternaliseM (Exp (Lore InternaliseM))]
 -> InternaliseM (Body (Lore InternaliseM)))
-> [InternaliseM (Exp (Lore InternaliseM))]
-> InternaliseM (Body (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ InternaliseM (ExpT SOACS) -> [InternaliseM (ExpT SOACS)]
forall (m :: * -> *) a. Monad m => a -> m a
return (InternaliseM (ExpT SOACS) -> [InternaliseM (ExpT SOACS)])
-> InternaliseM (ExpT SOACS) -> [InternaliseM (ExpT SOACS)]
forall a b. (a -> b) -> a -> b
$ Result -> CaseBase Info VName -> Body -> InternaliseM (ExpT SOACS)
generateCaseIf Result
ses CaseBase Info VName
c' Body
bf) Body
eLast' ([CaseBase Info VName] -> InternaliseM Body)
-> [CaseBase Info VName] -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$
            [CaseBase Info VName] -> [CaseBase Info VName]
forall a. [a] -> [a]
reverse ([CaseBase Info VName] -> [CaseBase Info VName])
-> [CaseBase Info VName] -> [CaseBase Info VName]
forall a b. (a -> b) -> a -> b
$ NonEmpty (CaseBase Info VName) -> [CaseBase Info VName]
forall a. NonEmpty a -> [a]
NE.init NonEmpty (CaseBase Info VName)
cs'
        String -> Exp (Lore InternaliseM) -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (ExpT SOACS -> InternaliseM Result)
-> InternaliseM (ExpT SOACS) -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Result -> CaseBase Info VName -> Body -> InternaliseM (ExpT SOACS)
generateCaseIf Result
ses CaseBase Info VName
c Body
bFalse
  StructType -> [VName] -> Result -> InternaliseM ()
bindExtSizes (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret) [VName]
retext Result
res
  Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
res

-- The "interesting" cases are over, now it's mostly boilerplate.

internaliseExp String
_ (E.Literal PrimValue
v SrcLoc
_) =
  Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return [PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimValue
internalisePrimValue PrimValue
v]
internaliseExp String
_ (E.IntLit Integer
v (Info PatternType
t) SrcLoc
_) =
  case PatternType
t of
    E.Scalar (E.Prim (E.Signed IntType
it)) ->
      Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return [PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
I.IntValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> IntValue
forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
v]
    E.Scalar (E.Prim (E.Unsigned IntType
it)) ->
      Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return [PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
I.IntValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> IntValue
forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
v]
    E.Scalar (E.Prim (E.FloatType FloatType
ft)) ->
      Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return [PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
I.FloatValue (FloatValue -> PrimValue) -> FloatValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ FloatType -> Integer -> FloatValue
forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
ft Integer
v]
    PatternType
_ -> String -> InternaliseM Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp: nonsensical type for integer literal: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PatternType -> String
forall a. Pretty a => a -> String
pretty PatternType
t
internaliseExp String
_ (E.FloatLit Double
v (Info PatternType
t) SrcLoc
_) =
  case PatternType
t of
    E.Scalar (E.Prim (E.FloatType FloatType
ft)) ->
      Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return [PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
I.FloatValue (FloatValue -> PrimValue) -> FloatValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ FloatType -> Double -> FloatValue
forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
ft Double
v]
    PatternType
_ -> String -> InternaliseM Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp: nonsensical type for float literal: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PatternType -> String
forall a. Pretty a => a -> String
pretty PatternType
t
internaliseExp String
desc (E.If ExpBase Info VName
ce ExpBase Info VName
te ExpBase Info VName
fe (Info PatternType
ret, Info [VName]
retext) SrcLoc
_) = do
  Result
ses <-
    String -> Exp (Lore InternaliseM) -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc
      (ExpT SOACS -> InternaliseM Result)
-> InternaliseM (ExpT SOACS) -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
(MonadBinder m, BranchType (Lore m) ~ ExtType) =>
m (Exp (Lore m))
-> m (Body (Lore m)) -> m (Body (Lore m)) -> m (Exp (Lore m))
eIf
        (BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS)
-> (SubExp -> BasicOp) -> SubExp -> ExpT SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
SubExp (SubExp -> ExpT SOACS)
-> InternaliseM SubExp -> InternaliseM (ExpT SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"cond" ExpBase Info VName
ce)
        (String -> ExpBase Info VName -> InternaliseM Body
internaliseBody (String
desc String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"_t") ExpBase Info VName
te)
        (String -> ExpBase Info VName -> InternaliseM Body
internaliseBody (String
desc String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"_f") ExpBase Info VName
fe)
  StructType -> [VName] -> Result -> InternaliseM ()
bindExtSizes (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret) [VName]
retext Result
ses
  Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
ses

-- Builtin operators are handled specially because they are
-- overloaded.
internaliseExp String
desc (E.Project Name
k ExpBase Info VName
e (Info PatternType
rt) SrcLoc
_) = do
  Int
n <- StructType -> InternaliseM Int
internalisedTypeSize (StructType -> InternaliseM Int) -> StructType -> InternaliseM Int
forall a b. (a -> b) -> a -> b
$ PatternType
rt PatternType -> () -> StructType
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` ()
  Int
i' <- ([Int] -> Int) -> InternaliseM [Int] -> InternaliseM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (InternaliseM [Int] -> InternaliseM Int)
-> InternaliseM [Int] -> InternaliseM Int
forall a b. (a -> b) -> a -> b
$
    (StructType -> InternaliseM Int)
-> [StructType] -> InternaliseM [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM StructType -> InternaliseM Int
internalisedTypeSize ([StructType] -> InternaliseM [Int])
-> [StructType] -> InternaliseM [Int]
forall a b. (a -> b) -> a -> b
$
      case ExpBase Info VName -> PatternType
E.typeOf ExpBase Info VName
e PatternType -> () -> StructType
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` () of
        E.Scalar (Record Map Name StructType
fs) ->
          ((Name, StructType) -> StructType)
-> [(Name, StructType)] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map (Name, StructType) -> StructType
forall a b. (a, b) -> b
snd ([(Name, StructType)] -> [StructType])
-> [(Name, StructType)] -> [StructType]
forall a b. (a -> b) -> a -> b
$ ((Name, StructType) -> Bool)
-> [(Name, StructType)] -> [(Name, StructType)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
/= Name
k) (Name -> Bool)
-> ((Name, StructType) -> Name) -> (Name, StructType) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, StructType) -> Name
forall a b. (a, b) -> a
fst) ([(Name, StructType)] -> [(Name, StructType)])
-> [(Name, StructType)] -> [(Name, StructType)]
forall a b. (a -> b) -> a -> b
$ Map Name StructType -> [(Name, StructType)]
forall a. Map Name a -> [(Name, a)]
sortFields Map Name StructType
fs
        StructType
t -> [StructType
t]
  Int -> Result -> Result
forall a. Int -> [a] -> [a]
take Int
n (Result -> Result) -> (Result -> Result) -> Result -> Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Result -> Result
forall a. Int -> [a] -> [a]
drop Int
i' (Result -> Result) -> InternaliseM Result -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc ExpBase Info VName
e
internaliseExp String
_ e :: ExpBase Info VName
e@E.BinOp {} =
  String -> InternaliseM Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp: Unexpected BinOp " String -> String -> String
forall a. [a] -> [a] -> [a]
++ ExpBase Info VName -> String
forall a. Pretty a => a -> String
pretty ExpBase Info VName
e
internaliseExp String
_ e :: ExpBase Info VName
e@E.Lambda {} =
  String -> InternaliseM Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp: Unexpected lambda at " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SrcLoc -> String
forall a. Located a => a -> String
locStr (ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
e)
internaliseExp String
_ e :: ExpBase Info VName
e@E.OpSection {} =
  String -> InternaliseM Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp: Unexpected operator section at " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SrcLoc -> String
forall a. Located a => a -> String
locStr (ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
e)
internaliseExp String
_ e :: ExpBase Info VName
e@E.OpSectionLeft {} =
  String -> InternaliseM Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp: Unexpected left operator section at " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SrcLoc -> String
forall a. Located a => a -> String
locStr (ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
e)
internaliseExp String
_ e :: ExpBase Info VName
e@E.OpSectionRight {} =
  String -> InternaliseM Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp: Unexpected right operator section at " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SrcLoc -> String
forall a. Located a => a -> String
locStr (ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
e)
internaliseExp String
_ e :: ExpBase Info VName
e@E.ProjectSection {} =
  String -> InternaliseM Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp: Unexpected projection section at " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SrcLoc -> String
forall a. Located a => a -> String
locStr (ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
e)
internaliseExp String
_ e :: ExpBase Info VName
e@E.IndexSection {} =
  String -> InternaliseM Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ String
"internaliseExp: Unexpected index section at " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SrcLoc -> String
forall a. Located a => a -> String
locStr (ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
e)

internaliseArg :: String -> (E.Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg :: String -> (ExpBase Info VName, Maybe VName) -> InternaliseM Result
internaliseArg String
desc (ExpBase Info VName
arg, Maybe VName
argdim) = do
  Result
arg' <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc ExpBase Info VName
arg
  case (Result
arg', Maybe VName
argdim) of
    ([SubExp
se], Just VName
d) -> [VName] -> Exp (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [VName
d] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
    (Result, Maybe VName)
_ -> () -> InternaliseM ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
  Result -> InternaliseM Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
arg'

generateCond :: E.Pattern -> [I.SubExp] -> InternaliseM (I.SubExp, [I.SubExp])
generateCond :: PatternBase Info VName -> Result -> InternaliseM (SubExp, Result)
generateCond PatternBase Info VName
orig_p Result
orig_ses = do
  (Result
cmps, Result
pertinent, Result
_) <- PatternBase Info VName
-> Result -> InternaliseM (Result, Result, Result)
forall vn.
(Eq vn, IsName vn) =>
PatternBase Info vn
-> Result -> InternaliseM (Result, Result, Result)
compares PatternBase Info VName
orig_p Result
orig_ses
  SubExp
cmp <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"matches" (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Result -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => Result -> m (Exp (Lore m))
eAll Result
cmps
  (SubExp, Result) -> InternaliseM (SubExp, Result)
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
cmp, Result
pertinent)
  where
    -- Literals are always primitive values.
    compares :: PatternBase Info vn
-> Result -> InternaliseM (Result, Result, Result)
compares (E.PatternLit PatLit
l Info PatternType
t SrcLoc
_) (SubExp
se : Result
ses) = do
      SubExp
e' <- case PatLit
l of
        PatLitPrim PrimValue
v -> SubExp -> InternaliseM SubExp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> InternaliseM SubExp) -> SubExp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
forall v. IsValue v => v -> SubExp
constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimValue
internalisePrimValue PrimValue
v
        PatLitInt Integer
x -> String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"constant" (ExpBase Info VName -> InternaliseM SubExp)
-> ExpBase Info VName -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ Integer -> Info PatternType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
Integer -> f PatternType -> SrcLoc -> ExpBase f vn
E.IntLit Integer
x Info PatternType
t SrcLoc
forall a. Monoid a => a
mempty
        PatLitFloat Double
x -> String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"constant" (ExpBase Info VName -> InternaliseM SubExp)
-> ExpBase Info VName -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ Double -> Info PatternType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
Double -> f PatternType -> SrcLoc -> ExpBase f vn
E.FloatLit Double
x Info PatternType
t SrcLoc
forall a. Monoid a => a
mempty
      PrimType
t' <- Type -> PrimType
forall shape u. TypeBase shape u -> PrimType
elemType (Type -> PrimType) -> InternaliseM Type -> InternaliseM PrimType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
se
      SubExp
cmp <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"match_lit" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
t') SubExp
e' SubExp
se
      (Result, Result, Result) -> InternaliseM (Result, Result, Result)
forall (m :: * -> *) a. Monad m => a -> m a
return ([SubExp
cmp], [SubExp
se], Result
ses)
    compares (E.PatternConstr Name
c (Info (E.Scalar (E.Sum Map Name [PatternType]
fs))) [PatternBase Info vn]
pats SrcLoc
_) (SubExp
se : Result
ses) = do
      ([TypeBase ExtShape Uniqueness]
payload_ts, Map Name (Int, [Int])
m) <- Map Name [StructType]
-> InternaliseM
     ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
internaliseSumType (Map Name [StructType]
 -> InternaliseM
      ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int])))
-> Map Name [StructType]
-> InternaliseM
     ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
forall a b. (a -> b) -> a -> b
$ ([PatternType] -> [StructType])
-> Map Name [PatternType] -> Map Name [StructType]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((PatternType -> StructType) -> [PatternType] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct) Map Name [PatternType]
fs
      case Name -> Map Name (Int, [Int]) -> Maybe (Int, [Int])
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
c Map Name (Int, [Int])
m of
        Just (Int
i, [Int]
payload_is) -> do
          let i' :: SubExp
i' = IntType -> Integer -> SubExp
intConst IntType
Int8 (Integer -> SubExp) -> Integer -> SubExp
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
i
          let (Result
payload_ses, Result
ses') = Int -> Result -> (Result, Result)
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) Result
ses
          SubExp
cmp <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"match_constr" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
int8) SubExp
i' SubExp
se
          (Result
cmps, Result
pertinent, Result
_) <- [PatternBase Info vn]
-> Result -> InternaliseM (Result, Result, Result)
comparesMany [PatternBase Info vn]
pats (Result -> InternaliseM (Result, Result, Result))
-> Result -> InternaliseM (Result, Result, Result)
forall a b. (a -> b) -> a -> b
$ (Int -> SubExp) -> [Int] -> Result
forall a b. (a -> b) -> [a] -> [b]
map (Result
payload_ses Result -> Int -> SubExp
forall a. [a] -> Int -> a
!!) [Int]
payload_is
          (Result, Result, Result) -> InternaliseM (Result, Result, Result)
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
cmp SubExp -> Result -> Result
forall a. a -> [a] -> [a]
: Result
cmps, Result
pertinent, Result
ses')
        Maybe (Int, [Int])
Nothing ->
          String -> InternaliseM (Result, Result, Result)
forall a. HasCallStack => String -> a
error String
"generateCond: missing constructor"
    compares (E.PatternConstr Name
_ (Info PatternType
t) [PatternBase Info vn]
_ SrcLoc
_) Result
_ =
      String -> InternaliseM (Result, Result, Result)
forall a. HasCallStack => String -> a
error (String -> InternaliseM (Result, Result, Result))
-> String -> InternaliseM (Result, Result, Result)
forall a b. (a -> b) -> a -> b
$ String
"generateCond: PatternConstr has nonsensical type: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PatternType -> String
forall a. Pretty a => a -> String
pretty PatternType
t
    compares (E.Id vn
_ Info PatternType
t SrcLoc
loc) Result
ses =
      PatternBase Info vn
-> Result -> InternaliseM (Result, Result, Result)
compares (Info PatternType -> SrcLoc -> PatternBase Info vn
forall (f :: * -> *) vn.
f PatternType -> SrcLoc -> PatternBase f vn
E.Wildcard Info PatternType
t SrcLoc
loc) Result
ses
    compares (E.Wildcard (Info PatternType
t) SrcLoc
_) Result
ses = do
      Int
n <- StructType -> InternaliseM Int
internalisedTypeSize (StructType -> InternaliseM Int) -> StructType -> InternaliseM Int
forall a b. (a -> b) -> a -> b
$ PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
t
      let (Result
id_ses, Result
rest_ses) = Int -> Result -> (Result, Result)
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n Result
ses
      (Result, Result, Result) -> InternaliseM (Result, Result, Result)
forall (m :: * -> *) a. Monad m => a -> m a
return ([], Result
id_ses, Result
rest_ses)
    compares (E.PatternParens PatternBase Info vn
pat SrcLoc
_) Result
ses =
      PatternBase Info vn
-> Result -> InternaliseM (Result, Result, Result)
compares PatternBase Info vn
pat Result
ses
    -- XXX: treat empty tuples and records as bool.
    compares (E.TuplePattern [] SrcLoc
loc) Result
ses =
      PatternBase Info vn
-> Result -> InternaliseM (Result, Result, Result)
compares (Info PatternType -> SrcLoc -> PatternBase Info vn
forall (f :: * -> *) vn.
f PatternType -> SrcLoc -> PatternBase f vn
E.Wildcard (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) Aliasing -> PatternType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar (ScalarTypeBase (DimDecl VName) Aliasing -> PatternType)
-> ScalarTypeBase (DimDecl VName) Aliasing -> PatternType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) Aliasing
forall dim as. PrimType -> ScalarTypeBase dim as
E.Prim PrimType
E.Bool) SrcLoc
loc) Result
ses
    compares (E.RecordPattern [] SrcLoc
loc) Result
ses =
      PatternBase Info vn
-> Result -> InternaliseM (Result, Result, Result)
compares (Info PatternType -> SrcLoc -> PatternBase Info vn
forall (f :: * -> *) vn.
f PatternType -> SrcLoc -> PatternBase f vn
E.Wildcard (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) Aliasing -> PatternType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar (ScalarTypeBase (DimDecl VName) Aliasing -> PatternType)
-> ScalarTypeBase (DimDecl VName) Aliasing -> PatternType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) Aliasing
forall dim as. PrimType -> ScalarTypeBase dim as
E.Prim PrimType
E.Bool) SrcLoc
loc) Result
ses
    compares (E.TuplePattern [PatternBase Info vn]
pats SrcLoc
_) Result
ses =
      [PatternBase Info vn]
-> Result -> InternaliseM (Result, Result, Result)
comparesMany [PatternBase Info vn]
pats Result
ses
    compares (E.RecordPattern [(Name, PatternBase Info vn)]
fs SrcLoc
_) Result
ses =
      [PatternBase Info vn]
-> Result -> InternaliseM (Result, Result, Result)
comparesMany (((Name, PatternBase Info vn) -> PatternBase Info vn)
-> [(Name, PatternBase Info vn)] -> [PatternBase Info vn]
forall a b. (a -> b) -> [a] -> [b]
map (Name, PatternBase Info vn) -> PatternBase Info vn
forall a b. (a, b) -> b
snd ([(Name, PatternBase Info vn)] -> [PatternBase Info vn])
-> [(Name, PatternBase Info vn)] -> [PatternBase Info vn]
forall a b. (a -> b) -> a -> b
$ Map Name (PatternBase Info vn) -> [(Name, PatternBase Info vn)]
forall a. Map Name a -> [(Name, a)]
E.sortFields (Map Name (PatternBase Info vn) -> [(Name, PatternBase Info vn)])
-> Map Name (PatternBase Info vn) -> [(Name, PatternBase Info vn)]
forall a b. (a -> b) -> a -> b
$ [(Name, PatternBase Info vn)] -> Map Name (PatternBase Info vn)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, PatternBase Info vn)]
fs) Result
ses
    compares (E.PatternAscription PatternBase Info vn
pat TypeDeclBase Info vn
_ SrcLoc
_) Result
ses =
      PatternBase Info vn
-> Result -> InternaliseM (Result, Result, Result)
compares PatternBase Info vn
pat Result
ses
    compares PatternBase Info vn
pat [] =
      String -> InternaliseM (Result, Result, Result)
forall a. HasCallStack => String -> a
error (String -> InternaliseM (Result, Result, Result))
-> String -> InternaliseM (Result, Result, Result)
forall a b. (a -> b) -> a -> b
$ String
"generateCond: No values left for pattern " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PatternBase Info vn -> String
forall a. Pretty a => a -> String
pretty PatternBase Info vn
pat

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

generateCaseIf :: [I.SubExp] -> Case -> I.Body -> InternaliseM I.Exp
generateCaseIf :: Result -> CaseBase Info VName -> Body -> InternaliseM (ExpT SOACS)
generateCaseIf Result
ses (CasePat PatternBase Info VName
p ExpBase Info VName
eCase SrcLoc
_) Body
bFail = do
  (SubExp
cond, Result
pertinent) <- PatternBase Info VName -> Result -> InternaliseM (SubExp, Result)
generateCond PatternBase Info VName
p Result
ses
  Body
eCase' <- PatternBase Info VName
-> Result
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM Body)
-> InternaliseM Body
forall a.
PatternBase Info VName
-> Result
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat' PatternBase Info VName
p Result
pertinent ExpBase Info VName
eCase (String -> ExpBase Info VName -> InternaliseM Body
internaliseBody String
"case")
  InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
(MonadBinder m, BranchType (Lore m) ~ ExtType) =>
m (Exp (Lore m))
-> m (Body (Lore m)) -> m (Body (Lore m)) -> m (Exp (Lore m))
eIf (SubExp -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *). MonadBinder m => SubExp -> m (Exp (Lore m))
eSubExp SubExp
cond) (Body -> InternaliseM Body
forall (m :: * -> *) a. Monad m => a -> m a
return Body
eCase') (Body -> InternaliseM Body
forall (m :: * -> *) a. Monad m => a -> m a
return Body
bFail)

internalisePat ::
  String ->
  E.Pattern ->
  E.Exp ->
  E.Exp ->
  (E.Exp -> InternaliseM a) ->
  InternaliseM a
internalisePat :: String
-> PatternBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat String
desc PatternBase Info VName
p ExpBase Info VName
e ExpBase Info VName
body ExpBase Info VName -> InternaliseM a
m = do
  Result
ses <- String -> ExpBase Info VName -> InternaliseM Result
internaliseExp String
desc' ExpBase Info VName
e
  PatternBase Info VName
-> Result
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
forall a.
PatternBase Info VName
-> Result
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat' PatternBase Info VName
p Result
ses ExpBase Info VName
body ExpBase Info VName -> InternaliseM a
m
  where
    desc' :: String
desc' = case Set (IdentBase Info VName) -> [IdentBase Info VName]
forall a. Set a -> [a]
S.toList (Set (IdentBase Info VName) -> [IdentBase Info VName])
-> Set (IdentBase Info VName) -> [IdentBase Info VName]
forall a b. (a -> b) -> a -> b
$ PatternBase Info VName -> Set (IdentBase Info VName)
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set (IdentBase f vn)
E.patternIdents PatternBase Info VName
p of
      [IdentBase Info VName
v] -> VName -> 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.Pattern ->
  [I.SubExp] ->
  E.Exp ->
  (E.Exp -> InternaliseM a) ->
  InternaliseM a
internalisePat' :: PatternBase Info VName
-> Result
-> ExpBase Info VName
-> (ExpBase Info VName -> InternaliseM a)
-> InternaliseM a
internalisePat' PatternBase Info VName
p Result
ses ExpBase Info VName
body ExpBase Info VName -> InternaliseM a
m = do
  [Type]
ses_ts <- (SubExp -> InternaliseM Type) -> Result -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType Result
ses
  PatternBase Info VName
-> [Type] -> ([VName] -> InternaliseM a) -> InternaliseM a
forall a.
PatternBase Info VName
-> [Type] -> ([VName] -> InternaliseM a) -> InternaliseM a
stmPattern PatternBase Info VName
p [Type]
ses_ts (([VName] -> InternaliseM a) -> InternaliseM a)
-> ([VName] -> InternaliseM a) -> InternaliseM a
forall a b. (a -> b) -> a -> b
$ \[VName]
pat_names -> do
    [(VName, SubExp)]
-> ((VName, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([VName] -> Result -> [(VName, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
pat_names Result
ses) (((VName, SubExp) -> InternaliseM ()) -> InternaliseM ())
-> ((VName, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \(VName
v, SubExp
se) ->
      [VName] -> Exp (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [VName
v] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
se
    ExpBase Info VName -> InternaliseM a
m ExpBase Info VName
body

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

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

-- Special-case an important common case that otherwise leads to horrible code.
internaliseDimIndex
  SubExp
w
  ( E.DimSlice
      Maybe (ExpBase Info VName)
Nothing
      Maybe (ExpBase Info VName)
Nothing
      (Just (E.Negate (E.IntLit Integer
1 Info PatternType
_ SrcLoc
_) SrcLoc
_))
    ) = do
    SubExp
w_minus_1 <-
      String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"w_minus_1" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
        BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) SubExp
w SubExp
one
    (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return
      ( SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
I.DimSlice SubExp
w_minus_1 SubExp
w (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 (-Integer
1),
        Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
True,
        [ErrorMsgPart SubExp]
forall a. Monoid a => a
mempty
      )
    where
      one :: SubExp
one = Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
1 :: Int64)
internaliseDimIndex SubExp
w (E.DimSlice Maybe (ExpBase Info VName)
i Maybe (ExpBase Info VName)
j Maybe (ExpBase Info VName)
s) = do
  SubExp
s' <- InternaliseM SubExp
-> (ExpBase Info VName -> InternaliseM SubExp)
-> Maybe (ExpBase Info VName)
-> InternaliseM SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (SubExp -> InternaliseM SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return SubExp
one) (((SubExp, IntType) -> SubExp)
-> InternaliseM (SubExp, IntType) -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SubExp, IntType) -> SubExp
forall a b. (a, b) -> a
fst (InternaliseM (SubExp, IntType) -> InternaliseM SubExp)
-> (ExpBase Info VName -> InternaliseM (SubExp, IntType))
-> ExpBase Info VName
-> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ExpBase Info VName -> InternaliseM (SubExp, IntType)
internaliseDimExp String
"s") Maybe (ExpBase Info VName)
s
  SubExp
s_sign <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"s_sign" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (IntType -> UnOp
I.SSignum IntType
Int64) SubExp
s'
  SubExp
backwards <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"backwards" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
int64) SubExp
s_sign SubExp
negone
  SubExp
w_minus_1 <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"w_minus_1" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) SubExp
w SubExp
one
  let i_def :: InternaliseM SubExp
i_def =
        String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"i_def" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          SubExp -> Body -> Body -> IfDec (BranchType SOACS) -> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If
            SubExp
backwards
            (Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody [SubExp
w_minus_1])
            (Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody [SubExp
zero])
            (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ [Type] -> IfDec ExtType
ifCommon [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64]
      j_def :: InternaliseM SubExp
j_def =
        String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"j_def" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          SubExp -> Body -> Body -> IfDec (BranchType SOACS) -> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If
            SubExp
backwards
            (Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody [SubExp
negone])
            (Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody [SubExp
w])
            (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ [Type] -> IfDec ExtType
ifCommon [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64]
  SubExp
i' <- InternaliseM SubExp
-> (ExpBase Info VName -> InternaliseM SubExp)
-> Maybe (ExpBase Info VName)
-> InternaliseM SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe InternaliseM SubExp
i_def (((SubExp, IntType) -> SubExp)
-> InternaliseM (SubExp, IntType) -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SubExp, IntType) -> SubExp
forall a b. (a, b) -> a
fst (InternaliseM (SubExp, IntType) -> InternaliseM SubExp)
-> (ExpBase Info VName -> InternaliseM (SubExp, IntType))
-> ExpBase Info VName
-> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ExpBase Info VName -> InternaliseM (SubExp, IntType)
internaliseDimExp String
"i") Maybe (ExpBase Info VName)
i
  SubExp
j' <- InternaliseM SubExp
-> (ExpBase Info VName -> InternaliseM SubExp)
-> Maybe (ExpBase Info VName)
-> InternaliseM SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe InternaliseM SubExp
j_def (((SubExp, IntType) -> SubExp)
-> InternaliseM (SubExp, IntType) -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SubExp, IntType) -> SubExp
forall a b. (a, b) -> a
fst (InternaliseM (SubExp, IntType) -> InternaliseM SubExp)
-> (ExpBase Info VName -> InternaliseM (SubExp, IntType))
-> ExpBase Info VName
-> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ExpBase Info VName -> InternaliseM (SubExp, IntType)
internaliseDimExp String
"j") Maybe (ExpBase Info VName)
j
  SubExp
j_m_i <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"j_m_i" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) SubExp
j' SubExp
i'
  -- Something like a division-rounding-up, but accomodating negative
  -- operands.
  let divRounding :: InternaliseM (ExpT SOACS)
-> InternaliseM (ExpT SOACS)
-> InternaliseM (Exp (Lore InternaliseM))
divRounding InternaliseM (ExpT SOACS)
x InternaliseM (ExpT SOACS)
y =
        BinOp
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> m (Exp (Lore m)) -> m (Exp (Lore m)) -> m (Exp (Lore m))
eBinOp
          (IntType -> Safety -> BinOp
SQuot IntType
Int64 Safety
Unsafe)
          ( BinOp
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> m (Exp (Lore m)) -> m (Exp (Lore m)) -> m (Exp (Lore m))
eBinOp
              (IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
I.OverflowWrap)
              InternaliseM (Exp (Lore InternaliseM))
InternaliseM (ExpT SOACS)
x
              (BinOp
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> m (Exp (Lore m)) -> m (Exp (Lore m)) -> m (Exp (Lore m))
eBinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) InternaliseM (Exp (Lore InternaliseM))
InternaliseM (ExpT SOACS)
y (InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
m (Exp (Lore m)) -> m (Exp (Lore m))
eSignum (InternaliseM (Exp (Lore InternaliseM))
 -> InternaliseM (Exp (Lore InternaliseM)))
-> InternaliseM (Exp (Lore InternaliseM))
-> InternaliseM (Exp (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ SubExp -> InternaliseM (Exp (Lore InternaliseM))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Lore m))
toExp SubExp
s'))
          )
          InternaliseM (Exp (Lore InternaliseM))
InternaliseM (ExpT SOACS)
y
  SubExp
n <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"n" (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (ExpT SOACS)
-> InternaliseM (ExpT SOACS)
-> InternaliseM (Exp (Lore InternaliseM))
divRounding (SubExp -> InternaliseM (Exp (Lore InternaliseM))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Lore m))
toExp SubExp
j_m_i) (SubExp -> InternaliseM (Exp (Lore InternaliseM))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Lore m))
toExp SubExp
s')

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

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

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

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

  SubExp
slice_ok <-
    String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"slice_ok" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      SubExp -> Body -> Body -> IfDec (BranchType SOACS) -> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If
        SubExp
backwards
        (Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody [SubExp
backwards_ok])
        (Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody [SubExp
forwards_ok])
        (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ [Type] -> IfDec ExtType
ifCommon [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.Bool]
  SubExp
ok_or_empty <-
    String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"ok_or_empty" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogOr SubExp
empty_slice SubExp
slice_ok

  let parts :: [ErrorMsgPart SubExp]
parts = case (Maybe (ExpBase Info VName)
i, Maybe (ExpBase Info VName)
j, Maybe (ExpBase Info VName)
s) of
        (Maybe (ExpBase Info VName)
_, Maybe (ExpBase Info VName)
_, Just {}) ->
          [ ErrorMsgPart SubExp
-> (ExpBase Info VName -> ErrorMsgPart SubExp)
-> Maybe (ExpBase Info VName)
-> ErrorMsgPart SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp
forall a b. a -> b -> a
const (ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp)
-> ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt64 SubExp
i') Maybe (ExpBase Info VName)
i,
            ErrorMsgPart SubExp
":",
            ErrorMsgPart SubExp
-> (ExpBase Info VName -> ErrorMsgPart SubExp)
-> Maybe (ExpBase Info VName)
-> ErrorMsgPart SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp
forall a b. a -> b -> a
const (ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp)
-> ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt64 SubExp
j') Maybe (ExpBase Info VName)
j,
            ErrorMsgPart SubExp
":",
            SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt64 SubExp
s'
          ]
        (Maybe (ExpBase Info VName)
_, Just {}, Maybe (ExpBase Info VName)
_) ->
          [ ErrorMsgPart SubExp
-> (ExpBase Info VName -> ErrorMsgPart SubExp)
-> Maybe (ExpBase Info VName)
-> ErrorMsgPart SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp
forall a b. a -> b -> a
const (ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp)
-> ErrorMsgPart SubExp -> ExpBase Info VName -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt64 SubExp
i') Maybe (ExpBase Info VName)
i,
            ErrorMsgPart SubExp
":",
            SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt64 SubExp
j'
          ]
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
-> (ExpBase Info VName -> [ErrorMsgPart SubExp])
-> Maybe (ExpBase Info VName)
-> [ErrorMsgPart SubExp]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [ErrorMsgPart SubExp]
forall a. Monoid a => a
mempty ([ErrorMsgPart SubExp]
-> ExpBase Info VName -> [ErrorMsgPart SubExp]
forall a b. a -> b -> a
const [ErrorMsgPart SubExp
":", SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt64 SubExp
s']) Maybe (ExpBase Info VName)
s
        (Maybe (ExpBase Info VName)
_, Maybe (ExpBase Info VName)
Nothing, Maybe (ExpBase Info VName)
Nothing) ->
          [SubExp -> ErrorMsgPart SubExp
forall a. a -> ErrorMsgPart a
ErrorInt64 SubExp
i', ErrorMsgPart SubExp
":"]
  (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
I.DimSlice SubExp
i' SubExp
n SubExp
s', SubExp
ok_or_empty, [ErrorMsgPart SubExp]
parts)
  where
    zero :: SubExp
zero = Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)
    negone :: SubExp
negone = Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (-Int64
1 :: Int64)
    one :: SubExp
one = Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
1 :: Int64)

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

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

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

  -- reshape return type of bucket function to have same size as neutral element
  -- (modulo the index)
  Param Type
bucket_param <- String -> Type -> InternaliseM (Param Type)
forall (m :: * -> *) dec.
MonadFreshNames m =>
String -> dec -> m (Param dec)
newParam String
"bucket_p" (Type -> InternaliseM (Param Type))
-> Type -> InternaliseM (Param Type)
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64
  [Param Type]
img_params <- (Type -> InternaliseM (Param Type))
-> [Type] -> InternaliseM [Param Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> Type -> InternaliseM (Param Type)
forall (m :: * -> *) dec.
MonadFreshNames m =>
String -> dec -> m (Param dec)
newParam String
"img_p" (Type -> InternaliseM (Param Type))
-> (Type -> Type) -> Type -> InternaliseM (Param Type)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
rowType) ([Type] -> InternaliseM [Param Type])
-> InternaliseM [Type] -> InternaliseM [Param Type]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (VName -> InternaliseM Type) -> [VName] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType [VName]
img'
  let params :: [Param Type]
params = Param Type
bucket_param Param Type -> [Param Type] -> [Param Type]
forall a. a -> [a] -> [a]
: [Param Type]
img_params
      rettype :: [Type]
rettype = PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64 Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
: [Type]
ne_ts
      body :: Body
body = Stms SOACS -> Result -> Body
forall lore. Bindable lore => Stms lore -> Result -> Body lore
mkBody Stms SOACS
forall a. Monoid a => a
mempty (Result -> Body) -> Result -> Body
forall a b. (a -> b) -> a -> b
$ (Param Type -> SubExp) -> [Param Type] -> Result
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
I.Var (VName -> SubExp) -> (Param Type -> VName) -> Param Type -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> VName
forall dec. Param dec -> VName
paramName) [Param Type]
params
  Body
body' <-
    Scope SOACS -> InternaliseM Body -> InternaliseM Body
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param Type] -> Scope SOACS
forall lore dec.
(LParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
scopeOfLParams [Param Type]
params) (InternaliseM Body -> InternaliseM Body)
-> InternaliseM Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$
      ErrorMsg SubExp -> SrcLoc -> [Type] -> Body -> InternaliseM Body
ensureResultShape
        ErrorMsg SubExp
"Row shape of value array does not match row shape of hist target"
        (ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
img)
        [Type]
rettype
        Body
body

  -- get sizes of histogram and image arrays
  SubExp
w_hist <- Int -> [Type] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 ([Type] -> SubExp) -> InternaliseM [Type] -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM Type) -> [VName] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType [VName]
hist'
  SubExp
w_img <- Int -> [Type] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 ([Type] -> SubExp) -> InternaliseM [Type] -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM Type) -> [VName] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType [VName]
img'

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

  String -> Exp (Lore InternaliseM) -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
    Op SOACS -> ExpT SOACS
forall lore. Op lore -> ExpT lore
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
      SubExp -> [HistOp SOACS] -> Lambda -> [VName] -> SOAC SOACS
forall lore.
SubExp -> [HistOp lore] -> Lambda lore -> [VName] -> SOAC lore
I.Hist SubExp
w_img [SubExp -> SubExp -> [VName] -> Result -> Lambda -> HistOp SOACS
forall lore.
SubExp -> SubExp -> [VName] -> Result -> Lambda lore -> HistOp lore
HistOp SubExp
w_hist SubExp
rf' [VName]
hist' Result
ne_shp Lambda
op'] ([LParam] -> Body -> [Type] -> Lambda
forall lore. [LParam lore] -> BodyT lore -> [Type] -> LambdaT lore
I.Lambda [Param Type]
[LParam]
params Body
body' [Type]
rettype) ([VName] -> SOAC SOACS) -> [VName] -> SOAC SOACS
forall a b. (a -> b) -> a -> b
$ VName
buckets'' VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: [VName]
img'

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

internaliseStreamRed ::
  String ->
  StreamOrd ->
  Commutativity ->
  E.Exp ->
  E.Exp ->
  E.Exp ->
  InternaliseM [SubExp]
internaliseStreamRed :: String
-> StreamOrd
-> Commutativity
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> InternaliseM Result
internaliseStreamRed String
desc StreamOrd
o Commutativity
comm ExpBase Info VName
lam0 ExpBase Info VName
lam ExpBase Info VName
arr = do
  [VName]
arrs <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"stream_input" ExpBase Info VName
arr
  [Type]
rowts <- (VName -> InternaliseM Type) -> [VName] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Type -> Type) -> InternaliseM Type -> InternaliseM Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
I.rowType (InternaliseM Type -> InternaliseM Type)
-> (VName -> InternaliseM Type) -> VName -> InternaliseM Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType) [VName]
arrs
  ([Param Type]
lam_params, Body
lam_body) <-
    InternaliseLambda
-> ExpBase Info VName -> [Type] -> InternaliseM ([LParam], Body)
internaliseStreamLambda InternaliseLambda
internaliseLambda ExpBase Info VName
lam [Type]
rowts
  let (Param Type
chunk_param, [Param Type]
_, [Param Type]
lam_val_params) =
        Int -> [Param Type] -> (Param Type, [Param Type], [Param Type])
forall dec.
Int -> [Param dec] -> (Param dec, [Param dec], [Param dec])
partitionChunkedFoldParameters Int
0 [Param Type]
lam_params

  -- Synthesize neutral elements by applying the fold function
  -- to an empty chunk.
  [VName] -> Exp (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [Param Type -> VName
forall dec. Param dec -> VName
I.paramName Param Type
chunk_param] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
    BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)
  [Param Type] -> (Param Type -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Param Type]
lam_val_params ((Param Type -> InternaliseM ()) -> InternaliseM ())
-> (Param Type -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \Param Type
p ->
    [VName] -> Exp (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [Param Type -> VName
forall dec. Param dec -> VName
I.paramName Param Type
p] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
        PrimType -> Result -> BasicOp
I.Scratch (Type -> PrimType
forall shape u. TypeBase shape u -> PrimType
I.elemType (Type -> PrimType) -> Type -> PrimType
forall a b. (a -> b) -> a -> b
$ Param Type -> Type
forall dec. Typed dec => Param dec -> Type
I.paramType Param Type
p) (Result -> BasicOp) -> Result -> BasicOp
forall a b. (a -> b) -> a -> b
$
          Type -> Result
forall u. TypeBase Shape u -> Result
I.arrayDims (Type -> Result) -> Type -> Result
forall a b. (a -> b) -> a -> b
$ Param Type -> Type
forall dec. Typed dec => Param dec -> Type
I.paramType Param Type
p
  Result
nes <- Body -> InternaliseM Result
forall (m :: * -> *). MonadBinder m => Body (Lore m) -> m Result
bodyBind (Body -> InternaliseM Result)
-> InternaliseM Body -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Body -> InternaliseM Body
forall lore (m :: * -> *).
(Renameable lore, MonadFreshNames m) =>
Body lore -> m (Body lore)
renameBody Body
lam_body

  [Type]
nes_ts <- (SubExp -> InternaliseM Type) -> Result -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
I.subExpType Result
nes
  SubExp
outsz <- Int -> [Type] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 ([Type] -> SubExp) -> InternaliseM [Type] -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM Type) -> [VName] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType [VName]
arrs
  let acc_arr_tps :: [Type]
acc_arr_tps = [Type -> Shape -> NoUniqueness -> Type
forall shape u_unused u.
ArrayShape shape =>
TypeBase shape u_unused -> shape -> u -> TypeBase shape u
I.arrayOf Type
t (Result -> Shape
forall d. [d] -> ShapeBase d
I.Shape [SubExp
outsz]) NoUniqueness
NoUniqueness | Type
t <- [Type]
nes_ts]
  Lambda
lam0' <- InternaliseLambda
-> ExpBase Info VName -> [Type] -> [Type] -> InternaliseM Lambda
internaliseFoldLambda InternaliseLambda
internaliseLambda ExpBase Info VName
lam0 [Type]
nes_ts [Type]
acc_arr_tps

  let lam0_acc_params :: [Param Type]
lam0_acc_params = Int -> [Param Type] -> [Param Type]
forall a. Int -> [a] -> [a]
take (Result -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Result
nes) ([Param Type] -> [Param Type]) -> [Param Type] -> [Param Type]
forall a b. (a -> b) -> a -> b
$ Lambda -> [LParam]
forall lore. LambdaT lore -> [LParam lore]
I.lambdaParams Lambda
lam0'
  [Param Type]
lam_acc_params <- [Param Type]
-> (Param Type -> InternaliseM (Param Type))
-> InternaliseM [Param Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Param Type]
lam0_acc_params ((Param Type -> InternaliseM (Param Type))
 -> InternaliseM [Param Type])
-> (Param Type -> InternaliseM (Param Type))
-> InternaliseM [Param Type]
forall a b. (a -> b) -> a -> b
$ \Param Type
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 Type -> VName
forall dec. Param dec -> VName
I.paramName Param Type
p
    Param Type -> InternaliseM (Param Type)
forall (m :: * -> *) a. Monad m => a -> m a
return Param Type
p {paramName :: VName
I.paramName = VName
name}

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

  Body
body_with_lam0 <-
    ErrorMsg SubExp -> SrcLoc -> [Type] -> Body -> InternaliseM Body
ensureResultShape
      ErrorMsg SubExp
"shape of result does not match shape of initial value"
      (ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
lam0)
      [Type]
nes_ts
      (Body -> InternaliseM Body)
-> (InternaliseM Body -> InternaliseM Body)
-> InternaliseM Body
-> InternaliseM Body
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< InternaliseM Body -> InternaliseM Body
forall (m :: * -> *).
MonadBinder m =>
m (Body (Lore m)) -> m (Body (Lore m))
insertStmsM
      (InternaliseM Body -> InternaliseM Body)
-> InternaliseM Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$ Scope SOACS -> InternaliseM Body -> InternaliseM Body
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param Type] -> Scope SOACS
forall lore dec.
(LParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
scopeOfLParams [Param Type]
lam_params') (InternaliseM Body -> InternaliseM Body)
-> InternaliseM Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$ do
        Result
lam_res <- Body (Lore InternaliseM) -> InternaliseM Result
forall (m :: * -> *). MonadBinder m => Body (Lore m) -> m Result
bodyBind Body (Lore InternaliseM)
Body
lam_body
        Result
lam_res' <-
          ErrorMsg SubExp
-> SrcLoc -> [VName] -> [Type] -> Result -> InternaliseM Result
forall u.
Typed (TypeBase Shape u) =>
ErrorMsg SubExp
-> SrcLoc
-> [VName]
-> [TypeBase Shape u]
-> Result
-> InternaliseM Result
ensureArgShapes
            ErrorMsg SubExp
"shape of chunk function result does not match shape of initial value"
            (ExpBase Info VName -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ExpBase Info VName
lam)
            []
            ((Param Type -> Type) -> [Param Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> Type
forall t. Typed t => t -> Type
I.typeOf ([Param Type] -> [Type]) -> [Param Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ Lambda -> [LParam]
forall lore. LambdaT lore -> [LParam lore]
I.lambdaParams Lambda
lam0')
            Result
lam_res
        Result
new_lam_res <-
          Lambda (Lore InternaliseM)
-> [InternaliseM (Exp (Lore InternaliseM))] -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
Lambda (Lore m) -> [m (Exp (Lore m))] -> m Result
eLambda Lambda (Lore InternaliseM)
Lambda
lam0' ([InternaliseM (Exp (Lore InternaliseM))] -> InternaliseM Result)
-> [InternaliseM (Exp (Lore InternaliseM))] -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
            (SubExp -> InternaliseM (ExpT SOACS))
-> Result -> [InternaliseM (ExpT SOACS)]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> InternaliseM (ExpT SOACS)
forall (m :: * -> *). MonadBinder m => SubExp -> m (Exp (Lore m))
eSubExp (Result -> [InternaliseM (ExpT SOACS)])
-> Result -> [InternaliseM (ExpT SOACS)]
forall a b. (a -> b) -> a -> b
$
              (Param Type -> SubExp) -> [Param Type] -> Result
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
I.Var (VName -> SubExp) -> (Param Type -> VName) -> Param Type -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> VName
forall dec. Param dec -> VName
paramName) [Param Type]
lam_acc_params Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ Result
lam_res'
        Body -> InternaliseM Body
forall (m :: * -> *) a. Monad m => a -> m a
return (Body -> InternaliseM Body) -> Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$ Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody Result
new_lam_res

  let form :: StreamForm SOACS
form = StreamOrd -> Commutativity -> Lambda -> StreamForm SOACS
forall lore.
StreamOrd -> Commutativity -> Lambda lore -> StreamForm lore
I.Parallel StreamOrd
o Commutativity
comm Lambda
lam0'
      lam' :: Lambda
lam' =
        Lambda :: forall lore. [LParam lore] -> BodyT lore -> [Type] -> LambdaT lore
I.Lambda
          { lambdaParams :: [LParam]
lambdaParams = [Param Type]
[LParam]
lam_params',
            lambdaBody :: Body
lambdaBody = Body
body_with_lam0,
            lambdaReturnType :: [Type]
lambdaReturnType = [Type]
nes_ts
          }
  SubExp
w <- Int -> [Type] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 ([Type] -> SubExp) -> InternaliseM [Type] -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM Type) -> [VName] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType [VName]
arrs
  String -> Exp (Lore InternaliseM) -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ Op SOACS -> ExpT SOACS
forall lore. Op lore -> ExpT lore
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp
-> [VName] -> StreamForm SOACS -> Result -> Lambda -> SOAC SOACS
forall lore.
SubExp
-> [VName] -> StreamForm lore -> Result -> Lambda lore -> SOAC lore
I.Stream SubExp
w [VName]
arrs StreamForm SOACS
form Result
nes Lambda
lam'

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

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

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

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

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

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

internaliseBinOp ::
  SrcLoc ->
  String ->
  E.BinOp ->
  I.SubExp ->
  I.SubExp ->
  E.PrimType ->
  E.PrimType ->
  InternaliseM [I.SubExp]
internaliseBinOp :: SrcLoc
-> String
-> BinOp
-> SubExp
-> SubExp
-> PrimType
-> PrimType
-> InternaliseM Result
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Plus SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  String -> BinOp -> SubExp -> SubExp -> InternaliseM Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result -> InternaliseM Result
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
    String -> BinOp -> SubExp -> SubExp -> InternaliseM Result
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 Result -> InternaliseM Result
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
    String -> BinOp -> SubExp -> SubExp -> InternaliseM Result
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 Result
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 Result
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 Result -> InternaliseM Result
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonnegative SrcLoc
loc IntType
t SubExp
y (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
    String -> BinOp -> SubExp -> SubExp -> InternaliseM Result
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 Result
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 Result -> InternaliseM Result
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
    String -> BinOp -> SubExp -> SubExp -> InternaliseM Result
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 Result -> InternaliseM Result
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
    String -> BinOp -> SubExp -> SubExp -> InternaliseM Result
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 Result
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 Result -> InternaliseM Result
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
    String -> BinOp -> SubExp -> SubExp -> InternaliseM Result
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 Result -> InternaliseM Result
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
    String -> BinOp -> SubExp -> SubExp -> InternaliseM Result
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 Result -> InternaliseM Result
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
    String -> BinOp -> SubExp -> SubExp -> InternaliseM Result
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 Result -> InternaliseM Result
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
    String -> BinOp -> SubExp -> SubExp -> InternaliseM Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp (String
desc String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"true") (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimType
internalisePrimType PrimType
t) SubExp
x SubExp
y
  (SubExp -> Result) -> InternaliseM SubExp -> InternaliseM Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> Result
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InternaliseM SubExp -> InternaliseM Result)
-> InternaliseM SubExp -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
desc (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
I.Not SubExp
eq
internaliseBinOp SrcLoc
_ String
desc BinOp
E.Less SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  String -> CmpOp -> SubExp -> SubExp -> InternaliseM Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
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 Result
forall a. HasCallStack => String -> a
error (String -> InternaliseM Result) -> String -> InternaliseM Result
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 Result
simpleBinOp String
desc BinOp
bop SubExp
x SubExp
y =
  String -> Exp (Lore InternaliseM) -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
bop SubExp
x SubExp
y

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

findFuncall ::
  E.Exp ->
  InternaliseM
    ( E.QualName VName,
      [(E.Exp, Maybe VName)],
      E.StructType,
      [VName]
    )
findFuncall :: ExpBase Info VName
-> InternaliseM
     (QualName VName, [(ExpBase Info VName, Maybe VName)], StructType,
      [VName])
findFuncall (E.Var QualName VName
fname (Info PatternType
t) SrcLoc
_) =
  (QualName VName, [(ExpBase Info VName, Maybe VName)], StructType,
 [VName])
-> InternaliseM
     (QualName VName, [(ExpBase Info VName, Maybe VName)], StructType,
      [VName])
forall (m :: * -> *) a. Monad m => a -> m a
return (QualName VName
fname, [], PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
t, [])
findFuncall (E.Apply ExpBase Info VName
f ExpBase Info VName
arg (Info (Diet
_, Maybe VName
argext)) (Info PatternType
ret, Info [VName]
retext) SrcLoc
_) = do
  (QualName VName
fname, [(ExpBase Info VName, Maybe VName)]
args, StructType
_, [VName]
_) <- ExpBase Info VName
-> InternaliseM
     (QualName VName, [(ExpBase Info VName, Maybe VName)], StructType,
      [VName])
findFuncall ExpBase Info VName
f
  (QualName VName, [(ExpBase Info VName, Maybe VName)], StructType,
 [VName])
-> InternaliseM
     (QualName VName, [(ExpBase Info VName, Maybe VName)], StructType,
      [VName])
forall (m :: * -> *) a. Monad m => a -> m a
return (QualName VName
fname, [(ExpBase Info VName, Maybe VName)]
args [(ExpBase Info VName, Maybe VName)]
-> [(ExpBase Info VName, Maybe VName)]
-> [(ExpBase Info VName, Maybe VName)]
forall a. [a] -> [a] -> [a]
++ [(ExpBase Info VName
arg, Maybe VName
argext)], PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatternType
ret, [VName]
retext)
findFuncall ExpBase Info VName
e =
  String
-> InternaliseM
     (QualName VName, [(ExpBase Info VName, Maybe VName)], StructType,
      [VName])
forall a. HasCallStack => String -> a
error (String
 -> InternaliseM
      (QualName VName, [(ExpBase Info VName, Maybe VName)], StructType,
       [VName]))
-> String
-> InternaliseM
     (QualName VName, [(ExpBase Info VName, Maybe VName)], StructType,
      [VName])
forall a b. (a -> b) -> a -> b
$ String
"Invalid function expression in application: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ ExpBase Info VName -> String
forall a. Pretty a => a -> String
pretty ExpBase Info VName
e

internaliseLambda :: InternaliseLambda
internaliseLambda :: InternaliseLambda
internaliseLambda (E.Parens ExpBase Info VName
e SrcLoc
_) [Type]
rowtypes =
  InternaliseLambda
internaliseLambda ExpBase Info VName
e [Type]
rowtypes
internaliseLambda (E.Lambda [PatternBase Info VName]
params ExpBase Info VName
body Maybe (TypeExp VName)
_ (Info (Aliasing
_, StructType
rettype)) SrcLoc
_) [Type]
rowtypes =
  [PatternBase Info VName]
-> [Type]
-> ([LParam] -> InternaliseM ([Param Type], Body, [Type]))
-> InternaliseM ([Param Type], Body, [Type])
forall a.
[PatternBase Info VName]
-> [Type] -> ([LParam] -> InternaliseM a) -> InternaliseM a
bindingLambdaParams [PatternBase Info VName]
params [Type]
rowtypes (([LParam] -> InternaliseM ([Param Type], Body, [Type]))
 -> InternaliseM ([Param Type], Body, [Type]))
-> ([LParam] -> InternaliseM ([Param Type], Body, [Type]))
-> InternaliseM ([Param Type], Body, [Type])
forall a b. (a -> b) -> a -> b
$ \[LParam]
params' -> do
    Body
body' <- String -> ExpBase Info VName -> InternaliseM Body
internaliseBody String
"lam" ExpBase Info VName
body
    [Type]
rettype' <- StructType -> InternaliseM [Type]
internaliseLambdaReturnType StructType
rettype
    ([Param Type], Body, [Type])
-> InternaliseM ([Param Type], Body, [Type])
forall (m :: * -> *) a. Monad m => a -> m a
return ([Param Type]
[LParam]
params', Body
body', [Type]
rettype')
internaliseLambda ExpBase Info VName
e [Type]
_ = String -> InternaliseM ([Param Type], Body, [Type])
forall a. HasCallStack => String -> a
error (String -> InternaliseM ([Param Type], Body, [Type]))
-> String -> InternaliseM ([Param Type], Body, [Type])
forall a b. (a -> b) -> a -> b
$ String
"internaliseLambda: unexpected expression:\n" String -> String -> String
forall a. [a] -> [a] -> [a]
++ ExpBase Info VName -> String
forall a. Pretty a => a -> String
pretty ExpBase Info VName
e

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

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

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

        doComparison :: String -> SubExp -> SubExp -> m SubExp
doComparison String
desc SubExp
x SubExp
y = do
          Type
x_t <- SubExp -> m Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
I.subExpType SubExp
x
          Type
y_t <- SubExp -> m Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
I.subExpType SubExp
y
          case Type
x_t of
            I.Prim PrimType
t -> String -> Exp (Lore m) -> m SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
desc (Exp (Lore m) -> m SubExp) -> Exp (Lore m) -> m SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
t) SubExp
x SubExp
y
            Type
_ -> do
              let x_dims :: Result
x_dims = Type -> Result
forall u. TypeBase Shape u -> Result
I.arrayDims Type
x_t
                  y_dims :: Result
y_dims = Type -> Result
forall u. TypeBase Shape u -> Result
I.arrayDims Type
y_t
              Result
dims_match <- [(SubExp, SubExp)] -> ((SubExp, SubExp) -> m SubExp) -> m Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (Result -> Result -> [(SubExp, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip Result
x_dims Result
y_dims) (((SubExp, SubExp) -> m SubExp) -> m Result)
-> ((SubExp, SubExp) -> m SubExp) -> m Result
forall a b. (a -> b) -> a -> b
$ \(SubExp
x_dim, SubExp
y_dim) ->
                String -> Exp (Lore m) -> m SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"dim_eq" (Exp (Lore m) -> m SubExp) -> Exp (Lore m) -> m SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
int64) SubExp
x_dim SubExp
y_dim
              SubExp
shapes_match <- String -> Exp (Lore m) -> m SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"shapes_match" (Exp (Lore m) -> m SubExp) -> m (Exp (Lore m)) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Result -> m (Exp (Lore m))
forall (m :: * -> *). MonadBinder m => Result -> m (Exp (Lore m))
eAll Result
dims_match
              Body (Lore m)
compare_elems_body <- Binder (Lore m) (Body (Lore m)) -> m (Body (Lore m))
forall lore (m :: * -> *) somelore.
(Bindable lore, MonadFreshNames m, HasScope somelore m,
 SameScope somelore lore) =>
Binder lore (Body lore) -> m (Body lore)
runBodyBinder (Binder (Lore m) (Body (Lore m)) -> m (Body (Lore m)))
-> Binder (Lore m) (Body (Lore m)) -> m (Body (Lore m))
forall a b. (a -> b) -> a -> b
$ do
                -- Flatten both x and y.
                SubExp
x_num_elems <-
                  String
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"x_num_elems"
                    (Exp (Lore m) -> BinderT (Lore m) (State VNameSource) SubExp)
-> BinderT (Lore m) (State VNameSource) (Exp (Lore m))
-> BinderT (Lore m) (State VNameSource) SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinOp
-> SubExp
-> Result
-> BinderT
     (Lore m)
     (State VNameSource)
     (Exp (Lore (BinderT (Lore m) (State VNameSource))))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> SubExp -> Result -> m (Exp (Lore m))
foldBinOp (IntType -> Overflow -> BinOp
I.Mul IntType
Int64 Overflow
I.OverflowUndef) (Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
1 :: Int64)) Result
x_dims
                VName
x' <- String
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m VName
letExp String
"x" (Exp (Lore (BinderT (Lore m) (State VNameSource)))
 -> BinderT (Lore m) (State VNameSource) VName)
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
x
                VName
y' <- String
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m VName
letExp String
"x" (Exp (Lore (BinderT (Lore m) (State VNameSource)))
 -> BinderT (Lore m) (State VNameSource) VName)
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
y
                VName
x_flat <- String
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m VName
letExp String
"x_flat" (Exp (Lore (BinderT (Lore m) (State VNameSource)))
 -> BinderT (Lore m) (State VNameSource) VName)
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
I.Reshape [SubExp -> DimChange SubExp
forall d. d -> DimChange d
I.DimNew SubExp
x_num_elems] VName
x'
                VName
y_flat <- String
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m VName
letExp String
"y_flat" (Exp (Lore (BinderT (Lore m) (State VNameSource)))
 -> BinderT (Lore m) (State VNameSource) VName)
-> Exp (Lore (BinderT (Lore m) (State VNameSource)))
-> BinderT (Lore m) (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
I.Reshape [SubExp -> DimChange SubExp
forall d. d -> DimChange d
I.DimNew SubExp
x_num_elems] VName
y'

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

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

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

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

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

      let conc :: VName -> VName -> ExpT SOACS
conc VName
xarr VName
yarr =
            BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ Int -> VName -> [VName] -> SubExp -> BasicOp
I.Concat Int
0 VName
xarr [VName
yarr] SubExp
ressize
      String -> [Exp (Lore InternaliseM)] -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
String -> [Exp (Lore m)] -> m Result
letSubExps String
desc ([Exp (Lore InternaliseM)] -> InternaliseM Result)
-> [Exp (Lore InternaliseM)] -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ (VName -> VName -> ExpT SOACS)
-> [VName] -> [VName] -> [ExpT SOACS]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith VName -> VName -> ExpT SOACS
conc [VName]
xs [VName]
ys
    handleRest [TupLit [ExpBase Info VName
offset, ExpBase Info VName
e] SrcLoc
_] String
"rotate" = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ \String
desc -> do
      SubExp
offset' <- String -> ExpBase Info VName -> InternaliseM SubExp
internaliseExp1 String
"rotation_offset" ExpBase Info VName
offset
      String
-> ExpBase Info VName
-> (VName -> InternaliseM BasicOp)
-> InternaliseM Result
internaliseOperation String
desc ExpBase Info VName
e ((VName -> InternaliseM BasicOp) -> InternaliseM Result)
-> (VName -> InternaliseM BasicOp) -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ \VName
v -> do
        Int
r <- Type -> Int
forall shape u. ArrayShape shape => TypeBase shape u -> Int
I.arrayRank (Type -> Int) -> InternaliseM Type -> InternaliseM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
v
        let zero :: SubExp
zero = IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0
            offsets :: Result
offsets = SubExp
offset' SubExp -> Result -> Result
forall a. a -> [a] -> [a]
: Int -> SubExp -> Result
forall a. Int -> a -> [a]
replicate (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) SubExp
zero
        BasicOp -> InternaliseM BasicOp
forall (m :: * -> *) a. Monad m => a -> m a
return (BasicOp -> InternaliseM BasicOp)
-> BasicOp -> InternaliseM BasicOp
forall a b. (a -> b) -> a -> b
$ Result -> VName -> BasicOp
I.Rotate Result
offsets VName
v
    handleRest [ExpBase Info VName
e] String
"transpose" = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      String
-> ExpBase Info VName
-> (VName -> InternaliseM BasicOp)
-> InternaliseM Result
internaliseOperation String
desc ExpBase Info VName
e ((VName -> InternaliseM BasicOp) -> InternaliseM Result)
-> (VName -> InternaliseM BasicOp) -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ \VName
v -> do
        Int
r <- Type -> Int
forall shape u. ArrayShape shape => TypeBase shape u -> Int
I.arrayRank (Type -> Int) -> InternaliseM Type -> InternaliseM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
v
        BasicOp -> InternaliseM BasicOp
forall (m :: * -> *) a. Monad m => a -> m a
return (BasicOp -> InternaliseM BasicOp)
-> BasicOp -> InternaliseM BasicOp
forall a b. (a -> b) -> a -> b
$ [Int] -> VName -> BasicOp
I.Rearrange ([Int
1, Int
0] [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
2 .. Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1]) VName
v
    handleRest [TupLit [ExpBase Info VName
x, ExpBase Info VName
y] SrcLoc
_] String
"zip" = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ \String
desc ->
      (VName -> InternaliseM SubExp) -> [VName] -> InternaliseM Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"zip_copy" (ExpT SOACS -> InternaliseM SubExp)
-> (VName -> ExpT SOACS) -> VName -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS)
-> (VName -> BasicOp) -> VName -> ExpT SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> BasicOp
Copy)
        ([VName] -> InternaliseM Result)
-> InternaliseM [VName] -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ( [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
(++)
                ([VName] -> [VName] -> [VName])
-> InternaliseM [VName] -> InternaliseM ([VName] -> [VName])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars (String
desc String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_zip_x") ExpBase Info VName
x
                InternaliseM ([VName] -> [VName])
-> InternaliseM [VName] -> InternaliseM [VName]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars (String
desc String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_zip_y") ExpBase Info VName
y
            )
    handleRest [ExpBase Info VName
x] String
"unzip" = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ (String -> ExpBase Info VName -> InternaliseM Result)
-> ExpBase Info VName -> String -> InternaliseM Result
forall a b c. (a -> b -> c) -> b -> a -> c
flip String -> ExpBase Info VName -> InternaliseM Result
internaliseExp ExpBase Info VName
x
    handleRest [ExpBase Info VName
x] String
"trace" = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ (String -> ExpBase Info VName -> InternaliseM Result)
-> ExpBase Info VName -> String -> InternaliseM Result
forall a b c. (a -> b -> c) -> b -> a -> c
flip String -> ExpBase Info VName -> InternaliseM Result
internaliseExp ExpBase Info VName
x
    handleRest [ExpBase Info VName
x] String
"break" = (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a. a -> Maybe a
Just ((String -> InternaliseM Result)
 -> Maybe (String -> InternaliseM Result))
-> (String -> InternaliseM Result)
-> Maybe (String -> InternaliseM Result)
forall a b. (a -> b) -> a -> b
$ (String -> ExpBase Info VName -> InternaliseM Result)
-> ExpBase Info VName -> String -> InternaliseM Result
forall a b c. (a -> b -> c) -> b -> a -> c
flip String -> ExpBase Info VName -> InternaliseM Result
internaliseExp ExpBase Info VName
x
    handleRest [ExpBase Info VName]
_ String
_ = Maybe (String -> InternaliseM Result)
forall a. Maybe a
Nothing

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

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

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

    scatterF :: Int
-> ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> String
-> InternaliseM Result
scatterF Int
dim ExpBase Info VName
a ExpBase Info VName
si ExpBase Info VName
v String
desc = do
      [VName]
si' <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"write_arg_i" ExpBase Info VName
si
      [VName]
svs <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"write_arg_v" ExpBase Info VName
v
      [VName]
sas <- String -> ExpBase Info VName -> InternaliseM [VName]
internaliseExpToVars String
"write_arg_a" ExpBase Info VName
a

      SubExp
si_w <- Int -> [Type] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
I.arraysSize Int
0 ([Type] -> SubExp) -> InternaliseM [Type] -> InternaliseM SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM Type) -> [VName] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType [VName]
si'
      [Type]
sv_ts <- (VName -> InternaliseM Type) -> [VName] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType [VName]
svs

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

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

      [Type]
indexType <- (Type -> Type) -> [Type] -> [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
rowType ([Type] -> [Type]) -> InternaliseM [Type] -> InternaliseM [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> InternaliseM Type) -> [VName] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType [VName]
si'
      [VName]
indexName <- (Type -> InternaliseM VName) -> [Type] -> InternaliseM [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\Type
_ -> String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"write_index") [Type]
indexType
      [VName]
valueNames <- Int -> InternaliseM VName -> InternaliseM [VName]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM ([Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
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"

      [Type]
sa_ts <- (VName -> InternaliseM Type) -> [VName] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> InternaliseM Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType [VName]
sas
      let bodyTypes :: [Type]
bodyTypes = [[Type]] -> [Type]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Int -> [Type] -> [[Type]]
forall a. Int -> a -> [a]
replicate ([Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
sv_ts) [Type]
indexType) [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Type -> Type
forall shape u.
ArrayShape shape =>
Int -> TypeBase shape u -> TypeBase shape u
I.stripArray Int
dim) [Type]
sa_ts
          paramTypes :: [Type]
paramTypes = [Type]
indexType [Type] -> [Type] -> [Type]
forall a. Semigroup a => a -> a -> a
<> (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
rowType [Type]
sv_ts
          bodyNames :: [VName]
bodyNames = [VName]
indexName [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
valueNames
          bodyParams :: [Param Type]
bodyParams = (VName -> Type -> Param Type) -> [VName] -> [Type] -> [Param Type]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith VName -> Type -> Param Type
forall dec. VName -> dec -> Param dec
I.Param [VName]
bodyNames [Type]
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
body <- Scope SOACS -> InternaliseM Body -> InternaliseM Body
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param Type] -> Scope SOACS
forall lore dec.
(LParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
scopeOfLParams [Param Type]
bodyParams) (InternaliseM Body -> InternaliseM Body)
-> InternaliseM Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$
        InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
m (Body (Lore m)) -> m (Body (Lore m))
insertStmsM (InternaliseM (Body (Lore InternaliseM))
 -> InternaliseM (Body (Lore InternaliseM)))
-> InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
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
          Result
results <- [VName] -> (VName -> InternaliseM SubExp) -> InternaliseM Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
outs ((VName -> InternaliseM SubExp) -> InternaliseM Result)
-> (VName -> InternaliseM SubExp) -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$ \VName
name ->
            String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"write_res" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
name
          ErrorMsg SubExp -> SrcLoc -> [Type] -> Body -> InternaliseM Body
ensureResultShape
            ErrorMsg SubExp
"scatter value has wrong size"
            SrcLoc
loc
            [Type]
bodyTypes
            (Body -> InternaliseM Body) -> Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$ Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody Result
results

      let lam :: Lambda
lam =
            Lambda :: forall lore. [LParam lore] -> BodyT lore -> [Type] -> LambdaT lore
I.Lambda
              { lambdaParams :: [LParam]
I.lambdaParams = [Param Type]
[LParam]
bodyParams,
                lambdaReturnType :: [Type]
I.lambdaReturnType = [Type]
bodyTypes,
                lambdaBody :: Body
I.lambdaBody = Body
body
              }
          sivs :: [VName]
sivs = [VName]
si' [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
svs'

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

funcall ::
  String ->
  QualName VName ->
  [SubExp] ->
  SrcLoc ->
  InternaliseM ([SubExp], [I.ExtType])
funcall :: String
-> QualName VName
-> Result
-> SrcLoc
-> InternaliseM (Result, [ExtType])
funcall String
desc (QualName [VName]
_ VName
fname) Result
args SrcLoc
loc = do
  ([VName]
shapes, [DeclType]
value_paramts, [Param DeclType]
fun_params, [(SubExp, Type)] -> Maybe [TypeBase ExtShape Uniqueness]
rettype_fun) <-
    VName -> InternaliseM FunInfo
lookupFunction VName
fname
  [Type]
argts <- (SubExp -> InternaliseM Type) -> Result -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType Result
args

  Result
shapeargs <- [VName] -> [FParam] -> [Type] -> InternaliseM Result
forall (m :: * -> *).
(HasScope SOACS m, Monad m) =>
[VName] -> [FParam] -> [Type] -> m Result
argShapes [VName]
shapes [Param DeclType]
[FParam]
fun_params [Type]
argts
  let diets :: [Diet]
diets =
        Int -> Diet -> [Diet]
forall a. Int -> a -> [a]
replicate (Result -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Result
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
  Result
args' <-
    ErrorMsg SubExp
-> SrcLoc -> [VName] -> [Type] -> Result -> InternaliseM Result
forall u.
Typed (TypeBase Shape u) =>
ErrorMsg SubExp
-> SrcLoc
-> [VName]
-> [TypeBase Shape u]
-> Result
-> InternaliseM Result
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 -> Type) -> [Param DeclType] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> Type
forall dec. Typed dec => Param dec -> Type
I.paramType [Param DeclType]
fun_params)
      (Result
shapeargs Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ Result
args)
  [Type]
argts' <- (SubExp -> InternaliseM Type) -> Result -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType Result
args'
  case [(SubExp, Type)] -> Maybe [TypeBase ExtShape Uniqueness]
rettype_fun ([(SubExp, Type)] -> Maybe [TypeBase ExtShape Uniqueness])
-> [(SubExp, Type)] -> Maybe [TypeBase ExtShape Uniqueness]
forall a b. (a -> b) -> a -> b
$ Result -> [Type] -> [(SubExp, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip Result
args' [Type]
argts' of
    Maybe [TypeBase ExtShape Uniqueness]
Nothing ->
      String -> InternaliseM (Result, [ExtType])
forall a. HasCallStack => String -> a
error (String -> InternaliseM (Result, [ExtType]))
-> String -> InternaliseM (Result, [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 (Result -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Result
args'),
            String
" arguments\n ",
            Result -> String
forall a. Pretty a => a -> String
pretty Result
args',
            String
"\nof types\n ",
            [Type] -> String
forall a. Pretty a => a -> String
pretty [Type]
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
      Result
ses <-
        Attrs -> InternaliseM Result -> InternaliseM Result
forall (m :: * -> *) a. MonadBinder m => Attrs -> m a -> m a
attributing Attrs
attrs (InternaliseM Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
          String -> Exp (Lore InternaliseM) -> InternaliseM Result
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m Result
letTupExp' String
desc (Exp (Lore InternaliseM) -> InternaliseM Result)
-> Exp (Lore InternaliseM) -> InternaliseM Result
forall a b. (a -> b) -> a -> b
$
            Name
-> [(SubExp, Diet)]
-> [RetType SOACS]
-> (Safety, SrcLoc, [SrcLoc])
-> ExpT SOACS
forall lore.
Name
-> [(SubExp, Diet)]
-> [RetType lore]
-> (Safety, SrcLoc, [SrcLoc])
-> ExpT lore
I.Apply (VName -> Name
internaliseFunName VName
fname) (Result -> [Diet] -> [(SubExp, Diet)]
forall a b. [a] -> [b] -> [(a, b)]
zip Result
args' [Diet]
diets) [TypeBase ExtShape Uniqueness]
[RetType SOACS]
ts (Safety
safety, SrcLoc
loc, [SrcLoc]
forall a. Monoid a => a
mempty)
      (Result, [ExtType]) -> InternaliseM (Result, [ExtType])
forall (m :: * -> *) a. Monad m => a -> m a
return (Result
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 :: E.StructType -> [VName] -> [SubExp] -> InternaliseM ()
bindExtSizes :: StructType -> [VName] -> Result -> InternaliseM ()
bindExtSizes StructType
ret [VName]
retext Result
ses = do
  [TypeBase ExtShape Uniqueness]
ts <- StructType -> InternaliseM [TypeBase ExtShape Uniqueness]
internaliseType StructType
ret
  [Type]
ses_ts <- (SubExp -> InternaliseM Type) -> Result -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType Result
ses

  let combine :: TypeBase ExtShape Uniqueness -> Type -> Map VName SubExp
combine TypeBase ExtShape Uniqueness
t1 Type
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] -> Result -> [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) (Type -> Result
forall u. TypeBase Shape u -> Result
arrayDims Type
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 -> Type -> Map VName SubExp)
-> [TypeBase ExtShape Uniqueness] -> [Type] -> [Map VName SubExp]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeBase ExtShape Uniqueness -> Type -> Map VName SubExp
combine [TypeBase ExtShape Uniqueness]
ts [Type]
ses_ts) (((VName, SubExp) -> InternaliseM ()) -> InternaliseM ())
-> ((VName, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \(VName
v, SubExp
se) ->
    [VName] -> Exp (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [VName
v] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se

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

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

  [Param Type]
add_lam_x_params <-
    Int -> InternaliseM (Param Type) -> InternaliseM [Param Type]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
k (InternaliseM (Param Type) -> InternaliseM [Param Type])
-> InternaliseM (Param Type) -> InternaliseM [Param Type]
forall a b. (a -> b) -> a -> b
$ VName -> Type -> Param Type
forall dec. VName -> dec -> Param dec
I.Param (VName -> Type -> Param Type)
-> InternaliseM VName -> InternaliseM (Type -> Param Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"x" InternaliseM (Type -> Param Type)
-> InternaliseM Type -> InternaliseM (Param Type)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> InternaliseM Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
  [Param Type]
add_lam_y_params <-
    Int -> InternaliseM (Param Type) -> InternaliseM [Param Type]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
k (InternaliseM (Param Type) -> InternaliseM [Param Type])
-> InternaliseM (Param Type) -> InternaliseM [Param Type]
forall a b. (a -> b) -> a -> b
$ VName -> Type -> Param Type
forall dec. VName -> dec -> Param dec
I.Param (VName -> Type -> Param Type)
-> InternaliseM VName -> InternaliseM (Type -> Param Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"y" InternaliseM (Type -> Param Type)
-> InternaliseM Type -> InternaliseM (Param Type)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> InternaliseM Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
  Body
add_lam_body <- Binder SOACS Body -> InternaliseM Body
forall lore (m :: * -> *) somelore.
(Bindable lore, MonadFreshNames m, HasScope somelore m,
 SameScope somelore lore) =>
Binder lore (Body lore) -> m (Body lore)
runBodyBinder (Binder SOACS Body -> InternaliseM Body)
-> Binder SOACS Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$
    Scope SOACS -> Binder SOACS Body -> Binder SOACS Body
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param Type] -> Scope SOACS
forall lore dec.
(LParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
scopeOfLParams ([Param Type] -> Scope SOACS) -> [Param Type] -> Scope SOACS
forall a b. (a -> b) -> a -> b
$ [Param Type]
add_lam_x_params [Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
++ [Param Type]
add_lam_y_params) (Binder SOACS Body -> Binder SOACS Body)
-> Binder SOACS Body -> Binder SOACS Body
forall a b. (a -> b) -> a -> b
$
      (Result -> Body)
-> BinderT SOACS (State VNameSource) Result -> Binder SOACS Body
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody (BinderT SOACS (State VNameSource) Result -> Binder SOACS Body)
-> BinderT SOACS (State VNameSource) Result -> Binder SOACS Body
forall a b. (a -> b) -> a -> b
$
        [(Param Type, Param Type)]
-> ((Param Type, Param Type)
    -> BinderT SOACS (State VNameSource) SubExp)
-> BinderT SOACS (State VNameSource) Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([Param Type] -> [Param Type] -> [(Param Type, Param Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param Type]
add_lam_x_params [Param Type]
add_lam_y_params) (((Param Type, Param Type)
  -> BinderT SOACS (State VNameSource) SubExp)
 -> BinderT SOACS (State VNameSource) Result)
-> ((Param Type, Param Type)
    -> BinderT SOACS (State VNameSource) SubExp)
-> BinderT SOACS (State VNameSource) Result
forall a b. (a -> b) -> a -> b
$ \(Param Type
x, Param Type
y) ->
          String
-> Exp (Lore (BinderT SOACS (State VNameSource)))
-> BinderT SOACS (State VNameSource) SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"z" (Exp (Lore (BinderT SOACS (State VNameSource)))
 -> BinderT SOACS (State VNameSource) SubExp)
-> Exp (Lore (BinderT SOACS (State VNameSource)))
-> BinderT SOACS (State VNameSource) SubExp
forall a b. (a -> b) -> a -> b
$
            BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
              BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp
                (IntType -> Overflow -> BinOp
I.Add IntType
Int64 Overflow
I.OverflowUndef)
                (VName -> SubExp
I.Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param Type -> VName
forall dec. Param dec -> VName
I.paramName Param Type
x)
                (VName -> SubExp
I.Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param Type -> VName
forall dec. Param dec -> VName
I.paramName Param Type
y)
  let add_lam :: Lambda
add_lam =
        Lambda :: forall lore. [LParam lore] -> BodyT lore -> [Type] -> LambdaT lore
I.Lambda
          { lambdaBody :: Body
I.lambdaBody = Body
add_lam_body,
            lambdaParams :: [LParam]
I.lambdaParams = [Param Type]
add_lam_x_params [Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
++ [Param Type]
add_lam_y_params,
            lambdaReturnType :: [Type]
I.lambdaReturnType = Int -> Type -> [Type]
forall a. Int -> a -> [a]
replicate Int
k (Type -> [Type]) -> Type -> [Type]
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64
          }
      nes :: Result
nes = Int -> SubExp -> Result
forall a. Int -> a -> [a]
replicate ([VName] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
increments) (SubExp -> Result) -> SubExp -> Result
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0

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

  -- We have the offsets for each of the partitions, but we also need
  -- the total sizes, which are the last elements in the offests.  We
  -- just have to be careful in case the array is empty.
  SubExp
last_index <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"last_index" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
Int64 Overflow
OverflowUndef) SubExp
w (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
1 :: Int64)
  Body
nonempty_body <- Binder SOACS Body -> InternaliseM Body
forall lore (m :: * -> *) somelore.
(Bindable lore, MonadFreshNames m, HasScope somelore m,
 SameScope somelore lore) =>
Binder lore (Body lore) -> m (Body lore)
runBodyBinder (Binder SOACS Body -> InternaliseM Body)
-> Binder SOACS Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$
    (Result -> Body)
-> BinderT SOACS (State VNameSource) Result -> Binder SOACS Body
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody (BinderT SOACS (State VNameSource) Result -> Binder SOACS Body)
-> BinderT SOACS (State VNameSource) Result -> Binder SOACS Body
forall a b. (a -> b) -> a -> b
$
      [VName]
-> (VName -> BinderT SOACS (State VNameSource) SubExp)
-> BinderT SOACS (State VNameSource) Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
all_offsets ((VName -> BinderT SOACS (State VNameSource) SubExp)
 -> BinderT SOACS (State VNameSource) Result)
-> (VName -> BinderT SOACS (State VNameSource) SubExp)
-> BinderT SOACS (State VNameSource) Result
forall a b. (a -> b) -> a -> b
$ \VName
offset_array ->
        String
-> Exp (Lore (BinderT SOACS (State VNameSource)))
-> BinderT SOACS (State VNameSource) SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"last_offset" (Exp (Lore (BinderT SOACS (State VNameSource)))
 -> BinderT SOACS (State VNameSource) SubExp)
-> Exp (Lore (BinderT SOACS (State VNameSource)))
-> BinderT SOACS (State VNameSource) SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ VName -> [DimIndex SubExp] -> BasicOp
I.Index VName
offset_array [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
I.DimFix SubExp
last_index]
  let empty_body :: Body
empty_body = Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody (Result -> Body) -> Result -> Body
forall a b. (a -> b) -> a -> b
$ Int -> SubExp -> Result
forall a. Int -> a -> [a]
replicate Int
k (SubExp -> Result) -> SubExp -> Result
forall a b. (a -> b) -> a -> b
$ Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)
  SubExp
is_empty <- String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"is_empty" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
CmpEq PrimType
int64) SubExp
w (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)
  [VName]
sizes <-
    String -> Exp (Lore InternaliseM) -> InternaliseM [VName]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [VName]
letTupExp String
"partition_size" (Exp (Lore InternaliseM) -> InternaliseM [VName])
-> Exp (Lore InternaliseM) -> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$
      SubExp -> Body -> Body -> IfDec (BranchType SOACS) -> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If SubExp
is_empty Body
empty_body Body
nonempty_body (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
        [Type] -> IfDec ExtType
ifCommon ([Type] -> IfDec ExtType) -> [Type] -> IfDec ExtType
forall a b. (a -> b) -> a -> b
$ Int -> Type -> [Type]
forall a. Int -> a -> [a]
replicate Int
k (Type -> [Type]) -> Type -> [Type]
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64

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

  -- Create scratch arrays for the result.
  [VName]
blanks <- [Type] -> (Type -> InternaliseM VName) -> InternaliseM [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Type]
arr_ts ((Type -> InternaliseM VName) -> InternaliseM [VName])
-> (Type -> InternaliseM VName) -> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$ \Type
arr_t ->
    String -> Exp (Lore InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m VName
letExp String
"partition_dest" (Exp (Lore InternaliseM) -> InternaliseM VName)
-> Exp (Lore InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
        PrimType -> Result -> BasicOp
Scratch (Type -> PrimType
forall shape u. TypeBase shape u -> PrimType
elemType Type
arr_t) (SubExp
w SubExp -> Result -> Result
forall a. a -> [a] -> [a]
: Int -> Result -> Result
forall a. Int -> [a] -> [a]
drop Int
1 (Type -> Result
forall u. TypeBase Shape u -> Result
I.arrayDims Type
arr_t))

  -- Now write into the result.
  Lambda
write_lam <- do
    Param Type
c_param <- VName -> Type -> Param Type
forall dec. VName -> dec -> Param dec
I.Param (VName -> Type -> Param Type)
-> InternaliseM VName -> InternaliseM (Type -> Param Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"c" InternaliseM (Type -> Param Type)
-> InternaliseM Type -> InternaliseM (Param Type)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> InternaliseM Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
    [Param Type]
offset_params <- Int -> InternaliseM (Param Type) -> InternaliseM [Param Type]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
k (InternaliseM (Param Type) -> InternaliseM [Param Type])
-> InternaliseM (Param Type) -> InternaliseM [Param Type]
forall a b. (a -> b) -> a -> b
$ VName -> Type -> Param Type
forall dec. VName -> dec -> Param dec
I.Param (VName -> Type -> Param Type)
-> InternaliseM VName -> InternaliseM (Type -> Param Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"offset" InternaliseM (Type -> Param Type)
-> InternaliseM Type -> InternaliseM (Param Type)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> InternaliseM Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
    [Param Type]
value_params <- [Type]
-> (Type -> InternaliseM (Param Type)) -> InternaliseM [Param Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Type]
arr_ts ((Type -> InternaliseM (Param Type)) -> InternaliseM [Param Type])
-> (Type -> InternaliseM (Param Type)) -> InternaliseM [Param Type]
forall a b. (a -> b) -> a -> b
$ \Type
arr_t ->
      VName -> Type -> Param Type
forall dec. VName -> dec -> Param dec
I.Param (VName -> Type -> Param Type)
-> InternaliseM VName -> InternaliseM (Type -> Param Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"v" InternaliseM (Type -> Param Type)
-> InternaliseM Type -> InternaliseM (Param Type)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> InternaliseM Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
I.rowType Type
arr_t)
    (SubExp
offset, Stms SOACS
offset_stms) <-
      InternaliseM SubExp
-> InternaliseM (SubExp, Stms (Lore InternaliseM))
forall (m :: * -> *) a.
MonadBinder m =>
m a -> m (a, Stms (Lore m))
collectStms (InternaliseM SubExp
 -> InternaliseM (SubExp, Stms (Lore InternaliseM)))
-> InternaliseM SubExp
-> InternaliseM (SubExp, Stms (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$
        Result -> SubExp -> Int -> [LParam] -> InternaliseM SubExp
mkOffsetLambdaBody
          ((VName -> SubExp) -> [VName] -> Result
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 Type -> VName
forall dec. Param dec -> VName
I.paramName Param Type
c_param)
          Int
0
          [Param Type]
[LParam]
offset_params
    Lambda -> InternaliseM Lambda
forall (m :: * -> *) a. Monad m => a -> m a
return
      Lambda :: forall lore. [LParam lore] -> BodyT lore -> [Type] -> LambdaT lore
I.Lambda
        { lambdaParams :: [LParam]
I.lambdaParams = Param Type
c_param Param Type -> [Param Type] -> [Param Type]
forall a. a -> [a] -> [a]
: [Param Type]
offset_params [Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
++ [Param Type]
value_params,
          lambdaReturnType :: [Type]
I.lambdaReturnType =
            Int -> Type -> [Type]
forall a. Int -> a -> [a]
replicate ([Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
arr_ts) (PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
              [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
I.rowType [Type]
arr_ts,
          lambdaBody :: Body
I.lambdaBody =
            Stms SOACS -> Result -> Body
forall lore. Bindable lore => Stms lore -> Result -> Body lore
mkBody Stms SOACS
offset_stms (Result -> Body) -> Result -> Body
forall a b. (a -> b) -> a -> b
$
              Int -> SubExp -> Result
forall a. Int -> a -> [a]
replicate ([Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
arr_ts) SubExp
offset
                Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ (Param Type -> SubExp) -> [Param Type] -> Result
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
I.Var (VName -> SubExp) -> (Param Type -> VName) -> Param Type -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> VName
forall dec. Param dec -> VName
I.paramName) [Param Type]
value_params
        }
  [VName]
results <-
    String -> Exp (Lore InternaliseM) -> InternaliseM [VName]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [VName]
letTupExp String
"partition_res" (Exp (Lore InternaliseM) -> InternaliseM [VName])
-> Exp (Lore InternaliseM) -> InternaliseM [VName]
forall a b. (a -> b) -> a -> b
$
      Op SOACS -> ExpT SOACS
forall lore. Op lore -> ExpT lore
I.Op (Op SOACS -> ExpT SOACS) -> Op SOACS -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
        SubExp -> Lambda -> [VName] -> [(Shape, Int, VName)] -> SOAC SOACS
forall lore.
SubExp
-> Lambda lore -> [VName] -> [(Shape, Int, VName)] -> SOAC lore
I.Scatter
          SubExp
w
          Lambda
write_lam
          (VName
classes VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: [VName]
all_offsets [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
arrs)
          ([(Shape, Int, VName)] -> SOAC SOACS)
-> [(Shape, Int, VName)] -> SOAC SOACS
forall a b. (a -> b) -> a -> b
$ [Shape] -> [Int] -> [VName] -> [(Shape, Int, VName)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 (Shape -> [Shape]
forall a. a -> [a]
repeat (Shape -> [Shape]) -> Shape -> [Shape]
forall a b. (a -> b) -> a -> b
$ Result -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
w]) (Int -> [Int]
forall a. a -> [a]
repeat Int
1) [VName]
blanks
  SubExp
sizes' <-
    String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"partition_sizes" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
        Result -> Type -> BasicOp
I.ArrayLit ((VName -> SubExp) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var [VName]
sizes) (Type -> BasicOp) -> Type -> BasicOp
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64
  (Result, Result) -> InternaliseM (Result, Result)
forall (m :: * -> *) a. Monad m => a -> m a
return ((VName -> SubExp) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var [VName]
results, [SubExp
sizes'])
  where
    mkOffsetLambdaBody ::
      [SubExp] ->
      SubExp ->
      Int ->
      [I.LParam] ->
      InternaliseM SubExp
    mkOffsetLambdaBody :: Result -> SubExp -> Int -> [LParam] -> InternaliseM SubExp
mkOffsetLambdaBody Result
_ SubExp
_ Int
_ [] =
      SubExp -> InternaliseM SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> InternaliseM SubExp) -> SubExp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (-Int64
1 :: Int64)
    mkOffsetLambdaBody Result
sizes SubExp
c Int
i (LParam
p : [LParam]
ps) = do
      SubExp
is_this_one <-
        String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"is_this_one" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
CmpEq PrimType
int64) SubExp
c (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$
              IntType -> Integer -> SubExp
intConst IntType
Int64 (Integer -> SubExp) -> Integer -> SubExp
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
i
      SubExp
next_one <- Result -> SubExp -> Int -> [LParam] -> InternaliseM SubExp
mkOffsetLambdaBody Result
sizes SubExp
c (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [LParam]
ps
      SubExp
this_one <-
        String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"this_offset"
          (ExpT SOACS -> InternaliseM SubExp)
-> InternaliseM (ExpT SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinOp -> SubExp -> Result -> InternaliseM (Exp (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> SubExp -> Result -> m (Exp (Lore m))
foldBinOp
            (IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
OverflowUndef)
            (Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (-Int64
1 :: Int64))
            (VName -> SubExp
I.Var (Param Type -> VName
forall dec. Param dec -> VName
I.paramName Param Type
LParam
p) SubExp -> Result -> Result
forall a. a -> [a] -> [a]
: Int -> Result -> Result
forall a. Int -> [a] -> [a]
take Int
i Result
sizes)
      String -> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"total_res" (Exp (Lore InternaliseM) -> InternaliseM SubExp)
-> Exp (Lore InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
        SubExp -> Body -> Body -> IfDec (BranchType SOACS) -> ExpT SOACS
forall lore.
SubExp
-> BodyT lore -> BodyT lore -> IfDec (BranchType lore) -> ExpT lore
I.If
          SubExp
is_this_one
          (Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody [SubExp
this_one])
          (Result -> Body
forall lore. Bindable lore => Result -> Body lore
resultBody [SubExp
next_one])
          (IfDec (BranchType SOACS) -> ExpT SOACS)
-> IfDec (BranchType SOACS) -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ [Type] -> IfDec ExtType
ifCommon [PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64]

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

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

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