{-# LANGUAGE TypeFamilies #-}

module Futhark.IR.Mem.Simplify
  ( simplifyProgGeneric,
    simplifyStmsGeneric,
    simpleGeneric,
    SimplifyMemory,
    memRuleBook,
  )
where

import Futhark.Analysis.SymbolTable qualified as ST
import Futhark.Analysis.UsageTable qualified as UT
import Futhark.Construct
import Futhark.IR.Mem
import Futhark.IR.Prop.Aliases (AliasedOp)
import Futhark.Optimise.Simplify qualified as Simplify
import Futhark.Optimise.Simplify.Engine qualified as Engine
import Futhark.Optimise.Simplify.Rep
import Futhark.Optimise.Simplify.Rule
import Futhark.Optimise.Simplify.Rules
import Futhark.Pass
import Futhark.Pass.ExplicitAllocations (simplifiable)

-- | Some constraints that must hold for the simplification rules to work.
type SimplifyMemory rep inner =
  ( Simplify.SimplifiableRep rep,
    LetDec rep ~ LetDecMem,
    ExpDec rep ~ (),
    BodyDec rep ~ (),
    CanBeWise (OpC rep),
    BuilderOps (Wise rep),
    OpReturns inner,
    ST.IndexOp (inner (Wise rep)),
    AliasedOp inner,
    Mem rep inner,
    CanBeWise inner,
    RephraseOp inner,
    ASTConstraints (inner (Engine.Wise rep))
  )

simpleGeneric ::
  (SimplifyMemory rep inner) =>
  (inner (Wise rep) -> UT.UsageTable) ->
  Simplify.SimplifyOp rep (inner (Wise rep)) ->
  Simplify.SimpleOps rep
simpleGeneric :: forall rep (inner :: * -> *).
SimplifyMemory rep inner =>
(inner (Wise rep) -> UsageTable)
-> SimplifyOp rep (inner (Wise rep)) -> SimpleOps rep
simpleGeneric = (inner (Wise rep) -> UsageTable)
-> (inner (Wise rep)
    -> SimpleM rep (inner (Wise rep), Stms (Wise rep)))
-> SimpleOps rep
forall rep (inner :: * -> *).
(SimplifiableRep rep, LetDec rep ~ LetDecMem, ExpDec rep ~ (),
 BodyDec rep ~ (), Mem (Wise rep) inner, CanBeWise inner,
 RephraseOp inner, IsOp inner, OpReturns inner, AliasedOp inner,
 IndexOp (inner (Wise rep))) =>
(inner (Wise rep) -> UsageTable)
-> (inner (Wise rep)
    -> SimpleM rep (inner (Wise rep), Stms (Wise rep)))
-> SimpleOps rep
simplifiable

simplifyProgGeneric ::
  (SimplifyMemory rep inner) =>
  RuleBook (Wise rep) ->
  Simplify.SimpleOps rep ->
  Prog rep ->
  PassM (Prog rep)
simplifyProgGeneric :: forall rep (inner :: * -> *).
SimplifyMemory rep inner =>
RuleBook (Wise rep)
-> SimpleOps rep -> Prog rep -> PassM (Prog rep)
simplifyProgGeneric RuleBook (Wise rep)
rules SimpleOps rep
ops =
  SimpleOps rep
-> RuleBook (Wise rep)
-> HoistBlockers rep
-> Prog rep
-> PassM (Prog rep)
forall rep.
SimplifiableRep rep =>
SimpleOps rep
-> RuleBook (Wise rep)
-> HoistBlockers rep
-> Prog rep
-> PassM (Prog rep)
Simplify.simplifyProg
    SimpleOps rep
ops
    RuleBook (Wise rep)
rules
    HoistBlockers rep
forall rep (inner :: * -> *).
(OpC rep ~ MemOp inner) =>
HoistBlockers rep
blockers {Engine.blockHoistBranch = blockAllocs}
  where
    blockAllocs :: SymbolTable rep -> p -> Stm rep -> Bool
blockAllocs SymbolTable rep
vtable p
_ (Let Pat (LetDec rep)
_ StmAux (ExpDec rep)
_ (Op Alloc {})) =
      Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ SymbolTable rep -> Bool
forall rep. SymbolTable rep -> Bool
ST.simplifyMemory SymbolTable rep
vtable
    -- Do not hoist statements that produce arrays.  This is
    -- because in the KernelsMem representation, multiple
    -- arrays can be located in the same memory block, and moving
    -- their creation out of a branch can thus cause memory
    -- corruption.  At this point in the compiler we have probably
    -- already moved all the array creations that matter.
    blockAllocs SymbolTable rep
_ p
_ (Let Pat (LetDec rep)
pat StmAux (ExpDec rep)
_ Exp rep
_) =
      Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (Type -> Bool) -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Type -> Bool
forall shape u. TypeBase shape u -> Bool
primType ([Type] -> Bool) -> [Type] -> Bool
forall a b. (a -> b) -> a -> b
$ Pat (LetDec rep) -> [Type]
forall dec. Typed dec => Pat dec -> [Type]
patTypes Pat (LetDec rep)
pat

simplifyStmsGeneric ::
  ( HasScope rep m,
    MonadFreshNames m,
    SimplifyMemory rep inner
  ) =>
  RuleBook (Wise rep) ->
  Simplify.SimpleOps rep ->
  Stms rep ->
  m (Stms rep)
simplifyStmsGeneric :: forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, MonadFreshNames m, SimplifyMemory rep inner) =>
RuleBook (Wise rep) -> SimpleOps rep -> Stms rep -> m (Stms rep)
simplifyStmsGeneric RuleBook (Wise rep)
rules SimpleOps rep
ops Stms rep
stms = do
  scope <- m (Scope rep)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  Simplify.simplifyStms ops rules blockers scope stms

isResultAlloc :: (OpC rep ~ MemOp op) => Engine.BlockPred rep
isResultAlloc :: forall rep (op :: * -> *). (OpC rep ~ MemOp op) => BlockPred rep
isResultAlloc SymbolTable rep
_ UsageTable
usage (Let (Pat [PatElem (LetDec rep)
pe]) StmAux (ExpDec rep)
_ (Op Alloc {})) =
  VName -> UsageTable -> Bool
UT.isInResult (PatElem (LetDec rep) -> VName
forall dec. PatElem dec -> VName
patElemName PatElem (LetDec rep)
pe) UsageTable
usage
isResultAlloc SymbolTable rep
_ UsageTable
_ Stm rep
_ = Bool
False

isAlloc :: (OpC rep ~ MemOp op) => Engine.BlockPred rep
isAlloc :: forall rep (op :: * -> *). (OpC rep ~ MemOp op) => BlockPred rep
isAlloc SymbolTable rep
_ UsageTable
_ (Let Pat (LetDec rep)
_ StmAux (ExpDec rep)
_ (Op Alloc {})) = Bool
True
isAlloc SymbolTable rep
_ UsageTable
_ Stm rep
_ = Bool
False

blockers ::
  (OpC rep ~ MemOp inner) =>
  Simplify.HoistBlockers rep
blockers :: forall rep (inner :: * -> *).
(OpC rep ~ MemOp inner) =>
HoistBlockers rep
blockers =
  HoistBlockers rep
forall {k} (rep :: k). HoistBlockers rep
Engine.noExtraHoistBlockers
    { Engine.blockHoistPar = isAlloc,
      Engine.blockHoistSeq = isResultAlloc,
      Engine.isAllocation = isAlloc mempty mempty
    }

-- If an allocation is statically known to be safe, then we can remove
-- the certificates on it.  This can help hoist things that would
-- otherwise be stuck inside loops or branches.
decertifySafeAlloc :: (SimplifyMemory rep inner) => TopDownRuleOp (Wise rep)
decertifySafeAlloc :: forall rep (inner :: * -> *).
SimplifyMemory rep inner =>
TopDownRuleOp (Wise rep)
decertifySafeAlloc TopDown (Wise rep)
_ Pat (LetDec (Wise rep))
pat (StmAux Certs
cs Attrs
attrs Provenance
_ ExpDec (Wise rep)
_) Op (Wise rep)
op
  | Certs
cs Certs -> Certs -> Bool
forall a. Eq a => a -> a -> Bool
/= Certs
forall a. Monoid a => a
mempty,
    [Mem Space
_] <- Pat (VarWisdom, LetDecMem) -> [Type]
forall dec. Typed dec => Pat dec -> [Type]
patTypes Pat (VarWisdom, LetDecMem)
Pat (LetDec (Wise rep))
pat,
    MemOp inner (Wise rep) -> Bool
forall rep. ASTRep rep => MemOp inner rep -> Bool
forall (op :: * -> *) rep. (IsOp op, ASTRep rep) => op rep -> Bool
safeOp Op (Wise rep)
MemOp inner (Wise rep)
op =
      RuleM (Wise rep) () -> Rule (Wise rep)
forall rep. RuleM rep () -> Rule rep
Simplify (RuleM (Wise rep) () -> Rule (Wise rep))
-> RuleM (Wise rep) () -> Rule (Wise rep)
forall a b. (a -> b) -> a -> b
$ Attrs -> RuleM (Wise rep) () -> RuleM (Wise rep) ()
forall (m :: * -> *) a. MonadBuilder m => Attrs -> m a -> m a
attributing Attrs
attrs (RuleM (Wise rep) () -> RuleM (Wise rep) ())
-> RuleM (Wise rep) () -> RuleM (Wise rep) ()
forall a b. (a -> b) -> a -> b
$ Pat (LetDec (Rep (RuleM (Wise rep))))
-> Exp (Rep (RuleM (Wise rep))) -> RuleM (Wise rep) ()
forall (m :: * -> *).
MonadBuilder m =>
Pat (LetDec (Rep m)) -> Exp (Rep m) -> m ()
letBind Pat (LetDec (Rep (RuleM (Wise rep))))
Pat (LetDec (Wise rep))
pat (Exp (Rep (RuleM (Wise rep))) -> RuleM (Wise rep) ())
-> Exp (Rep (RuleM (Wise rep))) -> RuleM (Wise rep) ()
forall a b. (a -> b) -> a -> b
$ Op (Wise rep) -> Exp (Wise rep)
forall rep. Op rep -> Exp rep
Op Op (Wise rep)
op
decertifySafeAlloc TopDown (Wise rep)
_ Pat (LetDec (Wise rep))
_ StmAux (ExpDec (Wise rep))
_ Op (Wise rep)
_ = Rule (Wise rep)
forall rep. Rule rep
Skip

--
-- copy(reshape(manifest(v0),s)) can be rewritten to just reshape(manifest(v0),s).
--
-- This is a pattern that can be produced by ExplicitAllocations when the
-- reshape would otherwise produce a layout that is not representable as an
-- LMAD. We have to be careful that the manifest writes to the same memory that
-- the original copy put it in.
copyManifest :: (SimplifyMemory rep inner) => TopDownRuleBasicOp (Wise rep)
copyManifest :: forall rep (inner :: * -> *).
SimplifyMemory rep inner =>
TopDownRuleBasicOp (Wise rep)
copyManifest TopDown (Wise rep)
vtable Pat (LetDec (Wise rep))
pat StmAux (ExpDec (Wise rep))
aux (Replicate (Shape []) (Var VName
v2))
  | Just (Reshape VName
v1 NewShape SubExp
s, Certs
v2_cs) <- VName -> TopDown (Wise rep) -> Maybe (BasicOp, Certs)
forall rep. VName -> SymbolTable rep -> Maybe (BasicOp, Certs)
ST.lookupBasicOp VName
v2 TopDown (Wise rep)
vtable,
    Just (Manifest VName
v0 [Int]
perm, Certs
v1_cs) <- VName -> TopDown (Wise rep) -> Maybe (BasicOp, Certs)
forall rep. VName -> SymbolTable rep -> Maybe (BasicOp, Certs)
ST.lookupBasicOp VName
v1 TopDown (Wise rep)
vtable,
    Pat [PatElem VName
_ (VarWisdom
_, MemArray PrimType
_ ShapeBase SubExp
_ NoUniqueness
_ (ArrayIn VName
mem LMAD
_))] <- Pat (LetDec (Wise rep))
pat =
      RuleM (Wise rep) () -> Rule (Wise rep)
forall rep. RuleM rep () -> Rule rep
Simplify (RuleM (Wise rep) () -> Rule (Wise rep))
-> RuleM (Wise rep) () -> Rule (Wise rep)
forall a b. (a -> b) -> a -> b
$ do
        ~(MemArray pt shape u (ArrayIn _ v1_lmad)) <- VName -> RuleM (Wise rep) LetDecMem
forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Mem rep inner) =>
VName -> m LetDecMem
lookupMemInfo VName
v1
        v0' <- newVName (baseString v1 <> "_manifest")
        let manifest_pat =
              [PatElem LetDecMem] -> Pat LetDecMem
forall dec. [PatElem dec] -> Pat dec
Pat [VName -> LetDecMem -> PatElem LetDecMem
forall dec. VName -> dec -> PatElem dec
PatElem VName
v0' (LetDecMem -> PatElem LetDecMem) -> LetDecMem -> PatElem LetDecMem
forall a b. (a -> b) -> a -> b
$ PrimType
-> ShapeBase SubExp -> NoUniqueness -> MemBind -> LetDecMem
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
pt ShapeBase SubExp
shape NoUniqueness
u (MemBind -> LetDecMem) -> MemBind -> LetDecMem
forall a b. (a -> b) -> a -> b
$ VName -> LMAD -> MemBind
ArrayIn VName
mem LMAD
v1_lmad]
            stm = Pat (LetDec rep)
-> StmAux (ExpDec rep) -> Exp (Wise rep) -> Stm (Wise rep)
forall rep.
Informing rep =>
Pat (LetDec rep)
-> StmAux (ExpDec rep) -> Exp (Wise rep) -> Stm (Wise rep)
mkWiseStm Pat (LetDec rep)
Pat LetDecMem
manifest_pat StmAux (ExpDec rep)
forall a. Monoid a => a
mempty (Exp (Wise rep) -> Stm (Wise rep))
-> Exp (Wise rep) -> Stm (Wise rep)
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Wise rep)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Wise rep)) -> BasicOp -> Exp (Wise rep)
forall a b. (a -> b) -> a -> b
$ VName -> [Int] -> BasicOp
Manifest VName
v0 [Int]
perm
        certifying (v1_cs <> v2_cs) $ addStm stm
        auxing aux $ letBind pat $ BasicOp $ Reshape v0' s
copyManifest TopDown (Wise rep)
_ Pat (LetDec (Wise rep))
_ StmAux (ExpDec (Wise rep))
_ BasicOp
_ = Rule (Wise rep)
forall rep. Rule rep
Skip

-- | Standard collection of simplification rules for representations
-- with memory.
memRuleBook :: (SimplifyMemory rep inner) => RuleBook (Wise rep)
memRuleBook :: forall rep (inner :: * -> *).
SimplifyMemory rep inner =>
RuleBook (Wise rep)
memRuleBook =
  RuleBook (Wise rep)
forall rep. (BuilderOps rep, TraverseOpStms rep) => RuleBook rep
standardRules
    RuleBook (Wise rep) -> RuleBook (Wise rep) -> RuleBook (Wise rep)
forall a. Semigroup a => a -> a -> a
<> [TopDownRule (Wise rep)]
-> [BottomUpRule (Wise rep)] -> RuleBook (Wise rep)
forall m. [TopDownRule m] -> [BottomUpRule m] -> RuleBook m
ruleBook
      [ RuleOp (Wise rep) (TopDown (Wise rep)) -> TopDownRule (Wise rep)
forall rep a. RuleOp rep a -> SimplificationRule rep a
RuleOp RuleOp (Wise rep) (TopDown (Wise rep))
forall rep (inner :: * -> *).
SimplifyMemory rep inner =>
TopDownRuleOp (Wise rep)
decertifySafeAlloc,
        RuleBasicOp (Wise rep) (TopDown (Wise rep))
-> TopDownRule (Wise rep)
forall rep a. RuleBasicOp rep a -> SimplificationRule rep a
RuleBasicOp RuleBasicOp (Wise rep) (TopDown (Wise rep))
forall rep (inner :: * -> *).
SimplifyMemory rep inner =>
TopDownRuleBasicOp (Wise rep)
copyManifest
      ]
      []