{-# LANGUAGE Strict #-}
{-# LANGUAGE TypeFamilies #-}

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

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

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

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

internaliseFunName :: VName -> Name
internaliseFunName :: VName -> Name
internaliseFunName = [Char] -> Name
nameFromString ([Char] -> Name) -> (VName -> [Char]) -> VName -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> [Char]
forall a. Pretty a => a -> [Char]
prettyString

shiftRetAls :: Int -> RetAls -> RetAls
shiftRetAls :: Int -> RetAls -> RetAls
shiftRetAls Int
d (RetAls [Int]
pals [Int]
rals) = [Int] -> [Int] -> RetAls
RetAls [Int]
pals ([Int] -> RetAls) -> [Int] -> RetAls
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
d) [Int]
rals

internaliseValBind :: VisibleTypes -> E.ValBind -> InternaliseM ()
internaliseValBind :: VisibleTypes -> ValBind -> InternaliseM ()
internaliseValBind VisibleTypes
types fb :: ValBind
fb@(E.ValBind Maybe (Info EntryPoint)
entry VName
fname Maybe (TypeExp Exp VName)
_ (Info ResRetType
rettype) [TypeParamBase VName]
tparams [PatBase Info VName ParamType]
params Exp
body Maybe DocComment
_ [AttrInfo VName]
attrs SrcLoc
_) = do
  [TypeParamBase VName]
-> [PatBase Info VName ParamType]
-> ([FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM ())
-> InternaliseM ()
forall a.
[TypeParamBase VName]
-> [PatBase Info VName ParamType]
-> ([FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM a)
-> InternaliseM a
bindingFParams [TypeParamBase VName]
tparams [PatBase Info VName ParamType]
params (([FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM ())
 -> InternaliseM ())
-> ([FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM ())
-> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \[FParam SOACS]
shapeparams [[Tree (FParam SOACS)]]
params' -> do
    let shapenames :: [VName]
shapenames = (Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
[FParam SOACS]
shapeparams
        all_params :: [Free [] (Param DeclType)]
all_params = (Param DeclType -> Free [] (Param DeclType))
-> [Param DeclType] -> [Free [] (Param DeclType)]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> Free [] (Param DeclType)
forall a. a -> Free [] a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Param DeclType]
[FParam SOACS]
shapeparams [Free [] (Param DeclType)]
-> [Free [] (Param DeclType)] -> [Free [] (Param DeclType)]
forall a. [a] -> [a] -> [a]
++ [[Free [] (Param DeclType)]] -> [Free [] (Param DeclType)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Free [] (Param DeclType)]]
[[Tree (FParam SOACS)]]
params'
        msg :: ErrorMsg SubExp
msg =
          [ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg
            [ ErrorMsgPart SubExp
"Internal runtime error.\n",
              ErrorMsgPart SubExp
"Return value of ",
              Text -> ErrorMsgPart SubExp
forall a. Text -> ErrorMsgPart a
ErrorString (VName -> Text
forall a. Pretty a => a -> Text
prettyText VName
fname),
              ErrorMsgPart SubExp
" does not match type shape.\n",
              ErrorMsgPart SubExp
"This is a bug in the Futhark compiler. Please report this:\n",
              ErrorMsgPart SubExp
"  https://github.com/diku-dk/futhark/issues"
            ]

    (body', rettype') <- InternaliseM (Result, [(TypeBase ExtShape Uniqueness, RetAls)])
-> InternaliseM
     (Body (Rep InternaliseM), [(TypeBase ExtShape Uniqueness, RetAls)])
forall (m :: * -> *) a.
MonadBuilder m =>
m (Result, a) -> m (Body (Rep m), a)
buildBody (InternaliseM (Result, [(TypeBase ExtShape Uniqueness, RetAls)])
 -> InternaliseM
      (Body (Rep InternaliseM),
       [(TypeBase ExtShape Uniqueness, RetAls)]))
-> InternaliseM (Result, [(TypeBase ExtShape Uniqueness, RetAls)])
-> InternaliseM
     (Body (Rep InternaliseM), [(TypeBase ExtShape Uniqueness, RetAls)])
forall a b. (a -> b) -> a -> b
$ do
      body_res <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp (VName -> [Char]
baseString VName
fname [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"_res") Exp
body
      (rettype', retals) <-
        first zeroExts . unzip . internaliseReturnType (map (fmap paramDeclType) all_params) rettype
          <$> mapM subExpType body_res

      when (null params') $
        bindExtSizes (E.AppRes (E.toStruct $ E.retType rettype) (E.retDims rettype)) body_res

      body_res' <-
        ensureResultExtShape msg (map I.fromDecl rettype') $ subExpsRes body_res
      let num_ctx = Set Int -> Int
forall a. Set a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([TypeBase ExtShape Uniqueness] -> Set Int
forall u. [TypeBase ExtShape u] -> Set Int
shapeContext [TypeBase ExtShape Uniqueness]
rettype')
      pure
        ( body_res',
          replicate num_ctx (I.Prim int64, mempty)
            ++ zip rettype' (map (shiftRetAls num_ctx) retals)
        )

    attrs' <- internaliseAttrs attrs

    let fd =
          Maybe EntryPoint
-> Attrs
-> Name
-> [(RetType SOACS, RetAls)]
-> [FParam SOACS]
-> Body SOACS
-> FunDef SOACS
forall rep.
Maybe EntryPoint
-> Attrs
-> Name
-> [(RetType rep, RetAls)]
-> [FParam rep]
-> Body rep
-> FunDef rep
I.FunDef
            Maybe EntryPoint
forall a. Maybe a
Nothing
            Attrs
attrs'
            (VName -> Name
internaliseFunName VName
fname)
            [(TypeBase ExtShape Uniqueness, RetAls)]
[(RetType SOACS, RetAls)]
rettype'
            ((Free [] (Param DeclType) -> [Param DeclType])
-> [Free [] (Param DeclType)] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Free [] (Param DeclType) -> [Param DeclType]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList [Free [] (Param DeclType)]
all_params)
            Body SOACS
body'

    if null params'
      then bindConstant fname fd
      else
        bindFunction
          fname
          fd
          ( shapenames,
            map declTypeOf $ foldMap (foldMap toList) params',
            foldMap toList all_params,
            fmap (`zip` map snd rettype')
              . applyRetType (map fst rettype') (foldMap toList all_params)
          )

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

generateEntryPoint :: VisibleTypes -> E.EntryPoint -> E.ValBind -> InternaliseM ()
generateEntryPoint :: VisibleTypes -> EntryPoint -> ValBind -> InternaliseM ()
generateEntryPoint VisibleTypes
types (E.EntryPoint [EntryParam]
e_params EntryType
e_rettype) ValBind
vb = do
  let (E.ValBind Maybe (Info EntryPoint)
_ VName
ofname Maybe (TypeExp Exp VName)
_ (Info ResRetType
rettype) [TypeParamBase VName]
tparams [PatBase Info VName ParamType]
params Exp
_ Maybe DocComment
_ [AttrInfo VName]
attrs SrcLoc
_) = ValBind
vb
  [TypeParamBase VName]
-> [PatBase Info VName ParamType]
-> ([FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM ())
-> InternaliseM ()
forall a.
[TypeParamBase VName]
-> [PatBase Info VName ParamType]
-> ([FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM a)
-> InternaliseM a
bindingFParams [TypeParamBase VName]
tparams [PatBase Info VName ParamType]
params (([FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM ())
 -> InternaliseM ())
-> ([FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM ())
-> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \[FParam SOACS]
shapeparams [[Tree (FParam SOACS)]]
params' -> do
    let all_params :: [Free [] (Param DeclType)]
all_params = (Param DeclType -> Free [] (Param DeclType))
-> [Param DeclType] -> [Free [] (Param DeclType)]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> Free [] (Param DeclType)
forall a. a -> Free [] a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Param DeclType]
[FParam SOACS]
shapeparams [Free [] (Param DeclType)]
-> [Free [] (Param DeclType)] -> [Free [] (Param DeclType)]
forall a. [a] -> [a] -> [a]
++ [[Free [] (Param DeclType)]] -> [Free [] (Param DeclType)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Free [] (Param DeclType)]]
[[Tree (FParam SOACS)]]
params'
        ([[TypeBase ExtShape Uniqueness]]
entry_rettype, [[RetAls]]
retals) =
          [([TypeBase ExtShape Uniqueness], [RetAls])]
-> ([[TypeBase ExtShape Uniqueness]], [[RetAls]])
forall a b. [(a, b)] -> ([a], [b])
unzip ([([TypeBase ExtShape Uniqueness], [RetAls])]
 -> ([[TypeBase ExtShape Uniqueness]], [[RetAls]]))
-> [([TypeBase ExtShape Uniqueness], [RetAls])]
-> ([[TypeBase ExtShape Uniqueness]], [[RetAls]])
forall a b. (a -> b) -> a -> b
$ ([(TypeBase ExtShape Uniqueness, RetAls)]
 -> ([TypeBase ExtShape Uniqueness], [RetAls]))
-> [[(TypeBase ExtShape Uniqueness, RetAls)]]
-> [([TypeBase ExtShape Uniqueness], [RetAls])]
forall a b. (a -> b) -> [a] -> [b]
map [(TypeBase ExtShape Uniqueness, RetAls)]
-> ([TypeBase ExtShape Uniqueness], [RetAls])
forall a b. [(a, b)] -> ([a], [b])
unzip ([[(TypeBase ExtShape Uniqueness, RetAls)]]
 -> [([TypeBase ExtShape Uniqueness], [RetAls])])
-> [[(TypeBase ExtShape Uniqueness, RetAls)]]
-> [([TypeBase ExtShape Uniqueness], [RetAls])]
forall a b. (a -> b) -> a -> b
$ [Tree DeclType]
-> ResRetType -> [[(TypeBase ExtShape Uniqueness, RetAls)]]
internaliseEntryReturnType ((Free [] (Param DeclType) -> Tree DeclType)
-> [Free [] (Param DeclType)] -> [Tree DeclType]
forall a b. (a -> b) -> [a] -> [b]
map ((Param DeclType -> DeclType)
-> Free [] (Param DeclType) -> Tree DeclType
forall a b. (a -> b) -> Free [] a -> Free [] b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Param DeclType -> DeclType
forall dec. DeclTyped dec => Param dec -> DeclType
paramDeclType) [Free [] (Param DeclType)]
all_params) ResRetType
rettype
        (EntryPoint
entry', OpaqueTypes
opaques) =
          VisibleTypes
-> Name
-> [(EntryParam, [Param DeclType])]
-> (EntryType, [[TypeBase Rank Uniqueness]])
-> (EntryPoint, OpaqueTypes)
entryPoint
            VisibleTypes
types
            (VName -> Name
baseName VName
ofname)
            ([EntryParam]
-> [[Param DeclType]] -> [(EntryParam, [Param DeclType])]
forall a b. [a] -> [b] -> [(a, b)]
zip [EntryParam]
e_params ([[Param DeclType]] -> [(EntryParam, [Param DeclType])])
-> [[Param DeclType]] -> [(EntryParam, [Param DeclType])]
forall a b. (a -> b) -> a -> b
$ ([Free [] (Param DeclType)] -> [Param DeclType])
-> [[Free [] (Param DeclType)]] -> [[Param DeclType]]
forall a b. (a -> b) -> [a] -> [b]
map ((Free [] (Param DeclType) -> [Param DeclType])
-> [Free [] (Param DeclType)] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Free [] (Param DeclType) -> [Param DeclType]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList) [[Free [] (Param DeclType)]]
[[Tree (FParam SOACS)]]
params')
            (EntryType
e_rettype, ([TypeBase ExtShape Uniqueness] -> [TypeBase Rank Uniqueness])
-> [[TypeBase ExtShape Uniqueness]] -> [[TypeBase Rank Uniqueness]]
forall a b. (a -> b) -> [a] -> [b]
map ((TypeBase ExtShape Uniqueness -> TypeBase Rank Uniqueness)
-> [TypeBase ExtShape Uniqueness] -> [TypeBase Rank Uniqueness]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase ExtShape Uniqueness -> TypeBase Rank Uniqueness
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase Rank u
I.rankShaped) [[TypeBase ExtShape Uniqueness]]
entry_rettype)
        args :: [SubExp]
args = (Param DeclType -> SubExp) -> [Param DeclType] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
I.Var (VName -> SubExp)
-> (Param DeclType -> VName) -> Param DeclType -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName) ([Param DeclType] -> [SubExp]) -> [Param DeclType] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ ([Free [] (Param DeclType)] -> [Param DeclType])
-> [[Free [] (Param DeclType)]] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((Free [] (Param DeclType) -> [Param DeclType])
-> [Free [] (Param DeclType)] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Free [] (Param DeclType) -> [Param DeclType]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList) [[Free [] (Param DeclType)]]
[[Tree (FParam SOACS)]]
params'

    OpaqueTypes -> InternaliseM ()
addOpaques OpaqueTypes
opaques

    (entry_body, ctx_ts) <- InternaliseM (Result, [(TypeBase ExtShape Uniqueness, RetAls)])
-> InternaliseM
     (Body (Rep InternaliseM), [(TypeBase ExtShape Uniqueness, RetAls)])
forall (m :: * -> *) a.
MonadBuilder m =>
m (Result, a) -> m (Body (Rep m), a)
buildBody (InternaliseM (Result, [(TypeBase ExtShape Uniqueness, RetAls)])
 -> InternaliseM
      (Body (Rep InternaliseM),
       [(TypeBase ExtShape Uniqueness, RetAls)]))
-> InternaliseM (Result, [(TypeBase ExtShape Uniqueness, RetAls)])
-> InternaliseM
     (Body (Rep InternaliseM), [(TypeBase ExtShape Uniqueness, RetAls)])
forall a b. (a -> b) -> a -> b
$ do
      -- Special case the (rare) situation where the entry point is
      -- not a function.
      maybe_const <- VName -> InternaliseM (Maybe [SubExp])
lookupConst VName
ofname
      vals <- case maybe_const of
        Just [SubExp]
ses ->
          [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
ses
        Maybe [SubExp]
Nothing ->
          [Char] -> QualName VName -> [SubExp] -> InternaliseM [SubExp]
funcall [Char]
"entry_result" (VName -> QualName VName
forall v. v -> QualName v
E.qualName VName
ofname) [SubExp]
args
      ctx <-
        extractShapeContext (zeroExts $ concat entry_rettype)
          <$> mapM (fmap I.arrayDims . subExpType) vals
      pure (subExpsRes $ ctx ++ vals, map (const (I.Prim int64, mempty)) ctx)

    attrs' <- internaliseAttrs attrs
    let num_ctx = [(TypeBase ExtShape Uniqueness, RetAls)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(TypeBase ExtShape Uniqueness, RetAls)]
ctx_ts
    addFunDef $
      I.FunDef
        (Just entry')
        attrs'
        ("entry_" <> baseName ofname)
        ( ctx_ts
            ++ zip
              (zeroExts (concat entry_rettype))
              (map (shiftRetAls num_ctx) $ concat retals)
        )
        (shapeparams ++ foldMap (foldMap toList) params')
        entry_body
  where
    zeroExts :: [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts [TypeBase ExtShape u]
ts = [TypeBase ExtShape u]
-> [TypeBase ExtShape u] -> [TypeBase ExtShape u]
forall u.
[TypeBase ExtShape u]
-> [TypeBase ExtShape u] -> [TypeBase ExtShape u]
generaliseExtTypes [TypeBase ExtShape u]
ts [TypeBase ExtShape u]
ts

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

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

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

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

internaliseAppExp :: String -> E.AppRes -> E.AppExp -> InternaliseM [I.SubExp]
internaliseAppExp :: [Char] -> AppRes -> AppExp -> InternaliseM [SubExp]
internaliseAppExp [Char]
desc AppRes
_ (E.Index Exp
e SliceBase Info VName
idxs SrcLoc
_) = do
  vs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"indexed" Exp
e
  dims <- case vs of
    [] -> [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [] -- Will this happen?
    VName
v : [VName]
_ -> TypeBase Shape NoUniqueness -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
I.arrayDims (TypeBase Shape NoUniqueness -> [SubExp])
-> InternaliseM (TypeBase Shape NoUniqueness)
-> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
v
  (idxs', cs) <- internaliseSlice dims idxs
  let index VName
v = do
        v_t <- VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
v
        pure $ I.BasicOp $ I.Index v $ fullSlice v_t idxs'
  certifying cs $ mapM (letSubExp desc <=< index) vs
internaliseAppExp [Char]
desc AppRes
_ (E.Range Exp
start Maybe Exp
maybe_second Inclusiveness Exp
end SrcLoc
_) = do
  start' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"range_start" Exp
start
  end' <- internaliseExp1 "range_end" $ case end of
    DownToExclusive Exp
e -> Exp
e
    ToInclusive Exp
e -> Exp
e
    UpToExclusive Exp
e -> Exp
e
  maybe_second' <-
    traverse (internaliseExp1 "range_second") maybe_second

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

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

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

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

  step_sign <- letSubExp "s_sign" $ BasicOp $ I.UnOp (I.SSignum it) step
  step_sign_i64 <- asIntS Int64 step_sign

  bounds_invalid_downwards <-
    letSubExp "bounds_invalid_downwards" $
      I.BasicOp $
        I.CmpOp lt_op start' end'
  bounds_invalid_upwards <-
    letSubExp "bounds_invalid_upwards" $
      I.BasicOp $
        I.CmpOp lt_op end' start'

  (distance, step_wrong_dir, bounds_invalid) <- case end of
    DownToExclusive {} -> do
      step_wrong_dir <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"step_wrong_dir" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
one
      distance <-
        letSubExp "distance" $
          I.BasicOp $
            I.BinOp (Sub it I.OverflowWrap) start' end'
      distance_i64 <- asIntS Int64 distance
      pure (distance_i64, step_wrong_dir, bounds_invalid_downwards)
    UpToExclusive {} -> do
      step_wrong_dir <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"step_wrong_dir" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
negone
      distance <- letSubExp "distance" $ I.BasicOp $ I.BinOp (Sub it I.OverflowWrap) end' start'
      distance_i64 <- asIntS Int64 distance
      pure (distance_i64, step_wrong_dir, bounds_invalid_upwards)
    ToInclusive {} -> do
      downwards <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"downwards" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
negone
      distance_downwards_exclusive <-
        letSubExp "distance_downwards_exclusive" $
          I.BasicOp $
            I.BinOp (Sub it I.OverflowWrap) start' end'
      distance_upwards_exclusive <-
        letSubExp "distance_upwards_exclusive" $
          I.BasicOp $
            I.BinOp (Sub it I.OverflowWrap) end' start'

      bounds_invalid <-
        letSubExp "bounds_invalid"
          =<< eIf
            (eSubExp downwards)
            (resultBodyM [bounds_invalid_downwards])
            (resultBodyM [bounds_invalid_upwards])
      distance_exclusive <-
        letSubExp "distance_exclusive"
          =<< eIf
            (eSubExp downwards)
            (resultBodyM [distance_downwards_exclusive])
            (resultBodyM [distance_upwards_exclusive])
      distance_exclusive_i64 <- asIntS Int64 distance_exclusive
      distance <-
        letSubExp "distance" $
          I.BasicOp $
            I.BinOp
              (Add Int64 I.OverflowWrap)
              distance_exclusive_i64
              (intConst Int64 1)
      pure (distance, constant False, bounds_invalid)

  step_invalid <-
    letSubExp "step_invalid" $
      I.BasicOp $
        I.BinOp I.LogOr step_wrong_dir step_zero

  invalid <-
    letSubExp "range_invalid" $
      I.BasicOp $
        I.BinOp I.LogOr step_invalid bounds_invalid
  valid <- letSubExp "valid" $ I.BasicOp $ I.UnOp (I.Neg I.Bool) invalid
  cs <- assert "range_valid_c" valid errmsg

  step_i64 <- asIntS Int64 step
  pos_step <-
    letSubExp "pos_step" $
      I.BasicOp $
        I.BinOp (Mul Int64 I.OverflowWrap) step_i64 step_sign_i64

  num_elems <-
    certifying cs $
      letSubExp "num_elems" $
        I.BasicOp $
          I.BinOp (SDivUp Int64 I.Unsafe) distance pos_step

  se <- letSubExp desc (I.BasicOp $ I.Iota num_elems start' step it)
  pure [se]
internaliseAppExp [Char]
desc (E.AppRes StructType
et [VName]
ext) e :: AppExp
e@E.Apply {} =
  case AppExp -> (Function, [(Exp, Maybe VName)])
findFuncall AppExp
e of
    (FunctionHole SrcLoc
loc, [(Exp, Maybe VName)]
_args) -> do
      -- The function we are supposed to call doesn't exist, but we
      -- have to synthesize some fake values of the right type.  The
      -- easy way to do this is to just ignore the arguments and
      -- create a hole whose type is the type of the entire
      -- application.  One caveat is that we need to replace any
      -- existential sizes, too (with zeroes, because they don't
      -- matter).
      let subst :: [(VName, Subst StructRetType)]
subst = (VName -> (VName, Subst StructRetType))
-> [VName] -> [(VName, Subst StructRetType)]
forall a b. (a -> b) -> [a] -> [b]
map (,Exp -> Subst StructRetType
forall t. Exp -> Subst t
E.ExpSubst (Integer -> SrcLoc -> Exp
E.sizeFromInteger Integer
0 SrcLoc
forall a. Monoid a => a
mempty)) [VName]
ext
          et' :: StructType
et' = TypeSubs -> StructType -> StructType
forall a. Substitutable a => TypeSubs -> a -> a
E.applySubst (VName
-> [(VName, Subst StructRetType)] -> Maybe (Subst StructRetType)
forall a b. Eq a => a -> [(a, b)] -> Maybe b
`lookup` [(VName, Subst StructRetType)]
subst) StructType
et
      [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc (Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn. f StructType -> SrcLoc -> ExpBase f vn
E.Hole (StructType -> Info StructType
forall a. a -> Info a
Info StructType
et') SrcLoc
loc)
    (FunctionName QualName VName
qfname, [(Exp, Maybe VName)]
args) -> do
      -- Argument evaluation is outermost-in so that any existential sizes
      -- created by function applications can be brought into scope.
      let fname :: Name
fname = [Char] -> Name
nameFromString ([Char] -> Name) -> [Char] -> Name
forall a b. (a -> b) -> a -> b
$ Name -> [Char]
forall a. Pretty a => a -> [Char]
prettyString (Name -> [Char]) -> Name -> [Char]
forall a b. (a -> b) -> a -> b
$ VName -> Name
baseName (VName -> Name) -> VName -> Name
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname
          arg_desc :: [Char]
arg_desc = Name -> [Char]
nameToString Name
fname [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"_arg"

      -- Some functions are magical (overloaded) and we handle that here.
      case () of
        ()
          -- Short-circuiting operators are magical.
          | VName -> Int
baseTag (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag,
            VName -> [Char]
baseString (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) [Char] -> [Char] -> Bool
forall a. Eq a => a -> a -> Bool
== [Char]
"&&",
            [(Exp
x, Maybe VName
_), (Exp
y, Maybe VName
_)] <- [(Exp, Maybe VName)]
args ->
              [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc (Exp -> InternaliseM [SubExp]) -> Exp -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
                AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
E.AppExp
                  (Exp -> Exp -> Exp -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
E.If Exp
x Exp
y (PrimValue -> SrcLoc -> Exp
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
E.Literal (Bool -> PrimValue
E.BoolValue Bool
False) SrcLoc
forall a. Monoid a => a
mempty) SrcLoc
forall a. Monoid a => a
mempty)
                  (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes (ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
E.Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
-> ScalarTypeBase Exp NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp NoUniqueness
forall dim u. PrimType -> ScalarTypeBase dim u
E.Prim PrimType
E.Bool) [])
          | VName -> Int
baseTag (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag,
            VName -> [Char]
baseString (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) [Char] -> [Char] -> Bool
forall a. Eq a => a -> a -> Bool
== [Char]
"||",
            [(Exp
x, Maybe VName
_), (Exp
y, Maybe VName
_)] <- [(Exp, Maybe VName)]
args ->
              [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc (Exp -> InternaliseM [SubExp]) -> Exp -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
                AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
E.AppExp
                  (Exp -> Exp -> Exp -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
E.If Exp
x (PrimValue -> SrcLoc -> Exp
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
E.Literal (Bool -> PrimValue
E.BoolValue Bool
True) SrcLoc
forall a. Monoid a => a
mempty) Exp
y SrcLoc
forall a. Monoid a => a
mempty)
                  (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes (ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
E.Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
-> ScalarTypeBase Exp NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp NoUniqueness
forall dim u. PrimType -> ScalarTypeBase dim u
E.Prim PrimType
E.Bool) [])
          -- Overloaded and intrinsic functions never take array
          -- arguments (except equality, but those cannot be
          -- existential), so we can safely ignore the existential
          -- dimensions.
          | Just [(StructType, [SubExp])] -> InternaliseM [SubExp]
internalise <- QualName VName
-> [Char]
-> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
isOverloadedFunction QualName VName
qfname [Char]
desc -> do
              let prepareArg :: (Exp, b) -> InternaliseM (StructType, [SubExp])
prepareArg (Exp
arg, b
_) =
                    (StructType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
E.toStruct (Exp -> StructType
E.typeOf Exp
arg),) ([SubExp] -> (StructType, [SubExp]))
-> InternaliseM [SubExp] -> InternaliseM (StructType, [SubExp])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"arg" Exp
arg
              [(StructType, [SubExp])] -> InternaliseM [SubExp]
internalise ([(StructType, [SubExp])] -> InternaliseM [SubExp])
-> InternaliseM [(StructType, [SubExp])] -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ((Exp, Maybe VName) -> InternaliseM (StructType, [SubExp]))
-> [(Exp, Maybe VName)] -> InternaliseM [(StructType, [SubExp])]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Exp, Maybe VName) -> InternaliseM (StructType, [SubExp])
forall {b}. (Exp, b) -> InternaliseM (StructType, [SubExp])
prepareArg [(Exp, Maybe VName)]
args
          | Just [Char] -> InternaliseM [SubExp]
internalise <- QualName VName -> [Exp] -> Maybe ([Char] -> InternaliseM [SubExp])
isIntrinsicFunction QualName VName
qfname (((Exp, Maybe VName) -> Exp) -> [(Exp, Maybe VName)] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (Exp, Maybe VName) -> Exp
forall a b. (a, b) -> a
fst [(Exp, Maybe VName)]
args) ->
              [Char] -> InternaliseM [SubExp]
internalise [Char]
desc
          | VName -> Int
baseTag (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag,
            Just (PrimType
rettype, [PrimType]
_) <- Name
-> Map Name (PrimType, [PrimType]) -> Maybe (PrimType, [PrimType])
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
fname Map Name (PrimType, [PrimType])
I.builtInFunctions -> do
              let tag :: [a] -> [(a, Diet)]
tag [a]
ses = [(a
se, Diet
I.Observe) | a
se <- [a]
ses]
              args' <- [[SubExp]] -> [[SubExp]]
forall a. [a] -> [a]
reverse ([[SubExp]] -> [[SubExp]])
-> InternaliseM [[SubExp]] -> InternaliseM [[SubExp]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Exp, Maybe VName) -> InternaliseM [SubExp])
-> [(Exp, Maybe VName)] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Char] -> (Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg [Char]
arg_desc) ([(Exp, Maybe VName)] -> [(Exp, Maybe VName)]
forall a. [a] -> [a]
reverse [(Exp, Maybe VName)]
args)
              let args'' = ([SubExp] -> [(SubExp, Diet)]) -> [[SubExp]] -> [(SubExp, Diet)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap [SubExp] -> [(SubExp, Diet)]
forall {a}. [a] -> [(a, Diet)]
tag [[SubExp]]
args'
              letValExp' desc $ I.Apply fname args'' [(I.Prim rettype, mempty)] Safe
          | Bool
otherwise -> do
              args' <- [[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[SubExp]] -> [SubExp])
-> ([[SubExp]] -> [[SubExp]]) -> [[SubExp]] -> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[SubExp]] -> [[SubExp]]
forall a. [a] -> [a]
reverse ([[SubExp]] -> [SubExp])
-> InternaliseM [[SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Exp, Maybe VName) -> InternaliseM [SubExp])
-> [(Exp, Maybe VName)] -> InternaliseM [[SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Char] -> (Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg [Char]
arg_desc) ([(Exp, Maybe VName)] -> [(Exp, Maybe VName)]
forall a. [a] -> [a]
reverse [(Exp, Maybe VName)]
args)
              funcall desc qfname args'
internaliseAppExp [Char]
desc AppRes
_ (E.LetPat [SizeBinder VName]
sizes PatBase Info VName StructType
pat Exp
e Exp
body SrcLoc
_) =
  [Char]
-> [SizeBinder VName]
-> PatBase Info VName StructType
-> Exp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
forall a.
[Char]
-> [SizeBinder VName]
-> PatBase Info VName StructType
-> Exp
-> InternaliseM a
-> InternaliseM a
internalisePat [Char]
desc [SizeBinder VName]
sizes PatBase Info VName StructType
pat Exp
e (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
body
internaliseAppExp [Char]
_ AppRes
_ (E.LetFun VName
ofname ([TypeParamBase VName], [PatBase Info VName ParamType],
 Maybe (TypeExp Exp VName), Info ResRetType, Exp)
_ Exp
_ SrcLoc
_) =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"Unexpected LetFun " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ VName -> [Char]
forall a. Pretty a => a -> [Char]
prettyString VName
ofname
internaliseAppExp [Char]
desc AppRes
_ (E.Loop [VName]
sparams PatBase Info VName ParamType
mergepat LoopInitBase Info VName
loopinit LoopFormBase Info VName
form Exp
loopbody SrcLoc
_) = do
  ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"loop_init" (Exp -> InternaliseM [SubExp]) -> Exp -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ LoopInitBase Info VName -> Exp
loopInitExp LoopInitBase Info VName
loopinit
  ((loopbody', (form', shapepat, mergepat', mergeinit')), initstms) <-
    collectStms $ handleForm ses form

  addStms initstms
  mergeinit_ts' <- mapM subExpType mergeinit'

  ctxinit <- argShapes (map I.paramName shapepat) mergepat' mergeinit_ts'

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

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

  -- As above, ensure that the result has the right shape.
  let merge = [Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([Param DeclType]
shapepat [Param DeclType] -> [Param DeclType] -> [Param DeclType]
forall a. [a] -> [a] -> [a]
++ [Param DeclType]
mergepat') [SubExp]
args'
      merge_ts = ((Param DeclType, SubExp) -> TypeBase Shape NoUniqueness)
-> [(Param DeclType, SubExp)] -> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map (Param DeclType -> TypeBase Shape NoUniqueness
forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
I.paramType (Param DeclType -> TypeBase Shape NoUniqueness)
-> ((Param DeclType, SubExp) -> Param DeclType)
-> (Param DeclType, SubExp)
-> TypeBase Shape NoUniqueness
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param DeclType, SubExp) -> Param DeclType
forall a b. (a, b) -> a
fst) [(Param DeclType, SubExp)]
merge
  loopbody'' <-
    localScope (scopeOfFParams (map fst merge) <> scopeOfLoopForm form') . buildBody_ $
      fmap subExpsRes
        . ensureArgShapes
          "shape of loop result does not match shapes in loop parameter"
          (map (I.paramName . fst) merge)
          merge_ts
        . map resSubExp
        =<< bodyBind loopbody'

  attrs <- asks envAttrs
  map I.Var . dropCond
    <$> attributing
      attrs
      (letValExp desc (I.Loop merge form' loopbody''))
  where
    sparams' :: [TypeParamBase VName]
sparams' = (VName -> TypeParamBase VName) -> [VName] -> [TypeParamBase VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SrcLoc -> TypeParamBase VName
forall vn. vn -> SrcLoc -> TypeParamBase vn
`TypeParamDim` SrcLoc
forall a. Monoid a => a
mempty) [VName]
sparams

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

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

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

    handleForm :: [SubExp]
-> LoopFormBase Info VName
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
handleForm [SubExp]
mergeinit (E.ForIn PatBase Info VName StructType
x Exp
arr) = do
      arr' <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"for_in_arr" Exp
arr
      arr_ts <- mapM lookupType arr'
      let w = Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [TypeBase Shape NoUniqueness]
arr_ts

      i <- newVName "i"

      ts <- mapM subExpType mergeinit
      bindingLoopParams sparams' mergepat ts $ \[FParam SOACS]
shapepat [FParam SOACS]
mergepat' ->
        [PatBase Info VName ParamType]
-> [TypeBase Shape NoUniqueness]
-> ([LParam SOACS]
    -> InternaliseM
         (Body SOACS,
          (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a.
[PatBase Info VName ParamType]
-> [TypeBase Shape NoUniqueness]
-> ([LParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLambdaParams [Diet -> StructType -> ParamType
forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
E.Observe (StructType -> ParamType)
-> PatBase Info VName StructType -> PatBase Info VName ParamType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PatBase Info VName StructType
x] ((TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
-> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall u. TypeBase Shape u -> TypeBase Shape u
rowType [TypeBase Shape NoUniqueness]
arr_ts) (([LParam SOACS]
  -> InternaliseM
       (Body SOACS,
        (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
 -> InternaliseM
      (Body SOACS,
       (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> ([LParam SOACS]
    -> InternaliseM
         (Body SOACS,
          (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ \[LParam SOACS]
x_params -> do
          let loopvars :: [(Param (TypeBase Shape NoUniqueness), VName)]
loopvars = [Param (TypeBase Shape NoUniqueness)]
-> [VName] -> [(Param (TypeBase Shape NoUniqueness), VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param (TypeBase Shape NoUniqueness)]
[LParam SOACS]
x_params [VName]
arr'
          [Param DeclType]
-> [Param DeclType]
-> [SubExp]
-> VName
-> [(Param (TypeBase Shape NoUniqueness), VName)]
-> LoopForm
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forLoop [Param DeclType]
[FParam SOACS]
mergepat' [Param DeclType]
[FParam SOACS]
shapepat [SubExp]
mergeinit VName
i [(Param (TypeBase Shape NoUniqueness), VName)]
loopvars (LoopForm
 -> InternaliseM
      (Body SOACS,
       (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> LoopForm
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ VName -> IntType -> SubExp -> LoopForm
I.ForLoop VName
i IntType
Int64 SubExp
w
    handleForm [SubExp]
mergeinit (E.For IdentBase Info VName StructType
i Exp
num_iterations) = do
      num_iterations' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"upper_bound" Exp
num_iterations
      num_iterations_t <- I.subExpType num_iterations'
      it <- case num_iterations_t of
        I.Prim (IntType IntType
it) -> IntType -> InternaliseM IntType
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure IntType
it
        TypeBase Shape NoUniqueness
_ -> [Char] -> InternaliseM IntType
forall a. HasCallStack => [Char] -> a
error [Char]
"internaliseExp Loop: invalid type"

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

        (loop_initial_cond, init_loop_cond_stms) <- collectStms $ do
          forM_ (zip shapepat shapeinit) $ \(Param DeclType
p, SubExp
se) ->
            [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName Param DeclType
p] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
          forM_ (zip mergepat' mergeinit) $ \(Param DeclType
p, SubExp
se) ->
            Bool -> InternaliseM () -> InternaliseM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (SubExp
se SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== VName -> SubExp
I.Var (Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName Param DeclType
p)) (InternaliseM () -> InternaliseM ())
-> InternaliseM () -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
              [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName Param DeclType
p] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
                case SubExp
se of
                  I.Var VName
v
                    | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoUniqueness -> Bool
forall shape u. TypeBase shape u -> Bool
primType (TypeBase Shape NoUniqueness -> Bool)
-> TypeBase Shape NoUniqueness -> Bool
forall a b. (a -> b) -> a -> b
$ Param DeclType -> TypeBase Shape NoUniqueness
forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
paramType Param DeclType
p ->
                        [SubExp] -> VName -> Exp SOACS
forall rep. [SubExp] -> VName -> Exp rep
shapeCoerce (TypeBase Shape NoUniqueness -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
I.arrayDims (TypeBase Shape NoUniqueness -> [SubExp])
-> TypeBase Shape NoUniqueness -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Param DeclType -> TypeBase Shape NoUniqueness
forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
paramType Param DeclType
p) VName
v
                  SubExp
_ -> BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se

          -- As the condition expression is inserted twice, we have to
          -- avoid shadowing (#1935).
          (cond_stms, cond') <-
            uncurry (flip renameStmsWith)
              =<< collectStms (internaliseExp1 "loop_cond" cond)
          addStms cond_stms
          pure cond'

        addStms init_loop_cond_stms

        bodyFromStms $ do
          ses <- loopBody
          sets <- mapM subExpType ses
          loop_while <- newParam "loop_while" $ I.Prim I.Bool
          shapeargs <- argShapes (map I.paramName shapepat) mergepat' sets

          -- Careful not to clobber anything.
          loop_end_cond_body <- renameBody <=< buildBody_ $ do
            forM_ (zip shapepat shapeargs) $ \(Param DeclType
p, SubExp
se) ->
              Bool -> InternaliseM () -> InternaliseM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (SubExp
se SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== VName -> SubExp
I.Var (Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName Param DeclType
p)) (InternaliseM () -> InternaliseM ())
-> InternaliseM () -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
                [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName Param DeclType
p] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
                  BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
                    SubExp -> BasicOp
SubExp SubExp
se
            forM_ (zip mergepat' ses) $ \(Param DeclType
p, SubExp
se) ->
              Bool -> InternaliseM () -> InternaliseM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (SubExp
se SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== VName -> SubExp
I.Var (Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName Param DeclType
p)) (InternaliseM () -> InternaliseM ())
-> InternaliseM () -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
                [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName Param DeclType
p] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
                  case SubExp
se of
                    I.Var VName
v
                      | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoUniqueness -> Bool
forall shape u. TypeBase shape u -> Bool
primType (TypeBase Shape NoUniqueness -> Bool)
-> TypeBase Shape NoUniqueness -> Bool
forall a b. (a -> b) -> a -> b
$ Param DeclType -> TypeBase Shape NoUniqueness
forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
paramType Param DeclType
p ->
                          [SubExp] -> VName -> Exp SOACS
forall rep. [SubExp] -> VName -> Exp rep
shapeCoerce (TypeBase Shape NoUniqueness -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
I.arrayDims (TypeBase Shape NoUniqueness -> [SubExp])
-> TypeBase Shape NoUniqueness -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Param DeclType -> TypeBase Shape NoUniqueness
forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
paramType Param DeclType
p) VName
v
                    SubExp
_ -> BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
            subExpsRes <$> internaliseExp "loop_cond" cond
          loop_end_cond <- bodyBind loop_end_cond_body

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

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

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

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

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

  let comb VName
sname SubExp
ve' = do
        sname_t <- VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
sname
        let full_slice = TypeBase Shape NoUniqueness -> [DimIndex SubExp] -> Slice SubExp
fullSlice TypeBase Shape NoUniqueness
sname_t [DimIndex SubExp]
idxs'
            rowtype = TypeBase Shape NoUniqueness
sname_t TypeBase Shape NoUniqueness
-> [SubExp] -> TypeBase Shape NoUniqueness
forall oldshape u.
TypeBase oldshape u -> [SubExp] -> TypeBase Shape u
`setArrayDims` Slice SubExp -> [SubExp]
forall d. Slice d -> [d]
sliceDims Slice SubExp
full_slice
        ve'' <-
          ensureShape errormsg rowtype "lw_val_correct_shape" ve'
        letInPlace desc sname full_slice $ BasicOp $ SubExp ve''
  certifying cs $ map I.Var <$> zipWithM comb srcs ves
internaliseExp [Char]
desc (E.RecordUpdate Exp
src [Name]
fields Exp
ve Info StructType
_ SrcLoc
loc) = SrcLoc -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. Located a => a -> InternaliseM b -> InternaliseM b
locating SrcLoc
loc (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ do
  src' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
src
  ve' <- internaliseExp desc ve
  replace (E.typeOf src) fields ve' src'
  where
    replace :: TypeBase Exp als -> [Name] -> [a] -> [a] -> m [a]
replace (E.Scalar (E.Record Map Name (TypeBase Exp als)
m)) (Name
f : [Name]
fs) [a]
ve' [a]
src'
      | Just TypeBase Exp als
t <- Name -> Map Name (TypeBase Exp als) -> Maybe (TypeBase Exp als)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f Map Name (TypeBase Exp als)
m = do
          let i :: Int
i =
                [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int)
-> ([(Name, TypeBase Exp als)] -> [Int])
-> [(Name, TypeBase Exp als)]
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Name, TypeBase Exp als) -> Int)
-> [(Name, TypeBase Exp als)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase Exp als -> Int
forall als. TypeBase Exp als -> Int
internalisedTypeSize (TypeBase Exp als -> Int)
-> ((Name, TypeBase Exp als) -> TypeBase Exp als)
-> (Name, TypeBase Exp als)
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, TypeBase Exp als) -> TypeBase Exp als
forall a b. (a, b) -> b
snd) ([(Name, TypeBase Exp als)] -> Int)
-> [(Name, TypeBase Exp als)] -> Int
forall a b. (a -> b) -> a -> b
$
                  ((Name, TypeBase Exp als) -> Bool)
-> [(Name, TypeBase Exp als)] -> [(Name, TypeBase Exp als)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
/= Name
f) (Name -> Bool)
-> ((Name, TypeBase Exp als) -> Name)
-> (Name, TypeBase Exp als)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, TypeBase Exp als) -> Name
forall a b. (a, b) -> a
fst) ([(Name, TypeBase Exp als)] -> [(Name, TypeBase Exp als)])
-> (Map Name (TypeBase Exp als) -> [(Name, TypeBase Exp als)])
-> Map Name (TypeBase Exp als)
-> [(Name, TypeBase Exp als)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name (TypeBase Exp als) -> [(Name, TypeBase Exp als)]
forall a. Map Name a -> [(Name, a)]
sortFields (Map Name (TypeBase Exp als) -> [(Name, TypeBase Exp als)])
-> Map Name (TypeBase Exp als) -> [(Name, TypeBase Exp als)]
forall a b. (a -> b) -> a -> b
$
                    Map Name (TypeBase Exp als)
m
              k :: Int
k = TypeBase Exp als -> Int
forall als. TypeBase Exp als -> Int
internalisedTypeSize TypeBase Exp als
t
              ([a]
bef, [a]
to_update, [a]
aft) = Int -> Int -> [a] -> ([a], [a], [a])
forall a. Int -> Int -> [a] -> ([a], [a], [a])
splitAt3 Int
i Int
k [a]
src'
          src'' <- TypeBase Exp als -> [Name] -> [a] -> [a] -> m [a]
replace TypeBase Exp als
t [Name]
fs [a]
ve' [a]
to_update
          pure $ bef ++ src'' ++ aft
    replace TypeBase Exp als
_ [Name]
_ [a]
ve' [a]
_ = [a] -> m [a]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [a]
ve'
internaliseExp [Char]
desc (E.Attr AttrInfo VName
attr Exp
e SrcLoc
loc) = do
  attr' <- AttrInfo VName -> InternaliseM Attr
internaliseAttr AttrInfo VName
attr
  e' <- local (f attr') $ internaliseExp desc e
  case attr' of
    Attr
"trace" ->
      Text -> [SubExp] -> InternaliseM [SubExp]
traceRes ([Char] -> Text
T.pack ([Char] -> Text) -> [Char] -> Text
forall a b. (a -> b) -> a -> b
$ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr SrcLoc
loc) [SubExp]
e'
    I.AttrComp Name
"trace" [I.AttrName Name
tag] ->
      Text -> [SubExp] -> InternaliseM [SubExp]
traceRes (Name -> Text
nameToText Name
tag) [SubExp]
e'
    Attr
"opaque" ->
      (SubExp -> InternaliseM SubExp)
-> [SubExp] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc (Exp SOACS -> InternaliseM SubExp)
-> (SubExp -> Exp SOACS) -> SubExp -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS)
-> (SubExp -> BasicOp) -> SubExp -> Exp SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpaqueOp -> SubExp -> BasicOp
Opaque OpaqueOp
OpaqueNil) [SubExp]
e'
    Attr
"scratch" -> do
      ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
e'
      forM (zip ts e') $ \(TypeBase Shape NoUniqueness
t, SubExp
se) ->
        case TypeBase Shape NoUniqueness
t of
          I.Array PrimType
pt Shape
shape NoUniqueness
_ ->
            [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ PrimType -> [SubExp] -> BasicOp
I.Scratch PrimType
pt ([SubExp] -> BasicOp) -> [SubExp] -> BasicOp
forall a b. (a -> b) -> a -> b
$ Shape -> [SubExp]
forall d. ShapeBase d -> [d]
I.shapeDims Shape
shape
          I.Prim PrimType
pt ->
            SubExp -> InternaliseM SubExp
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> InternaliseM SubExp) -> SubExp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
forall v. IsValue v => v -> SubExp
constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimValue
blankPrimValue PrimType
pt
          TypeBase Shape NoUniqueness
_ -> SubExp -> InternaliseM SubExp
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
se
    Attr
"blank" -> do
      ts <- (SubExp -> InternaliseM (TypeBase Shape NoUniqueness))
-> [SubExp] -> InternaliseM [TypeBase Shape NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> InternaliseM (TypeBase Shape NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
e'
      forM (zip ts e') $ \(TypeBase Shape NoUniqueness
t, SubExp
se) ->
        case TypeBase Shape NoUniqueness
t of
          I.Array PrimType
pt Shape
shape NoUniqueness
_ ->
            [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc (Exp SOACS -> InternaliseM SubExp)
-> (PrimValue -> Exp SOACS) -> PrimValue -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS)
-> (PrimValue -> BasicOp) -> PrimValue -> Exp SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shape -> SubExp -> BasicOp
I.Replicate Shape
shape (SubExp -> BasicOp)
-> (PrimValue -> SubExp) -> PrimValue -> BasicOp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimValue -> SubExp
forall v. IsValue v => v -> SubExp
constant (PrimValue -> InternaliseM SubExp)
-> PrimValue -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
              PrimType -> PrimValue
blankPrimValue PrimType
pt
          I.Prim PrimType
pt ->
            SubExp -> InternaliseM SubExp
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> InternaliseM SubExp) -> SubExp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
forall v. IsValue v => v -> SubExp
constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimValue
blankPrimValue PrimType
pt
          TypeBase Shape NoUniqueness
_ -> SubExp -> InternaliseM SubExp
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
se
    Attr
_ ->
      [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
e'
  where
    traceRes :: Text -> [SubExp] -> InternaliseM [SubExp]
traceRes Text
tag' =
      (SubExp -> InternaliseM SubExp)
-> [SubExp] -> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc (Exp SOACS -> InternaliseM SubExp)
-> (SubExp -> Exp SOACS) -> SubExp -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS)
-> (SubExp -> BasicOp) -> SubExp -> Exp SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpaqueOp -> SubExp -> BasicOp
Opaque (Text -> OpaqueOp
OpaqueTrace Text
tag'))
    f :: Attr -> InternaliseEnv -> InternaliseEnv
f Attr
attr' InternaliseEnv
env
      | Attr
attr' Attr -> Attr -> Bool
forall a. Eq a => a -> a -> Bool
== Attr
"unsafe",
        Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ InternaliseEnv -> Bool
envSafe InternaliseEnv
env =
          InternaliseEnv
env {envDoBoundsChecks = False}
      | Bool
otherwise =
          InternaliseEnv
env {envAttrs = envAttrs env <> oneAttr attr'}
internaliseExp [Char]
desc (E.Assert Exp
e1 Exp
e2 (Info Text
check) SrcLoc
loc) = SrcLoc -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. Located a => a -> InternaliseM b -> InternaliseM b
locating SrcLoc
loc (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ do
  e1' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"assert_cond" Exp
e1
  c <- assert "assert_c" e1' $ errorMsg [ErrorString $ "Assertion is false: " <> check]
  -- Make sure there are some bindings to certify.
  certifying c $ mapM rebind =<< internaliseExp desc e2
  where
    rebind :: SubExp -> m SubExp
rebind SubExp
v = do
      v' <- [Char] -> m VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"assert_res"
      letBindNames [v'] $ I.BasicOp $ I.SubExp v
      pure $ I.Var v'
internaliseExp [Char]
_ (E.Constr Name
c [Exp]
es (Info (E.Scalar (E.Sum Map Name [StructType]
fs))) SrcLoc
loc) = SrcLoc -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. Located a => a -> InternaliseM b -> InternaliseM b
locating SrcLoc
loc (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ do
  (ts, constr_map) <- Map Name [StructType]
-> InternaliseM ([TypeBase ExtShape Uniqueness], [(Name, [Int])])
internaliseSumType (Map Name [StructType]
 -> InternaliseM ([TypeBase ExtShape Uniqueness], [(Name, [Int])]))
-> Map Name [StructType]
-> InternaliseM ([TypeBase ExtShape Uniqueness], [(Name, [Int])])
forall a b. (a -> b) -> a -> b
$ ([StructType] -> [StructType])
-> Map Name [StructType] -> Map Name [StructType]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((StructType -> StructType) -> [StructType] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map StructType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
E.toStruct) Map Name [StructType]
fs
  es' <- concat <$> mapM (internaliseExp "payload") es

  let noExt p
_ = SubExp -> f SubExp
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> f SubExp) -> SubExp -> f SubExp
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0
  ts' <- instantiateShapes noExt $ map fromDecl ts

  case lookupWithIndex c constr_map of
    Just (Int
i, [Int]
js) ->
      (IntType -> Integer -> SubExp
intConst IntType
Int8 (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
i) SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
:) ([SubExp] -> [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int
-> [TypeBase Shape NoUniqueness]
-> [(Int, SubExp)]
-> InternaliseM [SubExp]
forall {f :: * -> *} {t}.
(Num t, MonadBuilder f, Eq t) =>
t -> [TypeBase Shape NoUniqueness] -> [(t, SubExp)] -> f [SubExp]
clauses Int
0 [TypeBase Shape NoUniqueness]
ts' ([Int] -> [SubExp] -> [(Int, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
js [SubExp]
es')
    Maybe (Int, [Int])
Nothing ->
      [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error [Char]
"internaliseExp Constr: missing constructor"
  where
    clauses :: t -> [TypeBase Shape NoUniqueness] -> [(t, SubExp)] -> f [SubExp]
clauses t
j (TypeBase Shape NoUniqueness
t : [TypeBase Shape NoUniqueness]
ts) [(t, SubExp)]
js_to_es
      | Just SubExp
e <- t
j t -> [(t, SubExp)] -> Maybe SubExp
forall a b. Eq a => a -> [(a, b)] -> Maybe b
`lookup` [(t, SubExp)]
js_to_es =
          (SubExp
e SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
:) ([SubExp] -> [SubExp]) -> f [SubExp] -> f [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> t -> [TypeBase Shape NoUniqueness] -> [(t, SubExp)] -> f [SubExp]
clauses (t
j t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) [TypeBase Shape NoUniqueness]
ts [(t, SubExp)]
js_to_es
      | Bool
otherwise = do
          blank <-
            -- Cannot use eBlank here for arrays, because when doing
            -- equality comparisons on sum types, we end up looking at
            -- the array elements. (#2081) This is a bit of an edge
            -- case, but arrays in sum types are known to be
            -- inefficient.
            [Char] -> Exp (Rep f) -> f SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"zero"
              (Exp (Rep f) -> f SubExp) -> f (Exp (Rep f)) -> f SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< case TypeBase Shape NoUniqueness
t of
                I.Array {} ->
                  Exp (Rep f) -> f (Exp (Rep f))
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp (Rep f) -> f (Exp (Rep f))) -> Exp (Rep f) -> f (Exp (Rep f))
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep f)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep f)) -> BasicOp -> Exp (Rep f)
forall a b. (a -> b) -> a -> b
$ Shape -> SubExp -> BasicOp
Replicate (TypeBase Shape NoUniqueness -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape TypeBase Shape NoUniqueness
t) (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimValue
blankPrimValue (PrimType -> PrimValue) -> PrimType -> PrimValue
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoUniqueness -> PrimType
forall shape u. TypeBase shape u -> PrimType
elemType TypeBase Shape NoUniqueness
t
                TypeBase Shape NoUniqueness
_ -> TypeBase Shape NoUniqueness -> f (Exp (Rep f))
forall (m :: * -> *).
MonadBuilder m =>
TypeBase Shape NoUniqueness -> m (Exp (Rep m))
eBlank TypeBase Shape NoUniqueness
t
          (blank :) <$> clauses (j + 1) ts js_to_es
    clauses t
_ [] [(t, SubExp)]
_ =
      [SubExp] -> f [SubExp]
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
internaliseExp [Char]
_ (E.Constr Name
_ [Exp]
_ (Info StructType
t) SrcLoc
loc) =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: constructor with type " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ StructType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString StructType
t [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr SrcLoc
loc
-- The "interesting" cases are over, now it's mostly boilerplate.

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

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

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

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

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

internalisePat ::
  String ->
  [E.SizeBinder VName] ->
  E.Pat StructType ->
  E.Exp ->
  InternaliseM a ->
  InternaliseM a
internalisePat :: forall a.
[Char]
-> [SizeBinder VName]
-> PatBase Info VName StructType
-> Exp
-> InternaliseM a
-> InternaliseM a
internalisePat [Char]
desc [SizeBinder VName]
sizes PatBase Info VName StructType
p Exp
e InternaliseM a
m = do
  ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc' Exp
e
  internalisePat' sizes p ses m
  where
    desc' :: [Char]
desc' = case PatBase Info VName StructType -> [IdentBase Info VName StructType]
forall (f :: * -> *) vn t. PatBase f vn t -> [IdentBase f vn t]
E.patIdents PatBase Info VName StructType
p of
      [IdentBase Info VName StructType
v] -> VName -> [Char]
baseString (VName -> [Char]) -> VName -> [Char]
forall a b. (a -> b) -> a -> b
$ IdentBase Info VName StructType -> VName
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
E.identName IdentBase Info VName StructType
v
      [IdentBase Info VName StructType]
_ -> [Char]
desc

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

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

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

-- Special-case an important common case that otherwise leads to horrible code.
internaliseDimIndex
  SubExp
w
  ( E.DimSlice
      Maybe Exp
Nothing
      Maybe Exp
Nothing
      (Just (E.Negate (E.IntLit Integer
1 Info StructType
_ SrcLoc
_) SrcLoc
_))
    ) = do
    w_minus_1 <-
      [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"w_minus_1" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
        BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
          BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) SubExp
w SubExp
one
    pure
      ( I.DimSlice w_minus_1 w $ intConst Int64 (-1),
        constant True,
        mempty
      )
    where
      one :: SubExp
one = Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
1 :: Int64)
internaliseDimIndex SubExp
w (E.DimSlice Maybe Exp
i Maybe Exp
j Maybe Exp
s) = do
  s' <- InternaliseM SubExp
-> (Exp -> InternaliseM SubExp) -> Maybe Exp -> InternaliseM SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (SubExp -> InternaliseM SubExp
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
one) (((SubExp, IntType) -> SubExp)
-> InternaliseM (SubExp, IntType) -> InternaliseM SubExp
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SubExp, IntType) -> SubExp
forall a b. (a, b) -> a
fst (InternaliseM (SubExp, IntType) -> InternaliseM SubExp)
-> (Exp -> InternaliseM (SubExp, IntType))
-> Exp
-> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp -> InternaliseM (SubExp, IntType)
internaliseSizeExp [Char]
"s") Maybe Exp
s
  s_sign <- letSubExp "s_sign" $ BasicOp $ I.UnOp (I.SSignum Int64) s'
  backwards <- letSubExp "backwards" $ I.BasicOp $ I.CmpOp (I.CmpEq int64) s_sign negone
  w_minus_1 <- letSubExp "w_minus_1" $ BasicOp $ I.BinOp (Sub Int64 I.OverflowWrap) w one
  let i_def =
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"i_def"
          (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
            (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
backwards)
            ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
w_minus_1])
            ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
zero])
      j_def =
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"j_def"
          (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
-> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
            (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
backwards)
            ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
negone])
            ([SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
w])
  i' <- maybe i_def (fmap fst . internaliseSizeExp "i") i
  j' <- maybe j_def (fmap fst . internaliseSizeExp "j") j
  j_m_i <- letSubExp "j_m_i" $ BasicOp $ I.BinOp (Sub Int64 I.OverflowWrap) j' i'
  -- Something like a division-rounding-up, but accomodating negative
  -- operands.
  let divRounding m (Exp (Rep m))
x m (Exp (Rep m))
y =
        BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eBinOp
          (IntType -> Safety -> BinOp
SQuot IntType
Int64 Safety
Safe)
          ( BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eBinOp
              (IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
I.OverflowWrap)
              m (Exp (Rep m))
x
              (BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eBinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) m (Exp (Rep m))
y (m (Exp (Rep m)) -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
m (Exp (Rep m)) -> m (Exp (Rep m))
eSignum m (Exp (Rep m))
y))
          )
          m (Exp (Rep m))
y
  n <- letSubExp "n" =<< divRounding (toExp j_m_i) (toExp s')

  zero_stride <- letSubExp "zero_stride" $ I.BasicOp $ I.CmpOp (CmpEq int64) s_sign zero
  nonzero_stride <- letSubExp "nonzero_stride" $ I.BasicOp $ I.UnOp (I.Neg I.Bool) zero_stride

  -- Bounds checks depend on whether we are slicing forwards or
  -- backwards.  If forwards, we must check '0 <= i && i <= j'.  If
  -- backwards, '-1 <= j && j <= i'.  In both cases, we check '0 <=
  -- i+n*s && i+(n-1)*s < w'.  We only check if the slice is nonempty.
  empty_slice <- letSubExp "empty_slice" $ I.BasicOp $ I.CmpOp (CmpEq int64) n zero

  m <- letSubExp "m" $ I.BasicOp $ I.BinOp (Sub Int64 I.OverflowWrap) n one
  m_t_s <- letSubExp "m_t_s" $ I.BasicOp $ I.BinOp (Mul Int64 I.OverflowWrap) m s'
  i_p_m_t_s <- letSubExp "i_p_m_t_s" $ I.BasicOp $ I.BinOp (Add Int64 I.OverflowWrap) i' m_t_s
  zero_leq_i_p_m_t_s <-
    letSubExp "zero_leq_i_p_m_t_s" $
      I.BasicOp $
        I.CmpOp (I.CmpSle Int64) zero i_p_m_t_s
  i_p_m_t_s_leq_w <-
    letSubExp "i_p_m_t_s_leq_w" $
      I.BasicOp $
        I.CmpOp (I.CmpSle Int64) i_p_m_t_s w
  i_p_m_t_s_lth_w <-
    letSubExp "i_p_m_t_s_leq_w" $
      I.BasicOp $
        I.CmpOp (I.CmpSlt Int64) i_p_m_t_s w

  zero_lte_i <- letSubExp "zero_lte_i" $ I.BasicOp $ I.CmpOp (I.CmpSle Int64) zero i'
  i_lte_j <- letSubExp "i_lte_j" $ I.BasicOp $ I.CmpOp (I.CmpSle Int64) i' j'
  forwards_ok <-
    letSubExp "forwards_ok"
      =<< eAll [zero_lte_i, i_lte_j, zero_leq_i_p_m_t_s, i_p_m_t_s_lth_w]

  negone_lte_j <- letSubExp "negone_lte_j" $ I.BasicOp $ I.CmpOp (I.CmpSle Int64) negone j'
  j_lte_i <- letSubExp "j_lte_i" $ I.BasicOp $ I.CmpOp (I.CmpSle Int64) j' i'
  backwards_ok <-
    letSubExp "backwards_ok"
      =<< eAll
        [negone_lte_j, j_lte_i, zero_leq_i_p_m_t_s, i_p_m_t_s_leq_w]

  slice_ok <-
    letSubExp "slice_ok"
      =<< eIf
        (eSubExp backwards)
        (resultBodyM [backwards_ok])
        (resultBodyM [forwards_ok])

  ok_or_empty <-
    letSubExp "ok_or_empty" $
      I.BasicOp $
        I.BinOp I.LogOr empty_slice slice_ok

  acceptable <-
    letSubExp "slice_acceptable" $
      I.BasicOp $
        I.BinOp I.LogAnd nonzero_stride ok_or_empty

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

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

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

  -- reshape neutral element to have same size as the destination array
  ne_shp <- forM (zip ne' hist') $ \(SubExp
n, VName
h) -> do
    rowtype <- Int -> TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall u. Int -> TypeBase Shape u -> TypeBase Shape u
I.stripArray Int
1 (TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
-> InternaliseM (TypeBase Shape NoUniqueness)
-> InternaliseM (TypeBase Shape NoUniqueness)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
h
    ensureShape
      "Row shape of destination array does not match shape of neutral element"
      rowtype
      "hist_ne_right_shape"
      n
  ne_ts <- mapM I.subExpType ne_shp
  his_ts <- mapM (fmap (I.stripArray (dim - 1)) . lookupType) hist'
  op' <- internaliseFoldLambda internaliseLambda op ne_ts his_ts

  -- reshape return type of bucket function to have same size as neutral element
  -- (modulo the index)
  bucket_params <- replicateM dim (newParam "bucket_p" $ I.Prim int64)
  img_params <- mapM (newParam "img_p" . rowType) =<< mapM lookupType img'
  let params = [Param (TypeBase Shape NoUniqueness)]
bucket_params [Param (TypeBase Shape NoUniqueness)]
-> [Param (TypeBase Shape NoUniqueness)]
-> [Param (TypeBase Shape NoUniqueness)]
forall a. [a] -> [a] -> [a]
++ [Param (TypeBase Shape NoUniqueness)]
img_params
      rettype = Int -> TypeBase Shape NoUniqueness -> [TypeBase Shape NoUniqueness]
forall a. Int -> a -> [a]
replicate Int
dim (PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64) [TypeBase Shape NoUniqueness]
-> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
forall a. [a] -> [a] -> [a]
++ [TypeBase Shape NoUniqueness]
ne_ts
      body = Stms SOACS -> Result -> Body SOACS
forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody Stms SOACS
forall a. Monoid a => a
mempty (Result -> Body SOACS) -> Result -> Body SOACS
forall a b. (a -> b) -> a -> b
$ [VName] -> Result
varsRes ([VName] -> Result) -> [VName] -> Result
forall a b. (a -> b) -> a -> b
$ (Param (TypeBase Shape NoUniqueness) -> VName)
-> [Param (TypeBase Shape NoUniqueness)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (TypeBase Shape NoUniqueness) -> VName
forall dec. Param dec -> VName
I.paramName [Param (TypeBase Shape NoUniqueness)]
params
  lam' <-
    mkLambda params $
      ensureResultShape
        "Row shape of value array does not match row shape of hist target"
        rettype
        =<< bodyBind body

  -- get sizes of histogram and image arrays
  shape_hist <- I.Shape . take dim . I.arrayDims <$> lookupType (head hist')
  w_img <- I.arraySize 0 <$> lookupType (head img')

  letValExp' desc . I.Op $
    I.Hist w_img (buckets' ++ img') [HistOp shape_hist rf' hist' ne_shp op'] lam'

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

  acc_cert_v <- newVName "acc_cert"
  dest_ts <- mapM lookupType dest'
  let dest_w = Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [TypeBase Shape NoUniqueness]
dest_ts
      acc_t = VName
-> Shape
-> [TypeBase Shape NoUniqueness]
-> NoUniqueness
-> TypeBase Shape NoUniqueness
forall shape u.
VName
-> Shape -> [TypeBase Shape NoUniqueness] -> u -> TypeBase shape u
Acc VName
acc_cert_v ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
I.Shape [SubExp
dest_w]) ((TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness)
-> [TypeBase Shape NoUniqueness] -> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Shape NoUniqueness -> TypeBase Shape NoUniqueness
forall u. TypeBase Shape u -> TypeBase Shape u
rowType [TypeBase Shape NoUniqueness]
dest_ts) NoUniqueness
NoUniqueness
  acc_p <- newParam "acc_p" acc_t
  withacc_lam <- mkLambda [Param mempty acc_cert_v (I.Prim I.Unit), acc_p] $ do
    bs_ts <- mapM lookupType bs'
    lam' <- internaliseLambdaCoerce lam $ map rowType $ paramType acc_p : bs_ts
    let w = Int -> [TypeBase Shape NoUniqueness] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [TypeBase Shape NoUniqueness]
bs_ts
    fmap subExpsRes . letValExp' "acc_res" $
      I.Op $
        I.Screma w (I.paramName acc_p : bs') (I.mapSOAC lam')

  op' <-
    case op of
      Just (Exp
op_lam, Exp
ne) -> do
        ne' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"hist_ne" Exp
ne
        ne_ts <- mapM I.subExpType ne'
        (lam_params, lam_body, lam_rettype) <-
          internaliseLambda op_lam $ ne_ts ++ ne_ts
        idxp <- newParam "idx" $ I.Prim int64
        let op_lam' = [LParam SOACS]
-> [TypeBase Shape NoUniqueness] -> Body SOACS -> Lambda SOACS
forall rep.
[LParam rep]
-> [TypeBase Shape NoUniqueness] -> Body rep -> Lambda rep
I.Lambda (Param (TypeBase Shape NoUniqueness)
idxp Param (TypeBase Shape NoUniqueness)
-> [Param (TypeBase Shape NoUniqueness)]
-> [Param (TypeBase Shape NoUniqueness)]
forall a. a -> [a] -> [a]
: [Param (TypeBase Shape NoUniqueness)]
lam_params) [TypeBase Shape NoUniqueness]
lam_rettype Body SOACS
lam_body
        pure $ Just (op_lam', ne')
      Maybe (Exp, Exp)
Nothing ->
        Maybe (Lambda SOACS, [SubExp])
-> InternaliseM (Maybe (Lambda SOACS, [SubExp]))
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Lambda SOACS, [SubExp])
forall a. Maybe a
Nothing

  destw <- arraysSize 0 <$> mapM lookupType dest'
  fmap (map I.Var) $
    letTupExp desc $
      WithAcc [(I.Shape [destw], dest', op')] withacc_lam

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

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

internaliseExpToVars :: String -> E.Exp -> InternaliseM [I.VName]
internaliseExpToVars :: [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
desc Exp
e =
  (SubExp -> InternaliseM VName) -> [SubExp] -> InternaliseM [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> InternaliseM VName
asIdent ([SubExp] -> InternaliseM [VName])
-> InternaliseM [SubExp] -> InternaliseM [VName]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
  where
    asIdent :: SubExp -> InternaliseM VName
asIdent (I.Var VName
v) = VName -> InternaliseM VName
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v
    asIdent SubExp
se = [Char] -> Exp (Rep InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM VName)
-> Exp (Rep InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
se

internaliseOperation ::
  String ->
  E.Exp ->
  (I.VName -> InternaliseM I.BasicOp) ->
  InternaliseM [I.SubExp]
internaliseOperation :: [Char]
-> Exp -> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
internaliseOperation [Char]
s Exp
e VName -> InternaliseM BasicOp
op = do
  vs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
s Exp
e
  mapM (letSubExp s . I.BasicOp <=< op) vs

certifyingNonzero ::
  IntType ->
  SubExp ->
  InternaliseM a ->
  InternaliseM a
certifyingNonzero :: forall a. IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero IntType
t SubExp
x InternaliseM a
m = do
  zero <-
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"zero" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
        CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp (PrimType -> CmpOp
CmpEq (IntType -> PrimType
IntType IntType
t)) SubExp
x (IntType -> Integer -> SubExp
intConst IntType
t Integer
0)
  nonzero <- letSubExp "nonzero" $ I.BasicOp $ UnOp (I.Neg I.Bool) zero
  c <- assert "nonzero_cert" nonzero "division by zero"
  certifying c m

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

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

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

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

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

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

findFuncall :: E.AppExp -> (Function, [(E.Exp, Maybe VName)])
findFuncall :: AppExp -> (Function, [(Exp, Maybe VName)])
findFuncall (E.Apply Exp
f NonEmpty (Info (Maybe VName), Exp)
args SrcLoc
_)
  | E.Var QualName VName
fname Info StructType
_ SrcLoc
_ <- Exp
f =
      (QualName VName -> Function
FunctionName QualName VName
fname, ((Info (Maybe VName), Exp) -> (Exp, Maybe VName))
-> [(Info (Maybe VName), Exp)] -> [(Exp, Maybe VName)]
forall a b. (a -> b) -> [a] -> [b]
map (Info (Maybe VName), Exp) -> (Exp, Maybe VName)
forall {b} {a}. (Info b, a) -> (a, b)
onArg ([(Info (Maybe VName), Exp)] -> [(Exp, Maybe VName)])
-> [(Info (Maybe VName), Exp)] -> [(Exp, Maybe VName)]
forall a b. (a -> b) -> a -> b
$ NonEmpty (Info (Maybe VName), Exp) -> [(Info (Maybe VName), Exp)]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Info (Maybe VName), Exp)
args)
  | E.Hole (Info StructType
_) SrcLoc
loc <- Exp
f =
      (SrcLoc -> Function
FunctionHole SrcLoc
loc, ((Info (Maybe VName), Exp) -> (Exp, Maybe VName))
-> [(Info (Maybe VName), Exp)] -> [(Exp, Maybe VName)]
forall a b. (a -> b) -> [a] -> [b]
map (Info (Maybe VName), Exp) -> (Exp, Maybe VName)
forall {b} {a}. (Info b, a) -> (a, b)
onArg ([(Info (Maybe VName), Exp)] -> [(Exp, Maybe VName)])
-> [(Info (Maybe VName), Exp)] -> [(Exp, Maybe VName)]
forall a b. (a -> b) -> a -> b
$ NonEmpty (Info (Maybe VName), Exp) -> [(Info (Maybe VName), Exp)]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Info (Maybe VName), Exp)
args)
  where
    onArg :: (Info b, a) -> (a, b)
onArg (Info b
argext, a
e) = (a
e, b
argext)
findFuncall AppExp
e =
  [Char] -> (Function, [(Exp, Maybe VName)])
forall a. HasCallStack => [Char] -> a
error ([Char] -> (Function, [(Exp, Maybe VName)]))
-> [Char] -> (Function, [(Exp, Maybe VName)])
forall a b. (a -> b) -> a -> b
$ [Char]
"Invalid function expression in application:\n" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ AppExp -> [Char]
forall a. Pretty a => a -> [Char]
prettyString AppExp
e

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

internaliseLambda :: InternaliseLambda
internaliseLambda :: InternaliseLambda
internaliseLambda (E.Parens Exp
e SrcLoc
_) [TypeBase Shape NoUniqueness]
rowtypes =
  InternaliseLambda
internaliseLambda Exp
e [TypeBase Shape NoUniqueness]
rowtypes
internaliseLambda (E.Lambda [PatBase Info VName ParamType]
params Exp
body Maybe (TypeExp Exp VName)
_ (Info (RetType [VName]
_ TypeBase Exp Uniqueness
rettype)) SrcLoc
_) [TypeBase Shape NoUniqueness]
rowtypes =
  [PatBase Info VName ParamType]
-> [TypeBase Shape NoUniqueness]
-> ([LParam SOACS]
    -> InternaliseM
         ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness]))
-> InternaliseM
     ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness])
forall a.
[PatBase Info VName ParamType]
-> [TypeBase Shape NoUniqueness]
-> ([LParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLambdaParams [PatBase Info VName ParamType]
params [TypeBase Shape NoUniqueness]
rowtypes (([LParam SOACS]
  -> InternaliseM
       ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness]))
 -> InternaliseM
      ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness]))
-> ([LParam SOACS]
    -> InternaliseM
         ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness]))
-> InternaliseM
     ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness])
forall a b. (a -> b) -> a -> b
$ \[LParam SOACS]
params' -> do
    body' <- [Char] -> Exp -> InternaliseM (Body SOACS)
internaliseBody [Char]
"lam" Exp
body
    rettype' <- internaliseLambdaReturnType rettype =<< bodyExtType body'
    pure (params', body', rettype')
internaliseLambda Exp
e [TypeBase Shape NoUniqueness]
_ = [Char]
-> InternaliseM
     ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness])
forall a. HasCallStack => [Char] -> a
error ([Char]
 -> InternaliseM
      ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness]))
-> [Char]
-> InternaliseM
     ([LParam SOACS], Body SOACS, [TypeBase Shape NoUniqueness])
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseLambda: unexpected expression:\n" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Exp -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Exp
e

internaliseLambdaCoerce :: E.Exp -> [Type] -> InternaliseM (I.Lambda SOACS)
internaliseLambdaCoerce :: Exp -> [TypeBase Shape NoUniqueness] -> InternaliseM (Lambda SOACS)
internaliseLambdaCoerce Exp
lam [TypeBase Shape NoUniqueness]
argtypes = do
  (params, body, rettype) <- InternaliseLambda
internaliseLambda Exp
lam [TypeBase Shape NoUniqueness]
argtypes
  mkLambda params $
    ensureResultShape
      (ErrorMsg [ErrorString "unexpected lambda result size"])
      rettype
      =<< bodyBind body

-- | Overloaded operators are treated here.
isOverloadedFunction ::
  E.QualName VName ->
  String ->
  Maybe ([(E.StructType, [SubExp])] -> InternaliseM [SubExp])
isOverloadedFunction :: QualName VName
-> [Char]
-> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
isOverloadedFunction QualName VName
qname [Char]
desc = do
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ VName -> Int
baseTag (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qname) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag
  [Char] -> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
handle ([Char]
 -> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp]))
-> [Char]
-> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ VName -> [Char]
baseString (VName -> [Char]) -> VName -> [Char]
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qname
  where
    -- Handle equality and inequality specially, to treat the case of
    -- arrays.
    handle :: [Char] -> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
handle [Char]
op
      | Just SubExp -> InternaliseM [SubExp]
cmp_f <- [Char] -> Maybe (SubExp -> InternaliseM [SubExp])
isEqlOp [Char]
op = ([(StructType, [SubExp])] -> InternaliseM [SubExp])
-> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([(StructType, [SubExp])] -> InternaliseM [SubExp])
 -> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp]))
-> ([(StructType, [SubExp])] -> InternaliseM [SubExp])
-> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[(StructType
_, [SubExp]
xe'), (StructType
_, [SubExp]
ye')] -> do
          rs <- (SubExp -> SubExp -> InternaliseM SubExp)
-> [SubExp] -> [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM SubExp -> SubExp -> InternaliseM SubExp
doComparison [SubExp]
xe' [SubExp]
ye'
          cmp_f =<< letSubExp "eq" =<< eAll rs
      where
        isEqlOp :: [Char] -> Maybe (SubExp -> InternaliseM [SubExp])
isEqlOp [Char]
"!=" = (SubExp -> InternaliseM [SubExp])
-> Maybe (SubExp -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((SubExp -> InternaliseM [SubExp])
 -> Maybe (SubExp -> InternaliseM [SubExp]))
-> (SubExp -> InternaliseM [SubExp])
-> Maybe (SubExp -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \SubExp
eq ->
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (PrimType -> UnOp
I.Neg PrimType
I.Bool) SubExp
eq
        isEqlOp [Char]
"==" = (SubExp -> InternaliseM [SubExp])
-> Maybe (SubExp -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((SubExp -> InternaliseM [SubExp])
 -> Maybe (SubExp -> InternaliseM [SubExp]))
-> (SubExp -> InternaliseM [SubExp])
-> Maybe (SubExp -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \SubExp
eq ->
          [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp
eq]
        isEqlOp [Char]
_ = Maybe (SubExp -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

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

                    -- Compare the elements.
                    cmp_lam <- cmpOpLambda $ I.CmpEq (elemType x_t)
                    cmps <-
                      letExp "cmps" $
                        I.Op $
                          I.Screma x_num_elems [x_flat, y_flat] (I.mapSOAC cmp_lam)

                    -- Check that all were equal.
                    and_lam <- binOpLambda I.LogAnd I.Bool
                    reduce <- I.reduceSOAC [Reduce Commutative and_lam [constant True]]
                    all_equal <- letSubExp "all_equal" $ I.Op $ I.Screma x_num_elems [cmps] reduce
                    pure $ subExpsRes [all_equal]

              letSubExp "arrays_equal"
                =<< eIf (eSubExp shapes_match) compare_elems_body (resultBodyM [constant False])
    handle [Char]
name
      | Just BinOp
bop <- (BinOp -> Bool) -> [BinOp] -> Maybe BinOp
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (([Char]
name [Char] -> [Char] -> Bool
forall a. Eq a => a -> a -> Bool
==) ([Char] -> Bool) -> (BinOp -> [Char]) -> BinOp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BinOp -> [Char]
forall a. Pretty a => a -> [Char]
prettyString) [BinOp
forall a. Bounded a => a
minBound .. BinOp
forall a. Bounded a => a
maxBound :: E.BinOp] =
          ([(StructType, [SubExp])] -> InternaliseM [SubExp])
-> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([(StructType, [SubExp])] -> InternaliseM [SubExp])
 -> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp]))
-> ([(StructType, [SubExp])] -> InternaliseM [SubExp])
-> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[(StructType
x_t, [SubExp
x']), (StructType
y_t, [SubExp
y'])] ->
            case (StructType
x_t, StructType
y_t) of
              (E.Scalar (E.Prim PrimType
t1), E.Scalar (E.Prim PrimType
t2)) ->
                [Char]
-> BinOp
-> SubExp
-> SubExp
-> PrimType
-> PrimType
-> InternaliseM [SubExp]
internaliseBinOp [Char]
desc BinOp
bop SubExp
x' SubExp
y' PrimType
t1 PrimType
t2
              (StructType, StructType)
_ -> [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise.internaliseExp: non-primitive type in BinOp."
    handle [Char]
_ = Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

-- | Handle intrinsic functions.  These are only allowed to be called
-- in the prelude, and their internalisation may involve inspecting
-- the AST.
isIntrinsicFunction ::
  E.QualName VName ->
  [E.Exp] ->
  Maybe (String -> InternaliseM [SubExp])
isIntrinsicFunction :: QualName VName -> [Exp] -> Maybe ([Char] -> InternaliseM [SubExp])
isIntrinsicFunction QualName VName
qname [Exp]
args = do
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ VName -> Int
baseTag (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qname) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag
  let handlers :: [[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])]
handlers =
        [ [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
forall {a}.
(Eq a, IsString a) =>
[Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleSign,
          [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
forall {f :: * -> *}.
Applicative f =>
[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM (f SubExp))
handleOps,
          [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
forall {a}.
(IsString a, Eq a) =>
[Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleSOACs,
          [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
forall {a}.
(Eq a, IsString a) =>
[Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleAccs,
          [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
forall {a}.
(IsString a, Eq a) =>
[Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleAD,
          [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
forall {a}.
(IsString a, Eq a) =>
[Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleRest
        ]
  [Maybe ([Char] -> InternaliseM [SubExp])]
-> Maybe ([Char] -> InternaliseM [SubExp])
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, MonadPlus m) =>
t (m a) -> m a
msum [[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
h [Exp]
args ([Char] -> Maybe ([Char] -> InternaliseM [SubExp]))
-> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ VName -> [Char]
baseString (VName -> [Char]) -> VName -> [Char]
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qname | [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
h <- [[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])]
handlers]
  where
    handleSign :: [Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleSign [Exp
x] a
"sign_i8" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
I.Int8 Exp
x
    handleSign [Exp
x] a
"sign_i16" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
I.Int16 Exp
x
    handleSign [Exp
x] a
"sign_i32" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
I.Int32 Exp
x
    handleSign [Exp
x] a
"sign_i64" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
I.Int64 Exp
x
    handleSign [Exp
x] a
"unsign_i8" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
I.Int8 Exp
x
    handleSign [Exp
x] a
"unsign_i16" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
I.Int16 Exp
x
    handleSign [Exp
x] a
"unsign_i32" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
I.Int32 Exp
x
    handleSign [Exp
x] a
"unsign_i64" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
I.Int64 Exp
x
    handleSign [Exp]
_ a
_ = Maybe ([Char] -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

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

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

    handleAccs :: [Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleAccs [Exp
dest, Exp
f, Exp
bs] a
"scatter_stream" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char]
-> Exp -> Maybe (Exp, Exp) -> Exp -> Exp -> InternaliseM [SubExp]
internaliseStreamAcc [Char]
desc Exp
dest Maybe (Exp, Exp)
forall a. Maybe a
Nothing Exp
f Exp
bs
    handleAccs [Exp
dest, Exp
op, Exp
ne, Exp
f, Exp
bs] a
"hist_stream" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
      [Char]
-> Exp -> Maybe (Exp, Exp) -> Exp -> Exp -> InternaliseM [SubExp]
internaliseStreamAcc [Char]
desc Exp
dest ((Exp, Exp) -> Maybe (Exp, Exp)
forall a. a -> Maybe a
Just (Exp
op, Exp
ne)) Exp
f Exp
bs
    handleAccs [Exp
acc, Exp
i, Exp
v] a
"acc_write" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      acc' <- [VName] -> VName
forall a. HasCallStack => [a] -> a
head ([VName] -> VName) -> InternaliseM [VName] -> InternaliseM VName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"acc" Exp
acc
      i' <- internaliseExp1 "acc_i" i
      vs <- internaliseExp "acc_v" v
      fmap pure $ letSubExp desc $ BasicOp $ UpdateAcc Safe acc' [i'] vs
    handleAccs [Exp]
_ a
_ = Maybe ([Char] -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

    handleAD :: [Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleAD [Exp
f, Exp
x, Exp
v] a
fname
      | a
fname a -> [a] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [a
"jvp2", a
"vjp2"] = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
          x' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"ad_x" Exp
x
          v' <- internaliseExp "ad_v" v
          lam <- internaliseLambdaCoerce f =<< mapM subExpType x'
          fmap (map I.Var) . letTupExp desc . Op $
            case fname of
              a
"jvp2" -> [SubExp] -> [SubExp] -> Lambda SOACS -> SOAC SOACS
forall rep. [SubExp] -> [SubExp] -> Lambda rep -> SOAC rep
JVP [SubExp]
x' [SubExp]
v' Lambda SOACS
lam
              a
_ -> [SubExp] -> [SubExp] -> Lambda SOACS -> SOAC SOACS
forall rep. [SubExp] -> [SubExp] -> Lambda rep -> SOAC rep
VJP [SubExp]
x' [SubExp]
v' Lambda SOACS
lam
    handleAD [Exp]
_ a
_ = Maybe ([Char] -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

    handleRest :: [Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleRest [Exp
a, Exp
si, Exp
v] a
"scatter" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ Int -> Exp -> Exp -> Exp -> [Char] -> InternaliseM [SubExp]
scatterF Int
1 Exp
a Exp
si Exp
v
    handleRest [Exp
a, Exp
si, Exp
v] a
"scatter_2d" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ Int -> Exp -> Exp -> Exp -> [Char] -> InternaliseM [SubExp]
scatterF Int
2 Exp
a Exp
si Exp
v
    handleRest [Exp
a, Exp
si, Exp
v] a
"scatter_3d" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ Int -> Exp -> Exp -> Exp -> [Char] -> InternaliseM [SubExp]
scatterF Int
3 Exp
a Exp
si Exp
v
    handleRest [Exp
n, Exp
m, Exp
arr] a
"unflatten" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"unflatten_arr" Exp
arr
      n' <- internaliseExp1 "n" n
      m' <- internaliseExp1 "m" m
      -- Each dimension must be nonnegative, and the unflattened
      -- dimension needs to have the same number of elements as the
      -- original dimension.
      old_dim <- I.arraysSize 0 <$> mapM lookupType arrs
      dim_ok <-
        letSubExp "dim_ok" <=< toExp $
          pe64 old_dim .==. pe64 n'
            * pe64 m'
              .&&. pe64 n'
              .>=. 0
              .&&. pe64 m'
              .>=. 0
      dim_ok_cert <-
        assert "dim_ok_cert" dim_ok $
          ErrorMsg
            [ "Cannot unflatten array of shape [",
              ErrorVal int64 old_dim,
              "] to array of shape [",
              ErrorVal int64 n',
              "][",
              ErrorVal int64 m',
              "]"
            ]
      certifying dim_ok_cert $
        forM arrs $ \VName
arr' -> do
          arr_t <- VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
arr'
          letSubExp desc . I.BasicOp $
            I.Reshape arr' $
              reshapeAll (I.arrayShape arr_t) $
                reshapeOuter (I.Shape [n', m']) 1 $
                  I.arrayShape arr_t
    handleRest [Exp
arr] a
"manifest" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"flatten_arr" Exp
arr
      forM arrs $ \VName
arr' -> do
        r <- TypeBase Shape NoUniqueness -> Int
forall shape u. ArrayShape shape => TypeBase shape u -> Int
I.arrayRank (TypeBase Shape NoUniqueness -> Int)
-> InternaliseM (TypeBase Shape NoUniqueness) -> InternaliseM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
arr'
        if r == 0
          then pure $ I.Var arr'
          else letSubExp desc $ I.BasicOp $ I.Manifest arr' [0 .. r - 1]
    handleRest [Exp
arr] a
"flatten" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"flatten_arr" Exp
arr
      forM arrs $ \VName
arr' -> do
        arr_t <- VName -> InternaliseM (TypeBase Shape NoUniqueness)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
arr'
        let n = Int -> TypeBase Shape NoUniqueness -> SubExp
forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
0 TypeBase Shape NoUniqueness
arr_t
            m = Int -> TypeBase Shape NoUniqueness -> SubExp
forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
1 TypeBase Shape NoUniqueness
arr_t
        k <- letSubExp "flat_dim" $ I.BasicOp $ I.BinOp (Mul Int64 I.OverflowUndef) n m
        letSubExp desc . I.BasicOp $
          I.Reshape arr' $
            reshapeAll (I.arrayShape arr_t) $
              reshapeOuter (I.Shape [k]) 2 $
                I.arrayShape arr_t
    handleRest [Exp
x, Exp
y] a
"concat" = ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([Char] -> InternaliseM [SubExp])
 -> Maybe ([Char] -> InternaliseM [SubExp]))
-> ([Char] -> InternaliseM [SubExp])
-> Maybe ([Char] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
      xs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"concat_x" Exp
x
      ys <- internaliseExpToVars "concat_y" y
      outer_size <- arraysSize 0 <$> mapM lookupType xs
      let sumdims SubExp
xsize SubExp
ysize =
            [Char] -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"conc_tmp" (Exp (Rep m) -> m SubExp) -> Exp (Rep m) -> m SubExp
forall a b. (a -> b) -> a -> b
$
              BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$
                BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Add IntType
I.Int64 Overflow
I.OverflowUndef) SubExp
xsize SubExp
ysize
      ressize <-
        foldM sumdims outer_size
          =<< mapM (fmap (arraysSize 0) . mapM lookupType) [ys]

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

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

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

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

      si_w <- I.arraysSize 0 <$> mapM lookupType si'
      sv_ts <- mapM lookupType svs

      svs' <- forM (zip svs sv_ts) $ \(VName
sv, TypeBase Shape NoUniqueness
sv_t) -> do
        let sv_shape :: Shape
sv_shape = TypeBase Shape NoUniqueness -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape TypeBase Shape NoUniqueness
sv_t
            sv_w :: SubExp
sv_w = Int -> TypeBase Shape NoUniqueness -> SubExp
forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
0 TypeBase Shape NoUniqueness
sv_t

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

      indexType <- fmap rowType <$> mapM lookupType si'
      indexName <- mapM (\TypeBase Shape NoUniqueness
_ -> [Char] -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"write_index") indexType
      valueNames <- replicateM (length sv_ts) $ newVName "write_value"

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

      -- This body is boring right now, as every input is exactly the output.
      -- But it can get funky later on if fused with something else.
      body <- localScope (scopeOfLParams bodyParams) . buildBody_ $ do
        let outs = [[VName]] -> [VName]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Int -> [VName] -> [[VName]]
forall a. Int -> a -> [a]
replicate ([VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
valueNames) [VName]
indexName) [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
valueNames
        results <- forM outs $ \VName
name ->
          [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"write_res" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
name
        ensureResultShape
          "scatter value has wrong size"
          bodyTypes
          (subExpsRes results)

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

      let sa_ws = (TypeBase Shape NoUniqueness -> Shape)
-> [TypeBase Shape NoUniqueness] -> [Shape]
forall a b. (a -> b) -> [a] -> [b]
map ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
I.Shape ([SubExp] -> Shape)
-> (TypeBase Shape NoUniqueness -> [SubExp])
-> TypeBase Shape NoUniqueness
-> Shape
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
take Int
dim ([SubExp] -> [SubExp])
-> (TypeBase Shape NoUniqueness -> [SubExp])
-> TypeBase Shape NoUniqueness
-> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase Shape NoUniqueness -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
arrayDims) [TypeBase Shape NoUniqueness]
sa_ts
          spec = [Shape] -> [Int] -> [VName] -> [(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
      letTupExp' desc $ I.Op $ I.Scatter si_w sivs spec lam

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

          span_and_lower <- letSubExp "span_and_lower" $ I.BasicOp $ I.BinOp (I.Add Int64 I.OverflowUndef) spn lower
          span_and_upper <- letSubExp "span_and_upper" $ I.BasicOp $ I.BinOp (I.Add Int64 I.OverflowUndef) spn upper

          lower' <- letSubExp "minimum" $ I.BasicOp $ I.BinOp (I.UMin Int64) span_and_lower lower
          upper' <- letSubExp "maximum" $ I.BasicOp $ I.BinOp (I.UMax Int64) span_and_upper upper

          pure (lower', upper')
      )
      (offset', offset')
      slices'
  min_in_bounds <- letSubExp "min_in_bounds" $ I.BasicOp $ I.CmpOp (I.CmpUle Int64) (intConst Int64 0) min_bound
  max_in_bounds <- letSubExp "max_in_bounds" $ I.BasicOp $ I.CmpOp (I.CmpUlt Int64) max_bound old_dim

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

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

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

          span_and_lower <- letSubExp "span_and_lower" $ I.BasicOp $ I.BinOp (I.Add Int64 I.OverflowUndef) spn lower
          span_and_upper <- letSubExp "span_and_upper" $ I.BasicOp $ I.BinOp (I.Add Int64 I.OverflowUndef) spn upper

          lower' <- letSubExp "minimum" $ I.BasicOp $ I.BinOp (I.UMin Int64) span_and_lower lower
          upper' <- letSubExp "maximum" $ I.BasicOp $ I.BinOp (I.UMax Int64) span_and_upper upper

          pure (lower', upper')
      )
      (offset', offset')
      slices'
  min_in_bounds <- letSubExp "min_in_bounds" $ I.BasicOp $ I.CmpOp (I.CmpUle Int64) (intConst Int64 0) min_bound
  max_in_bounds <- letSubExp "max_in_bounds" $ I.BasicOp $ I.CmpOp (I.CmpUlt Int64) max_bound old_dim

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

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

funcall ::
  String ->
  QualName VName ->
  [SubExp] ->
  InternaliseM [SubExp]
funcall :: [Char] -> QualName VName -> [SubExp] -> InternaliseM [SubExp]
funcall [Char]
desc (QualName [VName]
_ VName
fname) [SubExp]
args = do
  (shapes, value_paramts, fun_params, rettype_fun) <- VName -> InternaliseM FunInfo
lookupFunction VName
fname
  argts <- mapM subExpType args

  shapeargs <- argShapes shapes fun_params argts
  let diets =
        Int -> Diet -> [Diet]
forall a. Int -> a -> [a]
replicate ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
shapeargs) Diet
I.ObservePrim
          [Diet] -> [Diet] -> [Diet]
forall a. [a] -> [a] -> [a]
++ (DeclType -> Diet) -> [DeclType] -> [Diet]
forall a b. (a -> b) -> [a] -> [b]
map DeclType -> Diet
forall shape. TypeBase shape Uniqueness -> Diet
I.diet [DeclType]
value_paramts
  args' <-
    ensureArgShapes
      "function arguments of wrong shape"
      (map I.paramName fun_params)
      (map I.paramType fun_params)
      (shapeargs ++ args)
  argts' <- mapM subExpType args'
  case rettype_fun $ zip args' argts' of
    Maybe [(TypeBase ExtShape Uniqueness, RetAls)]
Nothing ->
      [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> ([[Char]] -> [Char]) -> [[Char]] -> InternaliseM [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[Char]] -> [Char]
unlines ([[Char]] -> InternaliseM [SubExp])
-> [[Char]] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
        [ [Char]
"Cannot apply "
            [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> VName -> [Char]
forall a. Pretty a => a -> [Char]
prettyString VName
fname
            [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" to "
            [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> [Char]
forall a. Show a => a -> [Char]
show ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
args')
            [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" arguments",
          [Char]
" " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [SubExp] -> [Char]
forall a. Pretty a => a -> [Char]
prettyString [SubExp]
args',
          [Char]
"of types",
          [Char]
" " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [TypeBase Shape NoUniqueness] -> [Char]
forall a. Pretty a => a -> [Char]
prettyString [TypeBase Shape NoUniqueness]
argts',
          [Char]
"Function has " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> [Char]
forall a. Show a => a -> [Char]
show ([Param DeclType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Param DeclType]
fun_params) [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" parameters",
          [Char]
" " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Param DeclType] -> [Char]
forall a. Pretty a => a -> [Char]
prettyString [Param DeclType]
fun_params
        ]
    Just [(TypeBase ExtShape Uniqueness, RetAls)]
ts -> do
      safety <- InternaliseM Safety
askSafety
      attrs <- asks envAttrs
      attributing attrs . letValExp' desc $
        I.Apply (internaliseFunName fname) (zip args' diets) ts safety

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

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

  forM_ (M.toList $ mconcat $ zipWith combine ts ses_ts) $ \(VName
v, SubExp
se) ->
    [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
v] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se

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

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

  add_lam_x_params <-
    replicateM k $ newParam "x" (I.Prim int64)
  add_lam_y_params <-
    replicateM k $ newParam "y" (I.Prim int64)
  add_lam_body <- runBodyBuilder $
    localScope (scopeOfLParams $ add_lam_x_params ++ add_lam_y_params) $
      fmap subExpsRes $
        forM (zip add_lam_x_params add_lam_y_params) $ \(Param (TypeBase Shape NoUniqueness)
x, Param (TypeBase Shape NoUniqueness)
y) ->
          [Char]
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"z" (Exp (Rep (BuilderT SOACS (State VNameSource)))
 -> BuilderT SOACS (State VNameSource) SubExp)
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) SubExp
forall a b. (a -> b) -> a -> b
$
            BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource))))
-> BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
forall a b. (a -> b) -> a -> b
$
              BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp
                (IntType -> Overflow -> BinOp
I.Add IntType
Int64 Overflow
I.OverflowUndef)
                (VName -> SubExp
I.Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param (TypeBase Shape NoUniqueness) -> VName
forall dec. Param dec -> VName
I.paramName Param (TypeBase Shape NoUniqueness)
x)
                (VName -> SubExp
I.Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param (TypeBase Shape NoUniqueness) -> VName
forall dec. Param dec -> VName
I.paramName Param (TypeBase Shape NoUniqueness)
y)
  let add_lam =
        I.Lambda
          { lambdaBody :: Body SOACS
I.lambdaBody = Body SOACS
add_lam_body,
            lambdaParams :: [LParam SOACS]
I.lambdaParams = [Param (TypeBase Shape NoUniqueness)]
add_lam_x_params [Param (TypeBase Shape NoUniqueness)]
-> [Param (TypeBase Shape NoUniqueness)]
-> [Param (TypeBase Shape NoUniqueness)]
forall a. [a] -> [a] -> [a]
++ [Param (TypeBase Shape NoUniqueness)]
add_lam_y_params,
            lambdaReturnType :: [TypeBase Shape NoUniqueness]
I.lambdaReturnType = Int -> TypeBase Shape NoUniqueness -> [TypeBase Shape NoUniqueness]
forall a. Int -> a -> [a]
replicate Int
k (TypeBase Shape NoUniqueness -> [TypeBase Shape NoUniqueness])
-> TypeBase Shape NoUniqueness -> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase Shape NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64
          }
      nes = Int -> SubExp -> [SubExp]
forall a. Int -> a -> [a]
replicate ([VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
increments) (SubExp -> [SubExp]) -> SubExp -> [SubExp]
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0

  scan <- I.scanSOAC [I.Scan add_lam nes]
  all_offsets <- letTupExp "offsets" $ I.Op $ I.Screma w increments 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.
  last_index <- letSubExp "last_index" $ I.BasicOp $ I.BinOp (I.Sub Int64 OverflowUndef) w $ constant (1 :: Int64)
  let nonempty_body = Builder SOACS Result -> InternaliseM (Body SOACS)
forall rep (m :: * -> *) somerep.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
 SameScope somerep rep) =>
Builder rep Result -> m (Body rep)
runBodyBuilder (Builder SOACS Result -> InternaliseM (Body SOACS))
-> Builder SOACS Result -> InternaliseM (Body SOACS)
forall a b. (a -> b) -> a -> b
$
        ([SubExp] -> Result)
-> BuilderT SOACS (State VNameSource) [SubExp]
-> Builder SOACS Result
forall a b.
(a -> b)
-> BuilderT SOACS (State VNameSource) a
-> BuilderT SOACS (State VNameSource) b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [SubExp] -> Result
subExpsRes (BuilderT SOACS (State VNameSource) [SubExp]
 -> Builder SOACS Result)
-> BuilderT SOACS (State VNameSource) [SubExp]
-> Builder SOACS Result
forall a b. (a -> b) -> a -> b
$
          [VName]
-> (VName -> BuilderT SOACS (State VNameSource) SubExp)
-> BuilderT SOACS (State VNameSource) [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
all_offsets ((VName -> BuilderT SOACS (State VNameSource) SubExp)
 -> BuilderT SOACS (State VNameSource) [SubExp])
-> (VName -> BuilderT SOACS (State VNameSource) SubExp)
-> BuilderT SOACS (State VNameSource) [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
offset_array ->
            [Char]
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"last_offset" (Exp (Rep (BuilderT SOACS (State VNameSource)))
 -> BuilderT SOACS (State VNameSource) SubExp)
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource))))
-> BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp
I.Index VName
offset_array (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ [DimIndex SubExp] -> Slice SubExp
forall d. [DimIndex d] -> Slice d
Slice [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
I.DimFix SubExp
last_index]
      empty_body = [SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM ([SubExp] -> InternaliseM (Body (Rep InternaliseM)))
-> [SubExp] -> InternaliseM (Body (Rep InternaliseM))
forall a b. (a -> b) -> a -> b
$ Int -> SubExp -> [SubExp]
forall a. Int -> a -> [a]
replicate Int
k (SubExp -> [SubExp]) -> SubExp -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)
  is_empty <- letSubExp "is_empty" $ I.BasicOp $ I.CmpOp (CmpEq int64) w $ constant (0 :: Int64)
  sizes <-
    letTupExp "partition_size" =<< eIf (eSubExp is_empty) empty_body nonempty_body

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

  -- Create scratch arrays for the result.
  blanks <- forM arr_ts $ \TypeBase Shape NoUniqueness
arr_t ->
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
"partition_dest" (Exp (Rep InternaliseM) -> InternaliseM VName)
-> Exp (Rep InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$
      BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
        PrimType -> [SubExp] -> BasicOp
Scratch (TypeBase Shape NoUniqueness -> PrimType
forall shape u. TypeBase shape u -> PrimType
I.elemType TypeBase Shape NoUniqueness
arr_t) (SubExp
w SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
drop Int
1 (TypeBase Shape NoUniqueness -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
I.arrayDims TypeBase Shape NoUniqueness
arr_t))

  -- Now write into the result.
  write_lam <- do
    c_param <- newParam "c" (I.Prim int64)
    offset_params <- replicateM k $ newParam "offset" (I.Prim int64)
    value_params <- mapM (newParam "v" . I.rowType) arr_ts
    (offset, offset_stms) <-
      collectStms $
        mkOffsetLambdaBody
          (map I.Var sizes)
          (I.Var $ I.paramName c_param)
          0
          offset_params
    pure
      I.Lambda
        { I.lambdaParams = c_param : offset_params ++ value_params,
          I.lambdaReturnType =
            replicate (length arr_ts) (I.Prim int64)
              ++ map I.rowType arr_ts,
          I.lambdaBody =
            mkBody offset_stms $
              replicate (length arr_ts) (subExpRes offset)
                ++ I.varsRes (map I.paramName value_params)
        }
  let spec = [Shape] -> [Int] -> [VName] -> [(Shape, Int, VName)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 (Shape -> [Shape]
forall a. a -> [a]
repeat (Shape -> [Shape]) -> Shape -> [Shape]
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Shape
forall d. [d] -> ShapeBase d
I.Shape [SubExp
w]) (Int -> [Int]
forall a. a -> [a]
repeat Int
1) [VName]
blanks
  results <-
    letTupExp "partition_res" . I.Op $
      I.Scatter w (classes : all_offsets ++ arrs) spec write_lam
  sizes' <-
    letSubExp "partition_sizes" $
      I.BasicOp $
        I.ArrayLit (map I.Var sizes) $
          I.Prim int64
  pure (map I.Var results, [sizes'])
  where
    mkOffsetLambdaBody ::
      [SubExp] ->
      SubExp ->
      Int ->
      [I.LParam SOACS] ->
      InternaliseM SubExp
    mkOffsetLambdaBody :: [SubExp] -> SubExp -> Int -> [LParam SOACS] -> InternaliseM SubExp
mkOffsetLambdaBody [SubExp]
_ SubExp
_ Int
_ [] =
      SubExp -> InternaliseM SubExp
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> InternaliseM SubExp) -> SubExp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (-Int64
1 :: Int64)
    mkOffsetLambdaBody [SubExp]
sizes SubExp
c Int
i (LParam SOACS
p : [LParam SOACS]
ps) = do
      is_this_one <-
        [Char] -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"is_this_one" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
CmpEq PrimType
int64) SubExp
c (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$
              IntType -> Integer -> SubExp
intConst IntType
Int64 (Integer -> SubExp) -> Integer -> SubExp
forall a b. (a -> b) -> a -> b
$
                Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
i
      next_one <- mkOffsetLambdaBody sizes c (i + 1) ps
      this_one <-
        letSubExp "this_offset"
          =<< foldBinOp
            (Add Int64 OverflowUndef)
            (constant (-1 :: Int64))
            (I.Var (I.paramName p) : take i sizes)
      letSubExp "total_res"
        =<< eIf
          (eSubExp is_this_one)
          (resultBodyM [this_one])
          (resultBodyM [next_one])

sizeExpForError :: E.Size -> InternaliseM [ErrorMsgPart SubExp]
sizeExpForError :: Exp -> InternaliseM [ErrorMsgPart SubExp]
sizeExpForError Exp
e
  | Exp
e Exp -> Exp -> Bool
forall a. Eq a => a -> a -> Bool
== Exp
anySize = [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [ErrorMsgPart SubExp
"[]"]
  | Bool
otherwise = do
      e' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"size" Exp
e
      pure ["[", ErrorVal int64 e', "]"]

typeExpForError :: E.TypeBase Size u -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError :: forall u. TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError (E.Scalar (E.Prim PrimType
t)) = [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Text -> ErrorMsgPart SubExp
forall a. Text -> ErrorMsgPart a
ErrorString (Text -> ErrorMsgPart SubExp) -> Text -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> Text
forall a. Pretty a => a -> Text
prettyText PrimType
t]
typeExpForError (E.Scalar (E.TypeVar u
_ QualName VName
v [TypeArg Exp]
args)) = do
  args' <- [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp])
-> InternaliseM [[ErrorMsgPart SubExp]]
-> InternaliseM [ErrorMsgPart SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeArg Exp -> InternaliseM [ErrorMsgPart SubExp])
-> [TypeArg Exp] -> InternaliseM [[ErrorMsgPart SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM TypeArg Exp -> InternaliseM [ErrorMsgPart SubExp]
onArg [TypeArg Exp]
args
  pure $ intersperse " " $ ErrorString (prettyText v) : args'
  where
    onArg :: TypeArg Exp -> InternaliseM [ErrorMsgPart SubExp]
onArg (TypeArgDim Exp
d) = Exp -> InternaliseM [ErrorMsgPart SubExp]
sizeExpForError Exp
d
    onArg (TypeArgType StructType
t) = StructType -> InternaliseM [ErrorMsgPart SubExp]
forall u. TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError StructType
t
typeExpForError (E.Scalar (E.Record Map Name (TypeBase Exp u)
fs))
  | Just [TypeBase Exp u]
ts <- Map Name (TypeBase Exp u) -> Maybe [TypeBase Exp u]
forall a. Map Name a -> Maybe [a]
E.areTupleFields Map Name (TypeBase Exp u)
fs = do
      ts' <- (TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp])
-> [TypeBase Exp u] -> InternaliseM [[ErrorMsgPart SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
forall u. TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError [TypeBase Exp u]
ts
      pure $ ["("] ++ intercalate [", "] ts' ++ [")"]
  | Bool
otherwise = do
      fs' <- ((Name, TypeBase Exp u) -> InternaliseM [ErrorMsgPart SubExp])
-> [(Name, TypeBase Exp u)] -> InternaliseM [[ErrorMsgPart SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Name, TypeBase Exp u) -> InternaliseM [ErrorMsgPart SubExp]
forall {a} {u}.
Pretty a =>
(a, TypeBase Exp u) -> InternaliseM [ErrorMsgPart SubExp]
onField ([(Name, TypeBase Exp u)] -> InternaliseM [[ErrorMsgPart SubExp]])
-> [(Name, TypeBase Exp u)] -> InternaliseM [[ErrorMsgPart SubExp]]
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase Exp u) -> [(Name, TypeBase Exp u)]
forall k a. Map k a -> [(k, a)]
M.toList Map Name (TypeBase Exp u)
fs
      pure $ ["{"] ++ intercalate [", "] fs' ++ ["}"]
  where
    onField :: (a, TypeBase Exp u) -> InternaliseM [ErrorMsgPart SubExp]
onField (a
k, TypeBase Exp u
te) =
      (Text -> ErrorMsgPart SubExp
forall a. Text -> ErrorMsgPart a
ErrorString (a -> Text
forall a. Pretty a => a -> Text
prettyText a
k Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
": ") ErrorMsgPart SubExp
-> [ErrorMsgPart SubExp] -> [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
<$> TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
forall u. TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeBase Exp u
te
typeExpForError (E.Array u
_ Shape Exp
shape ScalarTypeBase Exp NoUniqueness
et) = do
  shape' <- [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
forall a. Monoid a => [a] -> a
mconcat ([[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp])
-> InternaliseM [[ErrorMsgPart SubExp]]
-> InternaliseM [ErrorMsgPart SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> InternaliseM [ErrorMsgPart SubExp])
-> [Exp] -> InternaliseM [[ErrorMsgPart SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Exp -> InternaliseM [ErrorMsgPart SubExp]
sizeExpForError (Shape Exp -> [Exp]
forall dim. Shape dim -> [dim]
E.shapeDims Shape Exp
shape)
  et' <- typeExpForError $ Scalar et
  pure $ shape' ++ et'
typeExpForError (E.Scalar (E.Sum Map Name [TypeBase Exp u]
cs)) = do
  cs' <- ((Name, [TypeBase Exp u]) -> InternaliseM [ErrorMsgPart SubExp])
-> [(Name, [TypeBase Exp u])]
-> InternaliseM [[ErrorMsgPart SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Name, [TypeBase Exp u]) -> InternaliseM [ErrorMsgPart SubExp]
forall {a} {u}.
Pretty a =>
(a, [TypeBase Exp u]) -> InternaliseM [ErrorMsgPart SubExp]
onConstructor ([(Name, [TypeBase Exp u])]
 -> InternaliseM [[ErrorMsgPart SubExp]])
-> [(Name, [TypeBase Exp u])]
-> InternaliseM [[ErrorMsgPart SubExp]]
forall a b. (a -> b) -> a -> b
$ Map Name [TypeBase Exp u] -> [(Name, [TypeBase Exp u])]
forall k a. Map k a -> [(k, a)]
M.toList Map Name [TypeBase Exp u]
cs
  pure $ intercalate [" | "] cs'
  where
    onConstructor :: (a, [TypeBase Exp u]) -> InternaliseM [ErrorMsgPart SubExp]
onConstructor (a
c, [TypeBase Exp u]
ts) = do
      ts' <- (TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp])
-> [TypeBase Exp u] -> InternaliseM [[ErrorMsgPart SubExp]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
forall u. TypeBase Exp u -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError [TypeBase Exp u]
ts
      pure $ ErrorString ("#" <> prettyText c <> " ") : intercalate [" "] ts'
typeExpForError (E.Scalar Arrow {}) = [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [ErrorMsgPart SubExp
"#<fun>"]

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

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