horde-ad-0.3.0.0: src/HordeAd/Core/AstSimplify.hs
{-# LANGUAGE CPP #-}
#if MIN_VERSION_GLASGOW_HASKELL(9,12,1,0)
{-# OPTIONS_GHC -fno-expose-overloaded-unfoldings #-}
#endif
{-# LANGUAGE UnboxedTuples #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
{-# OPTIONS_GHC -Wno-orphans #-}
{-# OPTIONS_GHC -fmax-pmcheck-models=50 #-}
-- {-# OPTIONS_GHC -freduction-depth=10000 #-}
-- {-# OPTIONS_GHC -fconstraint-solver-iterations=10000 #-}
-- | This module holds smart constructors for AST, that is,
-- term-simplifying combinators corresponding to the Ast constructors.
-- The combinators simplify only on the basis of inspecting the roots of their
-- argument term trees. If the arguments get modified,
-- the modified forms are again inspected and potentially simplified again.
--
-- The limited simplification via combinators is enough to uncover redexes
-- for the factorization rules to fire and to undo some of the complication
-- introduced by vectorization. The intention is to leave intact as much
-- as possible of the original terms provided by the user while making
-- sure subterms introduced by vectorization are maximally simplified.
module HordeAd.Core.AstSimplify
( RewritePhase(..), SimplifyKnobs (..), defaultKnobs
, -- * The simplifying combinators, one for almost each AST constructor
astPair, astProject1, astProject2, astMapAccumLDer, astApply, astLet
, astPrimalPart, astDualPart, astPlainPart
, astPlusK, astTimesK, astN1K, astR1K, astR2K, astI2K, astConcreteK
, astFloorK, astFromIntegralK, astCastK
, astArgMinK, astArgMaxK, astIndexK, astIndexKnobsK
, astPlusS, astTimesS, astN1S, astR1S, astR2S, astI2S, astConcreteS
, astFloorS, astFromIntegralS, astCastS, astIndexS, astIndexKnobsS
, astCondKInitial, astCondK, astCondSInitial, astCondS
, astFromVectorK, astFromVectorS
, astSumK, astSumS, astScatterS, astScatterKnobsS
, astReplicateK, astReplicateS, astGatherS, astGatherKnobsS
, astAppendS, astSliceS, astReverseS, astTransposeS, astReshapeS
, astConvert
, astConvDownSFromR, astConvDownSFromX, astConvUpRFromS, astConvUpXFromS
, astSumKContract, astDot0, astDot1InS, astMatmul2S
, astBoolNotK, astBoolNotS, astBoolAndK, astBoolAndS, astLeqK, astLeq, astLeqS
-- * Helper combinators
, astConcrete, astLetFun, astLetRefresh, unRepl1
-- * Substitution operations
, substituteAst, substituteAstIxS
) where
import Prelude
import Control.Exception.Assert.Sugar
import Control.Monad (mapAndUnzipM, mplus)
import Data.Default
import Data.Foldable qualified as Foldable
import Data.List (elemIndex, findIndex)
import Data.Maybe (catMaybes, fromMaybe, isJust)
import Data.Proxy (Proxy (Proxy))
import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))
import Data.Type.Ord (Compare)
import Data.Vector.Generic qualified as V
import Data.Vector.Strict qualified as Data.Vector
import GHC.TypeLits
( Nat
, OrderingI (..)
, cmpNat
, sameNat
, type (*)
, type (+)
, type (-)
, type (<=)
, type (<=?)
)
import System.IO.Unsafe (unsafePerformIO)
import Type.Reflection (typeRep)
import Unsafe.Coerce (unsafeCoerce)
import Data.Array.Nested (MapJust, Replicate, type (++))
import Data.Array.Nested qualified as Nested
import Data.Array.Nested.Convert
(shrFromShS, shxFromShS, withShsFromShR, withShsFromShX)
import Data.Array.Nested.Lemmas
import Data.Array.Nested.Mixed.Shape
import Data.Array.Nested.Permutation (DropLen, Perm (..), TakeLen, permInverse)
import Data.Array.Nested.Permutation qualified as Permutation
import Data.Array.Nested.Ranked.Shape
import Data.Array.Nested.Shaped qualified as Shaped
import Data.Array.Nested.Shaped.Shape
import Data.Array.Nested.Types
(Head, Init, Last, Tail, fromSNat', pattern SZ, snatMul, snatPlus, unsafeCoerceRefl)
import HordeAd.Core.Ast
( AstTensor (AstConcreteK, AstConcreteS, AstPlusK, AstPlusS, AstTimesK, AstTimesS)
)
import HordeAd.Core.Ast hiding (AstTensor (..))
import HordeAd.Core.Ast qualified as Ast (AstTensor(..))
import HordeAd.Core.AstFreshId
import HordeAd.Core.AstTools
import HordeAd.Core.CarriersAst
(eqY, eqUnknownShapes, sunReplicate1, sunReplicateN, sunReplicate, unReplC, unAstK, unAstS)
import HordeAd.Core.CarriersConcrete
import HordeAd.Core.Conversion
import HordeAd.Core.ConvertTensor
import HordeAd.Core.Ops
import HordeAd.Core.TensorKind
import HordeAd.Core.Types
import HordeAd.Core.Unwind
data RewritePhase =
PhaseUnspecified
| PhaseVectorization
| PhaseSimplification
| PhaseExpansion
| PhaseContraction
deriving Eq
newtype SimplifyKnobs = SimplifyKnobs
{ knobPhase :: RewritePhase
}
defaultKnobs :: SimplifyKnobs
defaultKnobs = SimplifyKnobs PhaseUnspecified
-- @PhaseVectorization@ should only affect the topmost redex.
deVect :: SimplifyKnobs -> SimplifyKnobs
deVect (SimplifyKnobs PhaseVectorization) = SimplifyKnobs PhaseUnspecified
deVect knobs = knobs
-- * Expressing operations as Gather; introduces new variable names
-- | We keep AstTranspose terms for as long as possible, because
-- they are small and fuse nicely in many cases. For some forms of indexing
-- and nesting with reshape and gather they don't fuse, which is when
-- this function is invoked.
astTransposeAsGatherS
:: forall perm sh s r. KnownSpan s
=> SimplifyKnobs -> Permutation.Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (Permutation.PermutePrefix perm sh) r)
{-# NOINLINE astTransposeAsGatherS #-}
astTransposeAsGatherS knobs perm v =
let FTKS sh _ = ftkAst v
shn = shsDropLenPerm perm sh
shp = shsTakeLenPerm perm sh
shpPermuted = shsPermute perm shp
in funToVarsIxS @_ @AstMethodLet shpPermuted $ \vars ix ->
-- See astGatherCase.AstTransposeS for similar code with more comments.
gcastWith (lemRankMapJust shp) $
gcastWith (unsafeCoerceRefl :: Rank (TakeLen perm sh) :~: Rank perm) $
permInverse perm $ \(invperm :: Nested.Perm invperm) proof ->
case proof (ssxFromShX $ shxFromShS shp) of
Refl ->
gcastWith (unsafeCoerceRefl
:: DropLen invperm
(Permutation.Permute perm (TakeLen perm sh))
:~: '[]) $
gcastWith (lemAppNil
@(Permutation.Permute invperm
(Permutation.Permute perm (TakeLen perm sh)))) $
-- Seriously? This should be deduced from the above:
gcastWith (unsafeCoerceRefl
:: Permutation.PermutePrefix invperm
(Permutation.Permute perm (TakeLen perm sh))
:~: Permutation.Permute invperm
(Permutation.Permute perm (TakeLen perm sh))) $
-- This should follow from @proof@, if not for MapJust:
gcastWith (unsafeCoerceRefl
:: Permutation.Permute invperm
(Permutation.Permute perm (TakeLen perm sh))
:~: TakeLen perm sh) $
let asts :: AstIxS AstMethodLet (TakeLen perm sh)
!asts = ixsPermutePrefix invperm ix
in gcastWith (unsafeCoerceRefl
:: TakeLen perm sh ++ DropLen perm sh :~: sh) $
astGatherKnobsS knobs shpPermuted shn shp v (vars, asts)
-- TODO: To make this less disastrous, we need to add an extra constructor
-- to AstIndex with the semantics "this index reshaped from shIn to shOut"
-- that fuses perfectly with itself and absorbs normal indexes
-- by substitution. Or perhaps make this the only constructor, with normal
-- indexes represented as "this index reshaped from sh to sh".
-- Or only extend AstGather and possibly also AstIndex with the extra
-- shIn and shOut arguments. This complicates any code related to
-- AstGather and AstIndex, but often prevents nested reshapes from affecting
-- term size in any way. But we'd need to be careful to avoid breaking such
-- an index into components, because that forces index normalization,
-- e.g., index(gather) can no longer simplify recursively by one index
-- component at a time (probably possible only if the index is shorter
-- that the list of variables fo the gather). There are probably bad cases
-- where term size blowup can't be avoided, because the index has to be
-- normalized between each reshape.
--
-- | This generates big terms that don't simplify well,
-- so we keep the AstReshape form until simplification gets stuck.
-- In fact, to simplify the terms we'd need advanced solving of equations
-- in integer arithmetic modulo.
astReshapeAsGatherS
:: forall sh sh2 x s. KnownSpan s
=> SimplifyKnobs -> ShS sh2 -> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
{-# NOINLINE astReshapeAsGatherS #-}
astReshapeAsGatherS knobs shOut v | Refl <- lemAppNil @sh2
, Refl <- lemAppNil @sh
, FTKS shIn _ <- ftkAst v =
funToVarsIxS shOut $ \vars ix ->
let iUnshared :: AstInt AstMethodLet
iUnshared = ixsToLinear shOut ix
-- This can't be done, because i depends on vars:
-- astLetFun iUnshared $ \i ->
i = iUnshared -- sharing broken
asts :: AstIxS AstMethodLet sh
!asts = fromLinearIdxS shIn i
in astGatherKnobsS knobs shOut ZSS shIn v (vars, asts)
-- I can't switch to ixxFromLinear from ox-arrays
-- even just because IntegralH is not available in ox-arrays.
--
-- The inlines are not needed due to function arguments,
-- because there are none, but due to observed significant speedup
-- they provide in tests that intensively use these operations.
-- Maybe specialization doesn't quite work for them? Not verified.
--
-- | Given a linear index into the buffer, get the corresponding
-- multidimensional index.
--
-- If any of the dimensions is 0, the linear index has to be 0
-- (which we can't assert, because j may be a term and so == lies),
-- which is fine, because that's pointing at the start of the empty buffer.
fromLinearIdxS :: forall sh j. IntegralH j
=> ShS sh -> j -> IxS sh j
fromLinearIdxS = \sh lin -> case go sh lin of (# _, ix #) -> ix
where
go :: ShS sh1 -> j -> (# j, IxS sh1 j #)
go ZSS !n = (# n, ZIS #)
go ((:$$) n sh) lin =
let (# tensLin, idxInTens #) = go sh lin
tensLin' = tensLin `quotH` fromIntegral (fromSNat' n)
i = tensLin `remH` fromIntegral (fromSNat' n)
in (# tensLin', i :.$ idxInTens #)
-- * The simplifying combinators, one for almost each AST constructor
astPair :: KnownSpan s
=> AstTensor AstMethodLet s x -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKProduct x y)
astPair (Ast.AstProject1 t1) (Ast.AstProject2 t2)
| Just Refl <- eqUnknownShapes t1 t2 = t1
astPair (Ast.AstFromPrimal v1) (Ast.AstFromPrimal v2) =
fromPrimal $ astPair v1 v2
astPair (Ast.AstFromDual v1) (Ast.AstFromDual v2) =
fromDual $ astPair v1 v2
astPair (Ast.AstFromPlain v1) (Ast.AstFromPlain v2) =
fromPlain $ astPair v1 v2
astPair (AstConvUp c1 ftk1 v1) (AstConvUp c2 ftk2 v2) =
astConvertUp (ConvT2 c1 c2) (FTKProduct ftk1 ftk2) (astPair v1 v2)
astPair (AstConvUp c1 ftk1 v1) v2 =
astConvertUp (ConvT2 c1 ConvId) (FTKProduct ftk1 (ftkAst v2)) (astPair v1 v2)
astPair v1 (AstConvUp c2 ftk2 v2) =
astConvertUp (ConvT2 ConvId c2) (FTKProduct (ftkAst v1) ftk2) (astPair v1 v2)
astPair v1 v2 = Ast.AstPair v1 v2
astProject1
:: forall x z s. KnownSpan s
=> AstTensor AstMethodLet s (TKProduct x z) -> AstTensor AstMethodLet s x
astProject1 u = case u of
Ast.AstPair x _z -> x
Ast.AstMapAccumLDer k bftk eftk
(AstLambda varf vf)
(AstLambda vard vd)
(AstLambda varr vr)
acc0 es | Nothing <- matchingFTK bftk ftkUnit ->
let varf2 = varf
vf2 = astPair (astProject1 vf) (fromPlain $ AstConcreteK Z1)
vard2 = vard
vd2 = astPair (astProject1 vd) (fromPlain $ AstConcreteK Z1)
accftk = ftkAst acc0
ftkr2 = FTKProduct
(adFTK $ FTKProduct accftk ftkUnit)
(FTKProduct accftk eftk)
varr2 = reshapeVarName ftkr2 varr
astr2 = astVar varr2
zero = fromPlain $ astConcrete (adFTK bftk) (tdefTarget (adFTK bftk))
vr2 = substituteAst
(astPair (astPair (astProject1 (astProject1 astr2))
zero)
(astProject2 astr2))
varr vr
in Ast.AstProject1
$ astMapAccumLDer k ftkUnit eftk (AstLambda varf2 vf2)
(AstLambda vard2 vd2)
(AstLambda varr2 vr2) acc0 es
Ast.AstLet var t v -> astLet var t (astProject1 v)
Ast.AstPrimalPart v -> astPrimalPart $ astProject1 v
Ast.AstDualPart v -> astDualPart $ astProject1 v
Ast.AstPlainPart v -> astPlainPart $ astProject1 v
Ast.AstFromPrimal u1 -> fromPrimal $ astProject1 u1
Ast.AstFromDual u1 -> fromDual $ astProject1 u1
Ast.AstFromPlain u1 -> fromPlain $ astProject1 u1
Ast.AstConvert (ConvT2 c1 _c2) t -> astConvert c1 $ astProject1 t
-- Even Down conversions are not pushed down projections, so AstConvUp
-- is not necessary here for consistency of rewriting. The AstConvUp case
-- commented out below probably doesn't add anything, because ConvCmp is not
-- a normal form of conversions of products (and ConvZip is ruled out).
-- AstConvUp (FTKProduct zftk _) t | FTKProduct yftk _ <- ftkAst t ->
-- astConvert (convUp yftk zftk) (astProject1 t)
_ -> Ast.AstProject1 u
astProject2
:: forall x z s. KnownSpan s
=> AstTensor AstMethodLet s (TKProduct x z) -> AstTensor AstMethodLet s z
astProject2 u = case u of
Ast.AstPair _x z -> z
Ast.AstLet var t v -> astLet var t (astProject2 v)
Ast.AstPrimalPart v -> astPrimalPart $ astProject2 v
Ast.AstDualPart v -> astDualPart $ astProject2 v
Ast.AstPlainPart v -> astPlainPart $ astProject2 v
Ast.AstFromPrimal u1 -> fromPrimal $ astProject2 u1
Ast.AstFromDual u1 -> fromDual $ astProject2 u1
Ast.AstFromPlain u1 -> fromPlain $ astProject2 u1
Ast.AstConvert (ConvT2 _c1 c2) t -> astConvert c2 $ astProject2 t
_ -> Ast.AstProject2 u
astMapAccumLDer
:: forall accy by ey k s. KnownSpan s
=> SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s
(TKProduct accy ey) (TKProduct accy by)
-> AstHFun s
(TKProduct (ADTensorKind (TKProduct accy ey))
(TKProduct accy ey))
(ADTensorKind (TKProduct accy by))
-> AstHFun s
(TKProduct (ADTensorKind (TKProduct accy by))
(TKProduct accy ey))
(ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
astMapAccumLDer k bftk eftk (AstLambda varf vf)
(AstLambda vard vd)
(AstLambda varr vr)
(AstConvUp @accyFrom c accftk acc0From) es =
let accftkFrom = ftkAst acc0From
ftkf2 = FTKProduct accftkFrom eftk
varf2 = reshapeVarName ftkf2 varf
astf2 = astVar varf2
vf2 =
let subbed =
substituteAst
(astPair (astConvertUp c accftk (astProject1 astf2))
(astProject2 astf2))
varf vf
in astConvDown @(TKProduct accy by)
(FTKProduct accftkFrom bftk)
subbed
ftkd2 = FTKProduct
(adFTK $ FTKProduct accftkFrom eftk)
(FTKProduct accftkFrom eftk)
vard2 = reshapeVarName ftkd2 vard
astd2 = astVar vard2
vd2 =
let subbed =
substituteAst
(astPair (astPair (astConvUp @(ADTensorKind accyFrom)
(adFTK accftk)
(astProject1 (astProject1 astd2)))
-- TODO: adTKConversion c?
(astProject2 (astProject1 astd2)))
(astPair (astConvertUp c accftk
(astProject1 (astProject2 astd2)))
(astProject2 (astProject2 astd2))))
vard vd
in astConvDown @(ADTensorKind (TKProduct accy by))
(adFTK $ FTKProduct accftkFrom bftk)
subbed
ftkr2 = FTKProduct
(adFTK $ FTKProduct accftkFrom bftk)
(FTKProduct accftkFrom eftk)
varr2 = reshapeVarName ftkr2 varr
astr2 = astVar varr2
vr2 =
let subbed =
substituteAst
(astPair (astPair (astConvUp @(ADTensorKind accyFrom)
(adFTK accftk)
(astProject1 (astProject1 astr2)))
(astProject2 (astProject1 astr2)))
(astPair (astConvertUp c accftk
(astProject1 (astProject2 astr2)))
(astProject2 (astProject2 astr2))))
varr vr
in astConvDown @(ADTensorKind (TKProduct accy ey))
(adFTK $ FTKProduct accftkFrom eftk)
subbed
in astConvertUp (ConvT2 c ConvId)
(FTKProduct accftk (buildFTK k bftk))
$ astMapAccumLDer k bftk eftk (AstLambda varf2 vf2)
(AstLambda vard2 vd2)
(AstLambda varr2 vr2)
acc0From es
astMapAccumLDer k bftk eftk (AstLambda varf vf)
(AstLambda vard vd)
(AstLambda varr vr)
acc0 (AstConvUp @esShsFrom _c _esShsFTK esFrom) =
let accftk = ftkAst acc0
esShsFrom = ftkAst esFrom
esShsFromSTK = ftkToSTK esShsFrom
in case razeSTK esShsFromSTK of
(eftkFromSTK :: SingletonTK eyFrom) ->
gcastWith (unsafeCoerceRefl :: BuildTensorKind k eyFrom :~: esShsFrom) $
let eftkFrom = razeFTK k eftkFromSTK esShsFrom
ftkf2 = FTKProduct accftk eftkFrom
varf2 = reshapeVarName ftkf2 varf
astf2 = astVar varf2
vf2 =
let subbed =
substituteAst
(astPair (astProject1 astf2)
(astConvUp @eyFrom eftk (astProject2 astf2)))
varf vf
in subbed
ftkd2 = FTKProduct
(adFTK $ FTKProduct accftk eftkFrom)
(FTKProduct accftk eftkFrom)
vard2 = reshapeVarName ftkd2 vard
astd2 = astVar vard2
vd2 =
let subbed =
substituteAst
(astPair (astPair (astProject1 (astProject1 astd2))
(astConvUp @(ADTensorKind eyFrom)
(adFTK eftk)
(astProject2 (astProject1 astd2))))
(astPair (astProject1 (astProject2 astd2))
(astConvUp @eyFrom eftk
(astProject2 (astProject2 astd2)))))
vard vd
in subbed
ftkr2 = FTKProduct
(adFTK $ FTKProduct accftk bftk)
(FTKProduct accftk eftkFrom)
varr2 = reshapeVarName ftkr2 varr
astr2 = astVar varr2
vr2 =
let subbed =
substituteAst
(astPair (astProject1 astr2)
(astPair (astProject1 (astProject2 astr2))
(astConvUp @eyFrom eftk
(astProject2 (astProject2 astr2)))))
varr vr
in astConvDown @(ADTensorKind (TKProduct accy ey))
(adFTK $ FTKProduct accftk eftkFrom)
subbed
in astMapAccumLDer k bftk eftkFrom (AstLambda varf2 vf2)
(AstLambda vard2 vd2)
(AstLambda varr2 vr2)
acc0 esFrom
astMapAccumLDer k bftk eftk f df rf acc0 es =
Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es
astApply :: forall x z s. KnownSpan s
=> AstHFun s x z -> AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s z
astApply (AstLambda !var !v) u = astLet var u v
-- Invariant: if the variable has bounds, the expression can only have
-- values within the bounds (regardless of what the `bounds` call would say).
astLet :: forall y z s s2. KnownSpan s2
=> AstVarName '(s, y) -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
astLet _var _u v@AstConcreteK{} = v
astLet _var _u v@AstConcreteS{} = v
astLet _var _u v@Ast.AstIotaS{} = v
astLet var u v@(Ast.AstVar var2) =
if varNameToAstVarId var2 == varNameToAstVarId var
then case testEquality var var2 of
Just Refl -> u
_ -> error "astLet: wrong variable types at AstVar"
else v
astLet var u v@(Ast.AstPrimalPart (Ast.AstVar var2)) = -- a common noop
if varNameToAstVarId var2 == varNameToAstVarId var
then case testEquality var var2 of
Just Refl -> astPrimalPart u
_ -> error "astLet: wrong variable types at AstPrimalPart"
else v
astLet var u v@(Ast.AstDualPart (Ast.AstVar var2)) = -- a noop
if varNameToAstVarId var2 == varNameToAstVarId var
then case testEquality var var2 of
Just Refl -> astDualPart u
_ -> error "astLet: wrong variable types at AstDualPart"
else v
astLet var u v@(Ast.AstPlainPart (Ast.AstVar var2)) = -- a noop
if varNameToAstVarId var2 == varNameToAstVarId var
then case testEquality var var2 of
Just Refl -> astPlainPart u
_ -> error "astLet: wrong variable types at AstPlainPart"
else v
astLet var u v | astIsSmall True u =
substituteAst u var v
astLet var u (Ast.AstFromPrimal v0) = fromPrimal $ astLet var u v0
astLet var u (Ast.AstFromDual v0) = fromDual $ astLet var u v0
astLet var u (Ast.AstFromPlain v0) = fromPlain $ astLet var u v0
astLet var (Ast.AstFromPrimal u) v =
withKnownSpan (varNameToSpan var)
$ astLetFun u $ \ !ast1 -> substituteAst (fromPrimal ast1) var v
astLet var (Ast.AstFromDual u) v =
astLetFun u $ \ !ast1 -> substituteAst (fromDual ast1) var v
astLet var (Ast.AstFromPlain u) v =
withKnownSpan (varNameToSpan var)
$ astLetFun u $ \ !ast1 -> substituteAst (fromPlain ast1) var v
astLet var (Ast.AstPair u1 u2) v =
withKnownSpan (varNameToSpan var) $
astLetFun u1 $ \ !ast1 -> astLetFun u2 $ \ !ast2 ->
substituteAst (Ast.AstPair ast1 ast2) var v
astLet var (Ast.AstLet varN uN u12@Ast.AstPair{}) v =
withKnownSpan (varNameToSpan var) $
astLetRefresh varN uN u12 $ \u3 -> case u3 of
Ast.AstPair u1' u2' ->
astLetFun u1' $ \ !ast1 -> astLetFun u2' $ \ !ast2 ->
substituteAst (Ast.AstPair ast1 ast2) var v
_ -> astLet var (astLet varN uN u3) v
-- This is a common case, e.g., from representing conditionals.
astLet var (Ast.AstFromVectorK shm u) v | V.length u == 2 =
withKnownSpan (varNameToSpan var) $
astLetFun (u V.! 0) $ \ !ast1 -> astLetFun (u V.! 1) $ \ !ast2 ->
substituteAst (Ast.AstFromVectorK shm
$ V.fromListN 2 [ast1, ast2]) var v
astLet var (Ast.AstLet varN uN
u12@(Ast.AstFromVectorK _ u)) v | V.length u == 2 =
withKnownSpan (varNameToSpan var) $
astLetRefresh varN uN u12 $ \u3 -> case u3 of
Ast.AstFromVectorK shm u' ->
astLetFun (u' V.! 0) $ \ !ast1 -> astLetFun (u' V.! 1) $ \ !ast2 ->
substituteAst (Ast.AstFromVectorK shm
$ V.fromListN 2 [ast1, ast2]) var v
_ -> astLet var (astLet varN uN u3) v
astLet var (Ast.AstFromVectorS shm u) v | V.length u == 2 =
withKnownSpan (varNameToSpan var) $
astLetFun (u V.! 0) $ \ !ast1 -> astLetFun (u V.! 1) $ \ !ast2 ->
substituteAst (Ast.AstFromVectorS shm
$ V.fromListN 2 [ast1, ast2]) var v
astLet var (Ast.AstLet varN uN
u12@(Ast.AstFromVectorS _ u)) v | V.length u == 2 =
withKnownSpan (varNameToSpan var) $
astLetRefresh varN uN u12 $ \u3 -> case u3 of
Ast.AstFromVectorS shm u' ->
astLetFun (u' V.! 0) $ \ !ast1 -> astLetFun (u' V.! 1) $ \ !ast2 ->
substituteAst (Ast.AstFromVectorS shm
$ V.fromListN 2 [ast1, ast2]) var v
_ -> astLet var (astLet varN uN u3) v
astLet var (Ast.AstReplicateK shm a) v =
let var2 = reshapeVarName (ftkAst a) var
ast = Ast.AstReplicateK shm $ astVar var2
in astLet var2 a (substituteAst ast var v)
astLet var (Ast.AstReplicateS shm a) v =
let var2 = reshapeVarName (ftkAst a) var
ast = Ast.AstReplicateS shm $ astVar var2
in astLet var2 a (substituteAst ast var v)
astLet var (Ast.AstTransposeS perm a) v =
let var2 = reshapeVarName (ftkAst a) var
ast = Ast.AstTransposeS perm $ astVar var2
in astLet var2 a (substituteAst ast var v)
astLet var (AstConvUp c zftk a) v =
let var2 = reshapeVarName (ftkAst a) var
ast = withKnownSpan (varNameToSpan var)
$ astConvertUp c zftk $ astVar var2
in astLet var2 a (substituteAst ast var v)
astLet var u (AstConvUp c zftk a) =
astConvertUp c zftk $ astLet var u a
astLet var u v = Ast.AstLet var u v
astPrimalPart :: KnownSpan s
=> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet (PrimalStepSpan s) y
astPrimalPart t = case t of
Ast.AstPair t1 t2 -> astPair (astPrimalPart t1) (astPrimalPart t2)
-- We really need equality saturation...
Ast.AstProject1 u@Ast.AstMapAccumLDer{} -> astProject1 $ astPrimalPart u
Ast.AstProject2 u@Ast.AstMapAccumLDer{} -> astProject2 $ astPrimalPart u
Ast.AstProject1 u@(Ast.AstProject1 Ast.AstMapAccumLDer{}) ->
astProject1 $ astPrimalPart u
Ast.AstProject1 u@(Ast.AstProject2 Ast.AstMapAccumLDer{}) ->
astProject1 $ astPrimalPart u
Ast.AstProject2 u@(Ast.AstProject2 Ast.AstMapAccumLDer{}) ->
astProject2 $ astPrimalPart u
Ast.AstProject2 u@(Ast.AstProject1 Ast.AstMapAccumLDer{}) ->
astProject2 $ astPrimalPart u
Ast.AstProject1{} -> Ast.AstPrimalPart t
Ast.AstProject2{} -> Ast.AstPrimalPart t
Ast.AstMapAccumLDer k bftk eftk (AstLambda varf vf)
(AstLambda vard vd)
(AstLambda varr vr) acc0 es ->
let varf2 = respanVarName varf
astf2 = fromPrimal $ astVar varf2
vf2 = astPrimalPart $ substituteAst astf2 varf vf
vard2 = respanVarName vard
astd2 = fromPrimal $ astVar vard2
vd2 = astPrimalPart $ substituteAst astd2 vard vd
varr2 = respanVarName varr
astr2 = fromPrimal $ astVar varr2
vr2 = astPrimalPart $ substituteAst astr2 varr vr
in astMapAccumLDer k bftk eftk (AstLambda varf2 vf2)
(AstLambda vard2 vd2)
(AstLambda varr2 vr2)
(astPrimalPart acc0) (astPrimalPart es)
Ast.AstApply{} -> Ast.AstPrimalPart t
Ast.AstVar{} -> Ast.AstPrimalPart t -- the only normal form
Ast.AstBuild1 k stk (var, v) ->
let !v2 = astPrimalPart v
in Ast.AstBuild1 k stk (var, v2)
Ast.AstLet var u v -> astLet var u (astPrimalPart v)
Ast.AstPrimalPart{} -> Ast.AstPrimalPart t
Ast.AstDualPart{} -> Ast.AstPrimalPart t
Ast.AstPlainPart{} -> Ast.AstPrimalPart t -- impossible, but fine
Ast.AstFromPrimal v -> v
Ast.AstFromDual v ->
let ftk = ftkAst v
in fromPlain $ astConcrete ftk (tdefTarget ftk)
Ast.AstFromPlain v -> fromPlain v
AstPlusK u v -> astPrimalPart u `astPlusK` astPrimalPart v
AstTimesK u v -> astPrimalPart u `astTimesK` astPrimalPart v
Ast.AstN1K opCode u -> astN1K opCode (astPrimalPart u)
Ast.AstR1K opCode u -> astR1K opCode (astPrimalPart u)
Ast.AstR2K opCode u v -> astR2K opCode (astPrimalPart u) (astPrimalPart v)
Ast.AstI2K opCode u v -> astI2K opCode (astPrimalPart u) (astPrimalPart v)
AstConcreteK{} -> Ast.AstPrimalPart t
Ast.AstFloorK{} -> Ast.AstPrimalPart t
Ast.AstFromIntegralK{} -> Ast.AstPrimalPart t
Ast.AstCastK v -> astCastK $ astPrimalPart v
Ast.AstArgMinK{} -> Ast.AstPrimalPart t
Ast.AstArgMaxK{} -> Ast.AstPrimalPart t
Ast.AstIndexK v ix -> astIndexK (astPrimalPart v) ix
AstPlusS u v -> astPrimalPart u `astPlusS` astPrimalPart v
AstTimesS u v -> astPrimalPart u `astTimesS` astPrimalPart v
Ast.AstN1S opCode u -> astN1S opCode (astPrimalPart u)
Ast.AstR1S opCode u -> astR1S opCode (astPrimalPart u)
Ast.AstR2S opCode u v -> astR2S opCode (astPrimalPart u) (astPrimalPart v)
Ast.AstI2S opCode u v -> astI2S opCode (astPrimalPart u) (astPrimalPart v)
AstConcreteS{} -> Ast.AstPrimalPart t
Ast.AstFloorS{} -> Ast.AstPrimalPart t
Ast.AstFromIntegralS{} -> Ast.AstPrimalPart t
Ast.AstCastS v -> astCastS $ astPrimalPart v
Ast.AstArgMinS{} -> Ast.AstPrimalPart t
Ast.AstArgMaxS{} -> Ast.AstPrimalPart t
Ast.AstIndexS shn v ix -> astIndexS shn (astPrimalPart v) ix
Ast.AstCondK b a2 a3 -> astCondK b (astPrimalPart a2) (astPrimalPart a3)
Ast.AstCondS b a2 a3 -> astCondS b (astPrimalPart a2) (astPrimalPart a3)
Ast.AstFromVectorK shm l -> astFromVectorK shm (V.map astPrimalPart l)
Ast.AstFromVectorS shm l -> astFromVectorS shm (V.map astPrimalPart l)
Ast.AstSumK v -> astSumK $ astPrimalPart v
Ast.AstSumS shm v -> astSumS shm $ astPrimalPart v
Ast.AstScatterS shm shn shp v (vars, ix) ->
astScatterS shm shn shp (astPrimalPart v) (vars, ix)
Ast.AstReplicateK shm v -> astReplicateK shm (astPrimalPart v)
Ast.AstReplicateS shm v -> astReplicateS shm (astPrimalPart v)
Ast.AstGatherS shm shn shp v (vars, ix) ->
astGatherS shm shn shp (astPrimalPart v) (vars, ix)
Ast.AstIotaS{} -> Ast.AstPrimalPart t
Ast.AstAppendS x y -> astAppendS (astPrimalPart x) (astPrimalPart y)
Ast.AstSliceS i n k v -> astSliceS i n k (astPrimalPart v)
Ast.AstReverseS v -> astReverseS (astPrimalPart v)
Ast.AstTransposeS perm v -> astTransposeS perm (astPrimalPart v)
Ast.AstReshapeS sh v -> astReshapeS sh (astPrimalPart v)
-- Most conversions need to stay down here to cancel out.
AstConvUp c zftk a -> astConvertUp c zftk $ astPrimalPart a
Ast.AstConvert{} -> Ast.AstPrimalPart t
Ast.AstDot0 u v -> astDot0 (astPrimalPart u) (astPrimalPart v)
Ast.AstDot1InS sh n u v -> astDot1InS sh n (astPrimalPart u) (astPrimalPart v)
Ast.AstMatmul2S m n p u v ->
astMatmul2S m n p (astPrimalPart u) (astPrimalPart v)
Ast.AstBoolNotK{} -> Ast.AstPrimalPart t
Ast.AstBoolNotS{} -> Ast.AstPrimalPart t
Ast.AstBoolAndK{} -> Ast.AstPrimalPart t
Ast.AstBoolAndS{} -> Ast.AstPrimalPart t
Ast.AstLeqK{} -> Ast.AstPrimalPart t
Ast.AstLeq{} -> Ast.AstPrimalPart t
Ast.AstLeqS{} -> Ast.AstPrimalPart t
-- Note how this can't be pushed down into, say, multiplication, because it
-- multiplies the dual part by the primal part. Addition is fine, though.
astDualPart :: AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart t = case t of
Ast.AstPair t1 t2 -> astPair (astDualPart t1) (astDualPart t2)
Ast.AstProject1{} -> Ast.AstDualPart t
Ast.AstProject2{} -> Ast.AstDualPart t
Ast.AstMapAccumLDer k bftk eftk (AstLambda varf vf)
(AstLambda vard vd)
(AstLambda varr vr) acc0 es ->
let varf2 = respanVarName varf
astf2 = fromDual $ astVar varf2
vf2 = astDualPart $ substituteAst astf2 varf vf
vard2 = respanVarName vard
astd2 = fromDual $ astVar vard2
vd2 = astDualPart $ substituteAst astd2 vard vd
varr2 = respanVarName varr
astr2 = fromDual $ astVar varr2
vr2 = astDualPart $ substituteAst astr2 varr vr
in astMapAccumLDer k bftk eftk (AstLambda varf2 vf2)
(AstLambda vard2 vd2)
(AstLambda varr2 vr2)
(astDualPart acc0) (astDualPart es)
Ast.AstApply{} -> Ast.AstDualPart t
Ast.AstVar{} -> Ast.AstDualPart t
Ast.AstBuild1 k stk (var, v) ->
let !v2 = astDualPart v
in Ast.AstBuild1 k stk (var, v2)
Ast.AstLet var u v -> astLet var u (astDualPart v)
Ast.AstFromPrimal v ->
let ftk = ftkAst v
in Ast.AstDualPart $ fromPlain
$ astConcrete ftk (tdefTarget ftk)
-- let's hope this is smaller than v
Ast.AstFromDual v -> v
Ast.AstFromPlain v -> -- TODO: are deltas zeroed enough here?
let ftk = ftkAst v
in Ast.AstDualPart $ fromPlain
$ astConcrete ftk (tdefTarget ftk)
-- let's hope this is smaller than v
AstPlusK u v -> astDualPart u `astPlusK` astDualPart v
-- This one is mathematically wrong, dual numbers don't mult like that:
-- AstTimesK u v -> astDualPart u * astDualPart v
Ast.AstN1K NegateOp u -> negate (astDualPart u)
{- Some of these are wrong, so let's be conservative:
Ast.AstN1K AbsOp u -> abs (astDualPart u)
Ast.AstN1K SignumOp u -> signum (astDualPart u)
Ast.AstR1K opCode u -> Ast.AstR1K opCode (astDualPart u)
Ast.AstR2K opCode u v -> Ast.AstR2K opCode (astDualPart u) (astDualPart v)
Ast.AstI2K QuotOp u v -> quotH (astDualPart u) (astDualPart v)
Ast.AstI2K RemOp u v -> remH (astDualPart u) (astDualPart v)
-}
Ast.AstCastK v -> astCastK $ astDualPart v
Ast.AstIndexK v ix -> astIndexK (astDualPart v) ix
AstPlusS u v -> astDualPart u `astPlusS` astDualPart v
-- This one is mathematically wrong, dual numbers don't mult like that:
-- AstTimesS u v -> astDualPart u * astDualPart v
Ast.AstN1S NegateOp u -> negate (astDualPart u)
{- Some of these are wrong, so let's be conservative:
Ast.AstN1S AbsOp u -> abs (astDualPart u)
Ast.AstN1S SignumOp u -> signum (astDualPart u)
Ast.AstR1S opCode u -> Ast.AstR1S opCode (astDualPart u)
Ast.AstR2S opCode u v -> Ast.AstR2S opCode (astDualPart u)
(astDualPart v)
Ast.AstI2S opCode u v -> Ast.AstI2S opCode (astDualPart u)
(astDualPart v)
-}
Ast.AstCastS v -> astCastS $ astDualPart v
Ast.AstIndexS shn v ix -> astIndexS shn (astDualPart v) ix
Ast.AstCondK b a2 a3 -> astCondK b (astDualPart a2) (astDualPart a3)
Ast.AstCondS b a2 a3 -> astCondS b (astDualPart a2) (astDualPart a3)
Ast.AstFromVectorK shm l -> astFromVectorK shm (V.map astDualPart l)
Ast.AstFromVectorS shm l -> astFromVectorS shm (V.map astDualPart l)
Ast.AstSumK v -> astSumK $ astDualPart v
Ast.AstSumS shm v -> astSumS shm $ astDualPart v
Ast.AstScatterS shm shn shp v (vars, ix) ->
astScatterS shm shn shp (astDualPart v) (vars, ix)
Ast.AstReplicateK shm v -> astReplicateK shm (astDualPart v)
Ast.AstReplicateS shm v -> astReplicateS shm (astDualPart v)
Ast.AstGatherS shm shn shp v (vars, ix) ->
astGatherS shm shn shp (astDualPart v) (vars, ix)
Ast.AstAppendS x y -> astAppendS (astDualPart x) (astDualPart y)
Ast.AstSliceS i n k v -> astSliceS i n k (astDualPart v)
Ast.AstReverseS v -> astReverseS (astDualPart v)
Ast.AstTransposeS perm v -> astTransposeS perm (astDualPart v)
Ast.AstReshapeS sh v -> astReshapeS sh (astDualPart v)
-- Most conversions need to stay down here to cancel out.
AstConvUp c zftk a -> astConvertUp c zftk $ astDualPart a
Ast.AstConvert{} -> Ast.AstDualPart t
Ast.AstDot0 u v -> astDot0 (astDualPart u) (astDualPart v)
Ast.AstDot1InS sh n u v -> astDot1InS sh n (astDualPart u) (astDualPart v)
Ast.AstMatmul2S m n p u v ->
astMatmul2S m n p (astDualPart u) (astDualPart v)
_ -> Ast.AstDualPart t
astPlainPart :: forall y s. KnownSpan s
=> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet PlainSpan y
astPlainPart t = case t of
_ | SPlainSpan <- knownSpan @s -> t
Ast.AstPair t1 t2 -> astPair (astPlainPart t1) (astPlainPart t2)
-- We really need equality saturation...
Ast.AstProject1 u@Ast.AstMapAccumLDer{} -> astProject1 $ astPlainPart u
Ast.AstProject2 u@Ast.AstMapAccumLDer{} -> astProject2 $ astPlainPart u
Ast.AstProject1 u@(Ast.AstProject1 Ast.AstMapAccumLDer{}) ->
astProject1 $ astPlainPart u
Ast.AstProject1 u@(Ast.AstProject2 Ast.AstMapAccumLDer{}) ->
astProject1 $ astPlainPart u
Ast.AstProject2 u@(Ast.AstProject2 Ast.AstMapAccumLDer{}) ->
astProject2 $ astPlainPart u
Ast.AstProject2 u@(Ast.AstProject1 Ast.AstMapAccumLDer{}) ->
astProject2 $ astPlainPart u
Ast.AstProject1{} -> Ast.AstPlainPart t
Ast.AstProject2{} -> Ast.AstPlainPart t
Ast.AstMapAccumLDer k bftk eftk (AstLambda varf vf)
(AstLambda vard vd)
(AstLambda varr vr) acc0 es ->
let varf2 = respanVarName varf
astf2 = fromPlain $ astVar varf2
vf2 = astPlainPart $ substituteAst astf2 varf vf
vard2 = respanVarName vard
astd2 = fromPlain $ astVar vard2
vd2 = astPlainPart $ substituteAst astd2 vard vd
varr2 = respanVarName varr
astr2 = fromPlain $ astVar varr2
vr2 = astPlainPart $ substituteAst astr2 varr vr
in astMapAccumLDer k bftk eftk (AstLambda varf2 vf2)
(AstLambda vard2 vd2)
(AstLambda varr2 vr2)
(astPlainPart acc0) (astPlainPart es)
Ast.AstApply{} -> Ast.AstPlainPart t
Ast.AstVar{} -> Ast.AstPlainPart t -- the only normal form
Ast.AstBuild1 k stk (var, v) ->
let !v2 = astPlainPart v
in Ast.AstBuild1 k stk (var, v2)
Ast.AstLet var u v -> astLet var u (astPlainPart v)
Ast.AstPrimalPart v -> astPlainPart v
Ast.AstDualPart{} -> Ast.AstPlainPart t
Ast.AstFromPrimal v -> astPlainPart v
Ast.AstFromDual v ->
let ftk = ftkAst v
in astConcrete ftk (tdefTarget ftk)
Ast.AstFromPlain v -> v
AstPlusK u v -> astPlainPart u `astPlusK` astPlainPart v
AstTimesK u v -> astPlainPart u `astTimesK` astPlainPart v
Ast.AstN1K opCode u -> astN1K opCode (astPlainPart u)
Ast.AstR1K opCode u -> astR1K opCode (astPlainPart u)
Ast.AstR2K opCode u v -> astR2K opCode (astPlainPart u) (astPlainPart v)
Ast.AstI2K opCode u v -> astI2K opCode (astPlainPart u) (astPlainPart v)
Ast.AstCastK v -> astCastK $ astPlainPart v
Ast.AstIndexK v ix -> astIndexK (astPlainPart v) ix
AstPlusS u v -> astPlainPart u `astPlusS` astPlainPart v
AstTimesS u v -> astPlainPart u `astTimesS` astPlainPart v
Ast.AstN1S opCode u -> astN1S opCode (astPlainPart u)
Ast.AstR1S opCode u -> astR1S opCode (astPlainPart u)
Ast.AstR2S opCode u v -> astR2S opCode (astPlainPart u) (astPlainPart v)
Ast.AstI2S opCode u v -> astI2S opCode (astPlainPart u) (astPlainPart v)
Ast.AstCastS v -> astCastS $ astPlainPart v
Ast.AstIndexS shn v ix -> astIndexS shn (astPlainPart v) ix
Ast.AstCondK b a2 a3 -> astCondK b (astPlainPart a2) (astPlainPart a3)
Ast.AstCondS b a2 a3 -> astCondS b (astPlainPart a2) (astPlainPart a3)
Ast.AstFromVectorK shm l -> astFromVectorK shm (V.map astPlainPart l)
Ast.AstFromVectorS shm l -> astFromVectorS shm (V.map astPlainPart l)
Ast.AstSumK v -> astSumK $ astPlainPart v
Ast.AstSumS shm v -> astSumS shm $ astPlainPart v
Ast.AstScatterS shm shn shp v (vars, ix) ->
astScatterS shm shn shp (astPlainPart v) (vars, ix)
Ast.AstReplicateK shm v -> astReplicateK shm (astPlainPart v)
Ast.AstReplicateS shm v -> astReplicateS shm (astPlainPart v)
Ast.AstGatherS shm shn shp v (vars, ix) ->
astGatherS shm shn shp (astPlainPart v) (vars, ix)
Ast.AstAppendS x y -> astAppendS (astPlainPart x) (astPlainPart y)
Ast.AstSliceS i n k v -> astSliceS i n k (astPlainPart v)
Ast.AstReverseS v -> astReverseS (astPlainPart v)
Ast.AstTransposeS perm v -> astTransposeS perm (astPlainPart v)
Ast.AstReshapeS sh v -> astReshapeS sh (astPlainPart v)
-- Most conversions need to stay down here to cancel out.
AstConvUp c zftk a -> astConvertUp c zftk $ astPlainPart a
Ast.AstConvert{} -> Ast.AstPlainPart t
Ast.AstDot0 u v -> astDot0 (astPlainPart u) (astPlainPart v)
Ast.AstDot1InS sh n u v -> astDot1InS sh n (astPlainPart u) (astPlainPart v)
Ast.AstMatmul2S m n p u v ->
astMatmul2S m n p (astPlainPart u) (astPlainPart v)
-- TODO: perhaps aim for a polynomial normal form? but that requires global
-- inspection of the whole expression
-- TODO: let's aim at SOP (Sum-of-Products) form, just as
-- ghc-typelits-natnormalise does. Also, let's associate to the right
-- and let's push negation down.
--
-- Not considered are rules that would require comparing non-constant terms
-- or that would duplicate a non-constant term, as well as most rules
-- informed by inequalities, expressed via max or min, such as
-- max n (signum (abs x)) | n <= 0 --> signum (abs x).
--
-- | Integer terms need to be simplified, because large ones are often created
-- due to vectorization and simplification, e.g., via astTransposeAsGather
-- or astReshapeAsGather and can be a deciding factor in whether
-- the other tensor terms can be simplified in turn.
--
-- Additions are flattened into a list-like form.
-- AstConcreteK or AstFromPlain (AstConcreteK), if any, is always the first
-- element of the list of summands.
astPlusK :: (NumScalar r, KnownSpan s)
=> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
{-# INLINEABLE astPlusK #-}
astPlusK = \cases
(Ast.AstLet var n u) v -> astLetRefresh var n u $ \u' -> astPlusK u' v
u (Ast.AstLet var n v) -> astLetRefresh var n v $ \v' -> astPlusK u v'
(Ast.AstPrimalPart u) (Ast.AstPrimalPart v) -> primalPart $ astPlusK u v
(Ast.AstDualPart u) (Ast.AstDualPart v) -> dualPart $ astPlusK u v
(Ast.AstPlainPart @_ @s1 u) (Ast.AstPlainPart @_ @s2 v)
| Just Refl <- testEquality (knownSpan @s1) (knownSpan @s2) ->
plainPart $ astPlusK u v
-- This is, sadly, unsound, because @fromPrimal . primalPart@ is not identity.
-- u (Ast.AstFromPrimal v) -> fromPrimal $ astPlusK (astPrimalPart u) v
(Ast.AstFromPrimal u) (Ast.AstFromPrimal v) -> fromPrimal $ astPlusK u v
(Ast.AstFromDual u) (Ast.AstFromDual v) -> fromDual $ astPlusK u v
(Ast.AstFromPlain u) (Ast.AstFromPlain v) -> fromPlain $ astPlusK u v
(AstPlusK u w) v -> astPlusK u (astPlusK w v) -- flattening
u v | Just 0 <- unAstK u -> v
u v | Just 0 <- unAstK v -> u
u v | Just u0 <- unAstK u
, Just v0 <- unAstK v -> fromPlain $ AstConcreteK (u0 + v0)
u v | Just v0 <- unAstK v -> AstPlusK (fromPlain $ AstConcreteK v0) u
u (AstPlusK v w) | Just v0 <- unAstK v ->
astPlusK (fromPlain $ AstConcreteK v0) u + w
(Ast.AstN1K NegateOp t1) t2 | eqY t1 t2 -> fromPlain $ AstConcreteK 0
(Ast.AstN1K NegateOp t1) (AstPlusK t2 u) | eqY t1 t2 -> u
t2 (Ast.AstN1K NegateOp t1) | eqY t1 t2 -> fromPlain $ AstConcreteK 0
t2 (AstPlusK (Ast.AstN1K NegateOp t1) u) | eqY t1 t2 -> u
(Ast.AstI2K RemOp (Ast.AstN1K NegateOp t1) n)
(Ast.AstI2K RemOp t2 n')
| eqY t1 t2 && eqY n n' -> fromPlain $ AstConcreteK 0
(Ast.AstI2K RemOp (Ast.AstN1K NegateOp t1) n)
(AstPlusK (Ast.AstI2K RemOp t2 n') u)
| eqY t1 t2 && eqY n n' -> u
(Ast.AstI2K RemOp t1 n')
(Ast.AstI2K RemOp (Ast.AstN1K NegateOp t2) n)
| eqY t1 t2 && eqY n n' -> fromPlain $ AstConcreteK 0
(Ast.AstI2K RemOp t1 n')
(AstPlusK (Ast.AstI2K RemOp (Ast.AstN1K NegateOp t2) n) u)
| eqY t1 t2 && eqY n n' -> u
t1 t2 | eqY t1 t2 -> fromPlain (AstConcreteK 2) `astTimesK` t1
t1 (AstPlusK t2 w)
| eqY t1 t2 -> fromPlain (AstConcreteK 2) `astTimesK` t1 + w
t1 (AstTimesK n t2)
| Just n0 <- unAstK n
, eqY t1 t2 -> fromPlain (AstConcreteK (n0 + 1)) `astTimesK` t1
t1 (AstPlusK (AstTimesK n t2) w)
| Just n0 <- unAstK n
, eqY t1 t2 -> fromPlain (AstConcreteK (n0 + 1)) `astTimesK` t1 + w
(AstTimesK n t2) t1
| Just n0 <- unAstK n
, eqY t1 t2 -> fromPlain (AstConcreteK (n0 + 1)) `astTimesK` t1
(AstTimesK n t2) (AstPlusK t1 w)
| Just n0 <- unAstK n
, eqY t1 t2 -> fromPlain (AstConcreteK (n0 + 1)) `astTimesK` t1 + w
(AstTimesK n1 t1) (AstTimesK n2 t2)
| Just n10 <- unAstK n1
, Just n20 <- unAstK n2
, eqY t1 t2 -> fromPlain (AstConcreteK (n10 + n20)) `astTimesK` t1
(AstTimesK n1 t1) (AstPlusK (AstTimesK n2 t2) w)
| Just n10 <- unAstK n1
, Just n20 <- unAstK n2
, eqY t1 t2 -> fromPlain (AstConcreteK (n10 + n20)) `astTimesK` t1 + w
u v -> AstPlusK u v
-- Just as with AstPlusK, factors are flattened and a constant comes first.
-- Also, negation comes first.
astTimesK :: (NumScalar r, KnownSpan s)
=> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
{-# INLINEABLE astTimesK #-}
astTimesK = \cases
(Ast.AstLet var n u) v -> astLetRefresh var n u $ \u' -> astTimesK u' v
u (Ast.AstLet var n v) -> astLetRefresh var n v $ \v' -> astTimesK u v'
(Ast.AstPrimalPart u) (Ast.AstPrimalPart v) -> primalPart $ astTimesK u v
(Ast.AstPlainPart @_ @s1 u) (Ast.AstPlainPart @_ @s2 v)
| Just Refl <- testEquality (knownSpan @s1) (knownSpan @s2) ->
plainPart $ astTimesK u v
(Ast.AstFromPrimal u) (Ast.AstFromPrimal v) -> fromPrimal $ astTimesK u v
Ast.AstFromDual{} Ast.AstFromDual{} -> fromPlain $ AstConcreteK 0
(Ast.AstFromPlain u) (Ast.AstFromPlain v) -> fromPlain $ astTimesK u v
(AstTimesK u w) v -> astTimesK u (astTimesK w v) -- flattening
u _ | Just 0 <- unAstK u -> u
_ v | Just 0 <- unAstK v -> v
u v | Just 1 <- unAstK u -> v
u v | Just 1 <- unAstK v -> u
u v | Just (-1) <- unAstK u -> negate v
u v | Just (-1) <- unAstK v -> negate u
u v | Just u0 <- unAstK u
, Just v0 <- unAstK v -> fromPlain $ AstConcreteK (u0 * v0)
u v | Just v0 <- unAstK v -> AstTimesK (fromPlain $ AstConcreteK v0) u
u (AstTimesK v w) | Just v0 <- unAstK v ->
astTimesK (fromPlain $ AstConcreteK v0) u * w
(Ast.AstN1K NegateOp u) (Ast.AstN1K NegateOp v) -> astTimesK u v
(Ast.AstN1K NegateOp u) (AstTimesK (Ast.AstN1K NegateOp v) w) ->
astTimesK u v * w
u (Ast.AstN1K NegateOp v) -> astTimesK (negate u) v
u (AstTimesK (Ast.AstN1K NegateOp v) w) -> astTimesK (negate u) v * w
(Ast.AstR2K DivideOp u1 u2) v -> astR2K DivideOp (u1 * v) u2
u (Ast.AstR2K DivideOp v1 v2) -> astR2K DivideOp (u * v1) v2
(Ast.AstR1K RecipOp u) v -> astR2K DivideOp v u
u (Ast.AstR1K RecipOp v) -> astR2K DivideOp u v
-- This breaks sharing, because although u is concrete and so doesn't
-- have to be shared, the multiplication is not shared --- we end up
-- with one addition and two multiplications, not one. Similarly below.
-- However, three instead of two cheap scalar operations is benign
-- and there's no risk of exponential blowup via duplicating
-- variables, so the ability to simplify better is worth the overhead.
-- OTOH, we don't declare scalar operations small in astIsSmall,
-- because that could lead to such operations being inlined an unbounded
-- number of times into expressions.
-- Note that due to non-scalar versions of these rules being banned,
-- we get different terms depending on the form of conversions
-- and rank 0 arrays.
u (AstPlusK v w) | Just u0 <- unAstK u ->
astPlusK (astTimesK (fromPlain $ AstConcreteK u0) v)
(astTimesK (fromPlain $ AstConcreteK u0) w)
u (AstTimesK (AstPlusK v w) x) | Just u0 <- unAstK u ->
astPlusK (astTimesK (fromPlain $ AstConcreteK u0) v)
(astTimesK (fromPlain $ AstConcreteK u0) w) * x
{- TODO: such rules increase the number of occurrences of a variable
and trade multiplication and quotient for an equally problematic remnant,
so they are disabled until we find a way to profit from them.
-- With static shapes, the second argument to QuotOp and RemOp
-- is often a constant, which makes such rules worth including,
-- since they are likely to fire. To help them fire, we avoid changing
-- that constant, if possible, e.g., in rules for NegateOp.
(AstConcreteK n) (Ast.AstI2K QuotOp (Ast.AstVar var) (AstConcreteK n'))
| n == n' ->
AstPlusK
(astVar var)
(negate (Ast.AstI2K RemOp (astVar var) (AstConcreteK n))) -}
u v -> AstTimesK u v
astN1K :: (NumScalar r, KnownSpan s)
=> OpCodeNum1 -> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
astN1K opCode t = case t of
Ast.AstLet var n k -> astLet var n (astN1K opCode k)
Ast.AstPrimalPart n -> primalPart (astN1K opCode n)
Ast.AstDualPart n -> dualPart (astN1K opCode n)
Ast.AstPlainPart n -> plainPart (astN1K opCode n)
Ast.AstFromPrimal n -> fromPrimal (astN1K opCode n)
Ast.AstFromDual n -> fromDual (astN1K opCode n)
Ast.AstFromPlain n -> fromPlain (astN1K opCode n)
Ast.AstCondK b n k -> astCondK b (astN1K opCode n) (astN1K opCode k)
_ -> case (opCode, t) of
(NegateOp, AstConcreteK n) -> AstConcreteK (negate n)
(NegateOp, AstPlusK u v) -> AstPlusK (negate u) (negate v)
(NegateOp, AstTimesK u v) -> astTimesK (negate u) v
(NegateOp, Ast.AstN1K NegateOp u) -> u
(NegateOp, Ast.AstN1K SignumOp u) -> astN1K SignumOp (negate u)
(NegateOp, Ast.AstI2K QuotOp u v) -> astI2K QuotOp (negate u) v
-- v is positive and let's keep it so
(NegateOp, Ast.AstI2K RemOp u v) -> astI2K RemOp (negate u) v
(AbsOp, AstConcreteK n) -> AstConcreteK (abs n)
(AbsOp, Ast.AstN1K AbsOp u) -> astN1K AbsOp u
(AbsOp, Ast.AstN1K NegateOp u) -> astN1K AbsOp u
(SignumOp, AstConcreteK n) -> AstConcreteK (signum n)
(SignumOp, Ast.AstN1K SignumOp u) -> astN1K SignumOp u
_ -> Ast.AstN1K opCode t
astR1K :: (NumScalar r, Differentiable r, KnownSpan s)
=> OpCode1 -> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
astR1K opCode = \case
Ast.AstLet var n k -> astLet var n (astR1K opCode k)
Ast.AstPrimalPart u -> primalPart $ astR1K opCode u
Ast.AstPlainPart u -> plainPart $ astR1K opCode u
Ast.AstFromPrimal u -> fromPrimal $ astR1K opCode u
Ast.AstFromPlain u -> fromPlain $ astR1K opCode u
Ast.AstCondK b n k -> astCondK b (astR1K opCode n) (astR1K opCode k)
AstConcreteK u -> case opCode of
RecipOp -> AstConcreteK $ recip u
ExpOp -> AstConcreteK $ exp u
LogOp -> AstConcreteK $ log u
SqrtOp -> AstConcreteK $ sqrt u
SinOp -> AstConcreteK $ sin u
CosOp -> AstConcreteK $ cos u
TanOp -> AstConcreteK $ tan u
AsinOp -> AstConcreteK $ asin u
AcosOp -> AstConcreteK $ acos u
AtanOp -> AstConcreteK $ atan u
SinhOp -> AstConcreteK $ sinh u
CoshOp -> AstConcreteK $ cosh u
TanhOp -> AstConcreteK $ tanh u
AsinhOp -> AstConcreteK $ asinh u
AcoshOp -> AstConcreteK $ acosh u
AtanhOp -> AstConcreteK $ atanh u
u -> Ast.AstR1K opCode u
astR2K :: (NumScalar r, Differentiable r, KnownSpan s)
=> OpCode2
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
astR2K opCode = \cases
(Ast.AstLet var n u) v -> astLetRefresh var n u $ \u' -> astR2K opCode u' v
u (Ast.AstLet var n v) -> astLetRefresh var n v $ \v' -> astR2K opCode u v'
(Ast.AstPrimalPart u) (Ast.AstPrimalPart v) -> primalPart $ astR2K opCode u v
(Ast.AstPlainPart @_ @s1 u) (Ast.AstPlainPart @_ @s2 v)
| Just Refl <- testEquality (knownSpan @s1) (knownSpan @s2) ->
plainPart $ astR2K opCode u v
(Ast.AstFromPrimal u) (Ast.AstFromPrimal v) -> fromPrimal $ astR2K opCode u v
(Ast.AstFromPlain u) (Ast.AstFromPlain v) -> fromPlain $ astR2K opCode u v
(AstConcreteK u) (AstConcreteK v) -> case opCode of
DivideOp -> AstConcreteK $ u / v
PowerOp -> AstConcreteK $ u ** v
LogBaseOp -> AstConcreteK $ logBase u v
Atan2Op -> AstConcreteK $ atan2H u v
u v -> case opCode of
DivideOp -> case u of
_ | Just 0 <- unAstK u -> u
_ | Just 1 <- unAstK u -> recip v
_ | Just 1 <- unAstK v -> u
_ | Just 0 <- unAstK v -> u -- the partiality-removal hack
AstTimesK n t
| eqY n v -> t
_ -> Ast.AstR2K DivideOp u v
-- TODO: add other rules that are relatively numerically stable
_ -> Ast.AstR2K opCode u v
astI2K :: (NumScalar r, IntegralH r, Nested.IntElt r, KnownSpan s)
=> OpCodeIntegral2
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
astI2K opCode = \cases
(Ast.AstLet var n u) v -> astLetRefresh var n u $ \u' -> astI2K opCode u' v
u (Ast.AstLet var n v) -> astLetRefresh var n v $ \v' -> astI2K opCode u v'
(Ast.AstPrimalPart n) (Ast.AstPrimalPart k) ->
primalPart (astI2K opCode n k)
(Ast.AstPlainPart @_ @s1 n) (Ast.AstPlainPart @_ @s2 k)
| Just Refl <- testEquality (knownSpan @s1) (knownSpan @s2) ->
plainPart (astI2K opCode n k)
(Ast.AstFromPrimal n) (Ast.AstFromPrimal k) ->
fromPrimal (astI2K opCode n k)
(Ast.AstFromPlain n) (Ast.AstFromPlain k) ->
fromPlain (astI2K opCode n k)
u v | let t = Ast.AstI2K opCode u v
, Just (u1, u2) <- bounds t
, u1 == u2 ->
fromPlain $ AstConcreteK u1
u v -> case opCode of
QuotOp -> case (u, v) of
_ | Just u0 <- unAstK u
, Just v0 <- unAstK v -> fromPlain $ AstConcreteK (quotH u0 v0)
_ | Just 0 <- unAstK u -> u
_ | Just 1 <- unAstK v -> u
_ | Just 0 <- unAstK v -> u -- the partiality-removal hack
_ | Just v0 <- unAstK v, v0 < 0 ->
astI2K QuotOp (negate u) (fromPlain $ AstConcreteK $ negate v0)
(Ast.AstI2K RemOp _ k, _)
| Just k0 <- unAstK k
, Just v0 <- unAstK v
, v0 >= k0 && k0 >= 0 -> fromPlain $ AstConcreteK 0
(Ast.AstI2K QuotOp u0 v0, _) ->
astI2K QuotOp u0 (astTimesK v0 v)
(AstTimesK n t, _)
| eqY n v -> t
_ -> Ast.AstI2K QuotOp u v
RemOp -> case (u, v) of
_ | Just u0 <- unAstK u
, Just v0 <- unAstK v -> fromPlain $ AstConcreteK (remH u0 v0)
_ | Just 0 <- unAstK u -> u
_ | Just 1 <- unAstK v -> fromPlain $ AstConcreteK 0
_ | Just 0 <- unAstK v -> fromPlain $ AstConcreteK 0 -- the hack
_ | Just w <- unAstK v, w < 0 ->
astI2K RemOp u (fromPlain $ AstConcreteK $ negate w)
(Ast.AstI2K RemOp t k, _)
| Just k0 <- unAstK k
, Just v0 <- unAstK v
, v0 >= k0 && k0 >= 0 -> astI2K RemOp t k
(Ast.AstI2K RemOp t k, _)
| Just k0 <- unAstK k
, Just v0 <- unAstK v
, remH k0 v0 == 0 && k0 > 0 -> astI2K RemOp t v
(AstTimesK n _, _)
| Just n0 <- unAstK n
, Just v0 <- unAstK v
, remH n0 v0 == 0 -> fromPlain $ AstConcreteK 0
_ -> Ast.AstI2K RemOp u v
astConcreteK :: GoodScalar r
=> Concrete (TKScalar r)
-> AstTensor AstMethodLet PlainSpan (TKScalar r)
astConcreteK = AstConcreteK . unConcrete
astFloorK :: (NumScalar r1, Differentiable r1, NumScalar r2, Integral r2)
=> AstTensor AstMethodLet PlainSpan (TKScalar r1)
-> AstTensor AstMethodLet PlainSpan (TKScalar r2)
astFloorK t = case t of
Ast.AstLet var u v -> astLet var u (astFloorK v)
-- This increases work and term size, because conditional is eager.
-- Ast.AstCond b a2 a3 -> Ast.AstCond b (astFloorK a2) (astFloorK a3)
-- These values are small, so we can simplify them ASAP.
AstConcreteK k -> astConcreteK (tkfloor $ Concrete k)
Ast.AstFloorK v -> astFloorK v
Ast.AstFromIntegralK v -> astFromIntegralK v
Ast.AstCastK v -> astFloorK v
_ -> Ast.AstFloorK t
-- Beware that increasing the number of calls to this constructor
-- sometimes increases runtime, because not enough copies cancel out.
astFromIntegralK :: forall r1 r2. (NumScalar r1, Integral r1, NumScalar r2)
=> AstTensor AstMethodLet PlainSpan (TKScalar r1)
-> AstTensor AstMethodLet PlainSpan (TKScalar r2)
astFromIntegralK t = case t of
_ | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) -> t
Ast.AstLet var u v -> astLet var u (astFromIntegralK v)
AstConcreteK k -> astConcreteK (tkfromIntegral $ Concrete k)
Ast.AstN1K opCode u -> astN1K opCode (astFromIntegralK u)
Ast.AstFromIntegralK v -> astFromIntegralK v
_ -> Ast.AstFromIntegralK t
astCastK :: forall r1 r2 s.
( NumScalar r1, Differentiable r1, NumScalar r2, Differentiable r2
, KnownSpan s )
=> AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
astCastK t = case t of
_ | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) -> t
Ast.AstLet var u v -> astLet var u (astCastK v)
AstConcreteK k -> astConcreteK (tkcast $ Concrete k)
-- TODO: which should go deeper, casts or fromPrimal? Or maybe alternate
-- in different phases to make sure both can cancel out?
-- Rethink. For now, astFromPrimal is not called, to avoid loops.
-- The same with many others.
Ast.AstFromPrimal v -> fromPrimal $ astCastK v
Ast.AstFromDual v -> fromDual $ astCastK v
Ast.AstFromPlain v -> fromPlain $ astCastK v
Ast.AstN1K opCode u -> astN1K opCode (astCastK u)
Ast.AstR1K opCode u -> astR1K opCode (astCastK u)
Ast.AstFromIntegralK v -> astFromIntegralK v
Ast.AstCastK v -> astCastK v
_ -> Ast.AstCastK t
astArgMinK :: forall n r. NumScalar r
=> AstTensor AstMethodLet PlainSpan (TKS '[n] r)
-> AstTensor AstMethodLet PlainSpan (TKScalar Int)
astArgMinK t = case t of
Ast.AstLet var u v -> astLet var u (astArgMinK v)
AstConcreteS v -> astConcreteK (tkargMin $ Concrete v)
Ast.AstN1S NegateOp u -> astArgMaxK u
Ast.AstFromIntegralS v -> astArgMinK v
Ast.AstCastS v -> astArgMinK v
Ast.AstReplicateK{} -> 0 -- we mask undefined if n is zero; tough luck
Ast.AstReplicateS{} -> 0
Ast.AstIotaS{} -> 0
_ -> Ast.AstArgMinK t
astArgMaxK :: forall n r. NumScalar r
=> AstTensor AstMethodLet PlainSpan (TKS '[n] r)
-> AstTensor AstMethodLet PlainSpan (TKScalar Int)
astArgMaxK t = case t of
Ast.AstLet var u v -> astLet var u (astArgMaxK v)
AstConcreteS v -> astConcreteK (tkargMax $ Concrete v)
Ast.AstN1S NegateOp u -> astArgMinK u
Ast.AstFromIntegralS v -> astArgMaxK v
Ast.AstCastS v -> astArgMaxK v
Ast.AstReplicateK{} -> 0 -- we mask undefined if n is zero; tough luck
Ast.AstReplicateS{} -> 0
Ast.AstIotaS{} -> 0
_ -> Ast.AstArgMaxK t
astIndexK
:: forall shm s r. (GoodScalar r, KnownSpan s)
=> AstTensor AstMethodLet s (TKS shm r) -> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKScalar r)
astIndexK = astIndexKnobsK defaultKnobs
astIndexKnobsK
:: forall shm s r. (GoodScalar r, KnownSpan s)
=> SimplifyKnobs
-> AstTensor AstMethodLet s (TKS shm r) -> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKScalar r)
astIndexKnobsK _ v0 ZIS = kfromS v0
astIndexKnobsK _ v0 (i1 :.$ _)
| Just (lb, ub) <- intBounds i1
, FTKS (snat :$$ _) _ <- ftkAst v0
, ub < 0 || lb >= fromSNat' snat = fromPlain $ AstConcreteK def
astIndexKnobsK knobs u0 (Ast.AstCondK b v w :.$ rest0)
| knobPhase knobs `notElem` [PhaseUnspecified, PhaseVectorization] =
-- don't undo vectorization tweaks
astLetFun u0 $ \u ->
shareIx rest0 $ \rest ->
astCondK b (astIndexKnobsK knobs u (v :.$ rest))
(astIndexKnobsK knobs u (w :.$ rest))
astIndexKnobsK knobs v0 ix@(i1 :.$ rest1)
| FTKS ((:$$) @in1 @shm1 SNat shm1) _ <- ftkAst v0
, let tryRecursing :: AstTensor AstMethodLet s (TKS shm1 r)
-> AstTensor AstMethodLet s (TKScalar r)
tryRecursing u =
case 0 <=. i1 &&* i1 <=. valueOf @in1 - 1 of
AstConcreteK b ->
if b then astIndex u rest1 else fromPlain $ AstConcreteK def
_ -> Ast.AstIndexK v0 ix = case v0 of
Ast.AstProject1{} -> Ast.AstIndexK v0 ix
Ast.AstProject2{} -> Ast.AstIndexK v0 ix
Ast.AstApply{} -> Ast.AstIndexK v0 ix
Ast.AstVar{} -> Ast.AstIndexK v0 ix
Ast.AstBuild1 _ STKScalar (var2, v) | ZIS <- rest1 ->
tryRecursing $ sfromK $ astLet var2 i1 v
Ast.AstBuild1 _ STKS{} (var2, v) ->
tryRecursing $ astLet var2 i1 v
Ast.AstLet var u v -> astLetRefresh var u v $ \v' -> astIndex v' ix
Ast.AstPrimalPart{} -> Ast.AstIndexK v0 ix -- must be a NF
Ast.AstDualPart{} -> Ast.AstIndexK v0 ix
Ast.AstPlainPart{} -> Ast.AstIndexK v0 ix
Ast.AstFromPrimal v -> fromPrimal $ astIndex v ix
Ast.AstFromDual v -> fromDual $ astIndex v ix
Ast.AstFromPlain v -> fromPlain $ astIndex v ix
AstPlusS u v ->
shareIx ix $ \ !ix2 -> astPlusK (astIndex u ix2)
(astIndex v ix2)
AstTimesS u v ->
shareIx ix $ \ !ix2 -> astTimesK (astIndex u ix2)
(astIndex v ix2)
Ast.AstN1S opCode u -> astN1K opCode (astIndex u ix)
Ast.AstR1S opCode u -> astR1K opCode (astIndex u ix)
Ast.AstR2S opCode u v ->
shareIx ix $ \ !ix2 -> astR2K opCode (astIndex u ix2)
(astIndex v ix2)
Ast.AstI2S opCode u v ->
shareIx ix $ \ !ix2 -> astI2K opCode (astIndex u ix2)
(astIndex v ix2)
AstConcreteS a | AstConcreteK i <- i1 ->
let u = withKnownShS shm1 $
tsindex (Concrete a) (Concrete i :.$ ZIS)
in astIndex (astConcreteS u) rest1
AstConcreteS{} -> case unRepl1 v0 of
Just u -> tryRecursing u
_ -> Ast.AstIndexK v0 ix
Ast.AstFloorS v -> astFloorK $ astIndex v ix
Ast.AstFromIntegralS v -> astFromIntegralK $ astIndex v ix
Ast.AstCastS t -> astCastK $ astIndex t ix
Ast.AstArgMinS @k @shz v | FTKS nsh _ <- ftkAst v
, nl@(SNat @nl) <- shsLast nsh ->
gcastWith (unsafeCoerceRefl :: shm ++ '[nl] :~: k ': shz) $
astArgMinK @nl
$ astIndexKnobsS (deVect knobs) (nl :$$ ZSS) v ix
Ast.AstArgMaxS @k @shz v | FTKS nsh _ <- ftkAst v
, nl@(SNat @nl) <- shsLast nsh ->
gcastWith (unsafeCoerceRefl :: shm ++ '[nl] :~: k ': shz) $
astArgMaxK @nl
$ astIndexKnobsS (deVect knobs) (nl :$$ ZSS) v ix
Ast.AstIndexS _ v ix2 -> astIndex v (ix2 `ixsAppend` ix)
Ast.AstCondS b v w ->
shareIx ix $ \ !ix2 -> astCondK b (astIndex v ix2)
(astIndex w ix2)
Ast.AstFromVectorK (_ :$$ ZSS) l | AstConcreteK i <- i1, ZIS <- rest1 ->
tryRecursing $ sfromK (l V.! i)
Ast.AstFromVectorK{} -> -- normal form
Ast.AstIndexK v0 ix
Ast.AstFromVectorS (_ :$$ ZSS) l | AstConcreteK i <- i1 ->
tryRecursing (l V.! i)
Ast.AstFromVectorS @shm2 @shn2 shm2 l | shsLength shm2 < length ix
, SNat <- shsRank shm2 ->
gcastWith (unsafeCoerceRefl :: shn2 :~: Drop (Rank shm2) shm) $
gcastWith (unsafeCoerceRefl :: Take (Rank shm2) shm :~: shm2) $
shareIx ix $ \ !ix2 ->
let ixRest2 = ixsDrop @(Rank shm2) ix2
in Ast.AstIndexK (astFromVectorK @shm2 shm2
$ V.map (\a -> astIndex a ixRest2) l)
(ixsTake @(Rank shm2) ix2)
Ast.AstFromVectorS{} -> -- normal form (shsLength shm2 >= length ix)
Ast.AstIndexK v0 ix
Ast.AstSumS @shm2 shm2 v | Dict0 <- numFromTKAllNum (Proxy @r) ->
let perm1 = backpermCycleN (shsLength shm2) (shsLength shm2 + length ix)
in Permutation.permFromListCont perm1 $ \(perm :: Permutation.Perm perm) ->
gcastWith (unsafeCoerceRefl
:: (Rank perm <=? Rank (shm2 ++ shm)) :~: True) $
gcastWith (unsafeCoerceRefl
:: Permutation.PermutePrefix perm (shm2 ++ shm)
:~: shm ++ shm2) $
fromMaybe (error "astIndexKnobsK: impossible non-permutation")
$ Permutation.permCheckPermutation perm
$ astSumK
$ astIndexKnobsS @shm (deVect knobs) shm2 (astTransposeS perm v) ix
Ast.AstScatterS shm7 shn7 shp7 v (vars, i5 :.$ ix2) | eqY i5 i1 ->
astIndex (astScatter shm7 shn7 (shsTail shp7) v (vars, ix2)) rest1
Ast.AstScatterS _ _ _ _ (_, AstConcreteK{} :.$ _)
| AstConcreteK{} <- i1 ->
fromPlain $ AstConcreteK def -- from above we know i5 /= i1
Ast.AstScatterS{} -> -- normal form
Ast.AstIndexK v0 ix
Ast.AstReplicateK (_ :$$ shmRest) v ->
tryRecursing $ astReplicateK shmRest v
Ast.AstReplicateS ZSS v -> astIndex v ix
Ast.AstReplicateS (_ :$$ shmRest) v ->
tryRecursing $ astReplicateS shmRest v
-- This is not a possible normal form, but pattern needs to be exhaustive.
Ast.AstGatherS @_ @_ @shp' _ _ _ v (ZS, ix2) ->
astIndex @(shp' ++ shm) v (ix2 `ixsAppend` ix)
Ast.AstGatherS (_ :$$ (shm71 :: ShS shm71)) shn' shp'
v (var2 ::$ vars, ix2) ->
let w :: AstTensor AstMethodLet s (TKS shm1 r)
w = astGather shm71 shn' shp' v (vars, ix2)
u = astLet var2 i1 $ astIndex @shm1 w rest1
-- This let makes it impossible to use astCond when i1 is OOB.
-- No astLetRefresh is needed, because the smart constructor
-- for AstGatherS ensures var2 doesn't occur in v, even bound.
in case 0 <=. i1 &&* i1 <=. valueOf @in1 - 1 of
AstConcreteK b -> if b then u else fromPlain $ AstConcreteK def
_ -> Ast.AstIndexK v0 ix
Ast.AstIotaS _ ->
astLetFun i1 $ \i ->
astCondKInitial (0 <=. i &&* i <=. valueOf @in1 - 1)
(astFromIntegralK i)
0
Ast.AstAppendS u v | FTKS (SNat @m :$$ _) _ <- ftkAst u ->
if knobPhase knobs == PhaseExpansion then
astLetFun i1 $ \i ->
shareIx rest1 $ \ !rest2 ->
let ulen = valueOf @m
ix1 = i :.$ rest2
ix2 = i - ulen :.$ rest2
in case ulen <=. i of
AstConcreteK b -> if b then astIndex v ix2 else astIndex u ix1
bExpr ->
-- This results in a larger term, so we consider this late.
astCondKInitial bExpr (astIndex v ix2) (astIndex u ix1)
else
-- We don't need to share i1 and rest1 here and, in fact, we can't
-- because we'd loop via building AstLet around the resulting AstIndexS.
let ulen = valueOf @m
ix1 = i1 :.$ rest1
ix2 = i1 - ulen :.$ rest1
in case ulen <=. i1 of
AstConcreteK b -> if b then astIndex v ix2 else astIndex u ix1
_ -> Ast.AstIndexK v0 ix
Ast.AstSliceS i@(SNat @i) (SNat @n) k@SNat v ->
astLetFun i1 $ \iShared ->
let b = (if fromSNat' i == 0 then true else 0 <=. iShared)
&&* (if fromSNat' k == 0 then true else iShared <=. valueOf @n - 1)
ii = valueOf @i + iShared
in astCondKInitial b (astIndex v (ii :.$ rest1))
$ fromPlain $ AstConcreteK def
Ast.AstReverseS v ->
let iRev = valueOf @in1 - 1 - i1
in astIndex v (iRev :.$ rest1)
Ast.AstTransposeS @_ @sh2 perm v
| fromSNat' (Permutation.permRank perm) <= length ix ->
-- TODO: remake once there's an S version of permInverse:
permInverse perm $ \(permR :: Permutation.Perm permR) _ ->
let ix2 :: AstIxS AstMethodLet (Permutation.PermutePrefix permR shm)
ix2 = ixsPermutePrefix permR ix
in gcastWith (unsafeCoerceRefl
:: sh2 :~: Permutation.PermutePrefix permR shm) $
astIndex @(Permutation.PermutePrefix permR shm) v ix2
Ast.AstTransposeS @perm perm v ->
astIndex (astTransposeAsGatherS @perm (deVect knobs) perm v) ix
Ast.AstReshapeS sh v ->
astIndex (astReshapeAsGatherS (deVect knobs) sh v) ix
-- Conversions to shaped need to stay down, so this is NF.
Ast.AstConvert{} -> Ast.AstIndexK v0 ix
Ast.AstDot1InS _ n u v ->
astDot0 (astIndexKnobsS (deVect knobs) (n :$$ ZSS) u ix)
(astIndexKnobsS (deVect knobs) (n :$$ ZSS) v ix)
Ast.AstMatmul2S{} -> Ast.AstIndexK v0 ix -- TODO
Ast.AstBoolNotS a -> Ast.AstBoolNotK $ astIndex a ix
Ast.AstBoolAndS b1 b2 -> Ast.AstBoolAndK (astIndex b1 ix)
(astIndex b2 ix)
Ast.AstLeqS _shb sh arg1 arg2 ->
Ast.AstLeq (astIndexKnobsS (deVect knobs) sh arg1 ix)
(astIndexKnobsS (deVect knobs) sh arg2 ix)
where
astIndex
:: forall shm' s' r'. (GoodScalar r', KnownSpan s')
=> AstTensor AstMethodLet s' (TKS shm' r')
-> AstIxS AstMethodLet shm'
-> AstTensor AstMethodLet s' (TKScalar r')
astIndex = astIndexKnobsK (deVect knobs)
astGather
:: forall shm' shn' shp' x'.
ShS shm' -> ShS shn' -> ShS shp'
-> AstTensor AstMethodLet s (TKS2 (shp' ++ shn') x')
-> (AstVarListS shm', AstIxS AstMethodLet shp')
-> AstTensor AstMethodLet s (TKS2 (shm' ++ shn') x')
astGather = astGatherKnobsS (deVect knobs)
astScatter
:: forall shm' shn' shp' x'. TKAllNum x'
=> ShS shm' -> ShS shn' -> ShS shp'
-> AstTensor AstMethodLet s (TKS2 (shm' ++ shn') x')
-> (AstVarListS shm', AstIxS AstMethodLet shp')
-> AstTensor AstMethodLet s (TKS2 (shp' ++ shn') x')
astScatter = astScatterKnobsS (deVect knobs)
-- Just as with AstPlusK, summands are flattened and a constant comes first.
-- We don't sum fromVectors, because it's costly and even if everything
-- nicely simplifies, we end up with an equaly big vector of the results.
-- OTOH, AstSum* and AstDot* eliminate at least one dimension, so they
-- are worth it, at least for small enough arguments and if the result
-- simplifies decently.
astPlusS :: (NumScalar r, KnownSpan s)
=> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
astPlusS = \cases
u v | FTKS (snat :$$ _) _ <- ftkAst u
, Just u0 <- unRepl1 u
, Just v0 <- unRepl1 v ->
astReplicateS (snat :$$ ZSS) $ astPlusS u0 v0
(Ast.AstPrimalPart u) (Ast.AstPrimalPart v) -> primalPart $ astPlusS u v
(Ast.AstDualPart u) (Ast.AstDualPart v) -> dualPart $ astPlusS u v
(Ast.AstPlainPart @_ @s1 u) (Ast.AstPlainPart @_ @s2 v)
| Just Refl <- testEquality (knownSpan @s1) (knownSpan @s2) ->
plainPart $ astPlusS u v
(Ast.AstFromPrimal u) (Ast.AstFromPrimal v) -> fromPrimal $ astPlusS u v
(Ast.AstFromDual u) (Ast.AstFromDual v) -> fromDual $ astPlusS u v
(Ast.AstFromPlain u) (Ast.AstFromPlain v) -> fromPlain $ astPlusS u v
(AstPlusS u w) v -> astPlusS u (astPlusS w v) -- flattening
u v | Just 0 <- unReplC u -> v
u v | Just 0 <- unReplC v -> u
u v | Just u0 <- unAstS u
, Just v0 <- unAstS v -> fromPlain $ AstConcreteS (u0 + v0)
u v | Just v0 <- unAstS v -> AstPlusS (fromPlain $ AstConcreteS v0) u
u (AstPlusS v w) | Just v0 <- unAstS v ->
astPlusS (fromPlain $ AstConcreteS v0) u + w
(Ast.AstN1S NegateOp t1) t2 | eqY t1 t2 ->
fromPlain $ AstConcreteS $ defTargetRep $ ftkAst t1
(Ast.AstN1S NegateOp t1) (AstPlusS t2 u) | eqY t1 t2 -> u
t2 (Ast.AstN1S NegateOp t1) | eqY t1 t2 ->
fromPlain $ AstConcreteS $ defTargetRep $ ftkAst t1
t2 (AstPlusS (Ast.AstN1S NegateOp t1) u) | eqY t1 t2 -> u
(AstTimesS n1 t1) (AstTimesS n2 t2)
| Just n10 <- unAstS n1
, Just n20 <- unAstS n2
, eqY t1 t2 -> fromPlain (AstConcreteS (n10 + n20)) `astTimesS` t1
(AstTimesS n1 t1) (AstPlusS (AstTimesS n2 t2) w)
| Just n10 <- unAstS n1
, Just n20 <- unAstS n2
, eqY t1 t2 -> fromPlain (AstConcreteS (n10 + n20)) `astTimesS` t1 + w
(AstConvUpSFromK u) (AstConvUpSFromK v) -> sfromK $ astPlusK u v
u (AstConvUpSFromK v) -> sfromK $ astPlusK (kfromS u) v
(AstConvUpSFromK u) v -> sfromK $ astPlusK u (kfromS v)
u@Ast.AstScatterS{} v@Ast.AstScatterS{} ->
fromMaybe (AstPlusS u v) $ astPlusSMaybe u v
u@Ast.AstScatterS{} v@(AstPlusS u2@Ast.AstScatterS{} w) ->
maybe (AstPlusS u v) (+ w) $ astPlusSMaybe u u2
u@(Ast.AstTransposeS perm1 a1@(Ast.AstScatterS @_ @shn1 @shp1 _ _ _ _ _))
v@(Ast.AstTransposeS perm2 a2@(Ast.AstScatterS @_ @shn2 @shp2 _ _ _ _ _))
| Just Refl <- testEquality perm1 perm2 ->
gcastWith (unsafeCoerceRefl :: shp1 ++ shn1 :~: shp2 ++ shn2) $
maybe (AstPlusS u v) (astTransposeS perm1) $ astPlusSMaybe a1 a2
u@(Ast.AstTransposeS perm1 a1@(Ast.AstScatterS @_ @shn1 @shp1 _ _ _ _ _))
v@(AstPlusS (Ast.AstTransposeS perm2 a2@(Ast.AstScatterS
@_ @shn2 @shp2 _ _ _ _ _)) w)
| Just Refl <- testEquality perm1 perm2 ->
gcastWith (unsafeCoerceRefl :: shp1 ++ shn1 :~: shp2 ++ shn2) $
maybe (AstPlusS u v) (\res -> astTransposeS perm1 res + w)
$ astPlusSMaybe a1 a2
u v -> AstPlusS u v
astPlusSMaybe :: (NumScalar r, KnownSpan s)
=> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> Maybe (AstTensor AstMethodLet s (TKS sh r))
astPlusSMaybe = \cases
(Ast.AstScatterS @_ @_ @shp (m :$$ shmRest) shn shp
u (var ::$ vars, ix))
(Ast.AstScatterS @_ @_ @shp2 (m2 :$$ shmRest2) shn2 shp2
u2 (var2 ::$ vars2, ix2))
| Just Refl <- testEquality shmRest shmRest2
, Just Refl <- testEquality shn shn2 ->
gcastWith (unsafeCoerceRefl :: shp :~: shp2) $
fuseScatters2 (m :$$ shmRest) shn shp u (var ::$ vars, ix)
(m2 :$$ shmRest2) shn2 shp2 u2 (var2 ::$ vars2, ix2)
-- Same as above, but with arguments reversed, in case it works better:
`mplus`
fuseScatters2 (m2 :$$ shmRest2) shn2 shp2 u2 (var2 ::$ vars2, ix2)
(m :$$ shmRest) shn shp u (var ::$ vars, ix)
{- TODO: not needed ATM, though this could simplify astScatterS
(Ast.AstScatterS @_ @_ @shp ZSS shn shp
u (_, ix))
(Ast.AstScatterS @_ @_ @shp2 ZSS shn2 shp2
u2 (_, ix2))
| Just Refl <- testEquality shn shn2 ->
gcastWith (unsafeCoerceRefl :: shp :~: shp2) $
fuseScatters0 shn shp u ix
shn2 shp2 u2 ix2
(Ast.AstScatterS @_ @_ @shp (m :$$ shmRest) shn shp
u (var ::$ vars, ix))
(Ast.AstScatterS @_ @_ @shp2 shm2 shn2 shp2
u2 (vars2, ix2))
| Just Refl <- testEquality shmRest shm2
, Just Refl <- testEquality shn shn2 ->
gcastWith (unsafeCoerceRefl :: shp :~: shp2) $
fuseScatters1 (m :$$ shmRest) shn shp u (var ::$ vars, ix)
shm2 shn2 shp2 u2 (vars2, ix2)
(Ast.AstScatterS @_ @_ @shp2 shm2 shn2 shp2
u2 (vars2, ix2))
(Ast.AstScatterS @_ @_ @shp (m :$$ shmRest) shn shp
u (var ::$ vars, ix))
| Just Refl <- testEquality shmRest shm2
, Just Refl <- testEquality shn shn2 ->
gcastWith (unsafeCoerceRefl :: shp :~: shp2) $
fuseScatters1 (m :$$ shmRest) shn shp u (var ::$ vars, ix)
shm2 shn2 shp2 u2 (vars2, ix2) -}
_ _ -> Nothing
{- TODO: not needed ATM
fuseScatters0 :: forall shn shp x s. (TKAllNum x, KnownSpan s)
=> ShS shn -> ShS shp
-> AstTensor AstMethodLet s (TKS2 shn x)
-> AstIxS AstMethodLet shp
-> ShS shn -> ShS shp
-> AstTensor AstMethodLet s (TKS2 shn x)
-> AstIxS AstMethodLet shp
-> Maybe (AstTensor AstMethodLet s (TKS2 (shp ++ shn) x))
fuseScatters0 shn shp u ix
_shn2 _shp2 u2 ix2 =
funToAstIntMaybe (Just (0, 1)) $ \(var', ast') ->
let g i i2 = astCond (AstConcreteK 1 <=. ast')
i2 i
ix3 = ixsZipWith g ix ix2
a = astFromVector (SNat @2) (STKS shn (stkAstX u)) (V.fromList [u, u2])
in if any isCond ix3
then Nothing -- a risk of nested conditional build-up
else Just $ astScatterS (SNat @2 :$$ ZSS) shn shp
a (var' ::$ ZS, ix3)
fuseScatters1 :: forall m shm shn shp x s. (TKAllNum x, KnownSpan s)
=> ShS (m : shm) -> ShS shn -> ShS shp
-> AstTensor AstMethodLet s (TKS2 ((m : shm) ++ shn) x)
-> (AstVarListS (m : shm), AstIxS AstMethodLet shp)
-> ShS shm -> ShS shn -> ShS shp
-> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> Maybe (AstTensor AstMethodLet s (TKS2 (shp ++ shn) x))
fuseScatters1 (m :$$ shmRest) shn shp u (var ::$ vars, ix)
_shm2 _shn2 _shp2 u2 (vars2, ix2) =
let ub = fromSNat' m
var' = reboundsVarName (0, ub) var
f i2 = foldr (\(v, v2) -> substituteAst (astVar v) v2) i2
(ixsZip vars vars2)
g i i2 = astCond (AstConcreteK (fromSNat' m) <=. astVar var')
(f i2) i -- i has, effectively, varFalse in it
ix3 = ixsZipWith g ix ix2
a = astReplicateS (SNat @1 :$$ ZSS) u2
in if any isCond ix3
then Nothing -- a risk of nested conditional build-up
else Just $ astScatterS (SS m :$$ shmRest) shn shp
(astAppendS u a) (var' ::$ vars, ix3) -}
fuseScatters2 :: forall m m2 shm shn shp x s. (TKAllNum x, KnownSpan s)
=> ShS (m : shm) -> ShS shn -> ShS shp
-> AstTensor AstMethodLet s (TKS2 ((m : shm) ++ shn) x)
-> (AstVarListS (m : shm), AstIxS AstMethodLet shp)
-> ShS (m2 : shm) -> ShS shn -> ShS shp
-> AstTensor AstMethodLet s (TKS2 ((m2 : shm) ++ shn) x)
-> (AstVarListS (m2 : shm), AstIxS AstMethodLet shp)
-> Maybe (AstTensor AstMethodLet s (TKS2 (shp ++ shn) x))
fuseScatters2 (m :$$ shmRest) shn shp u (var ::$ vars, ix)
(m2 :$$ _shmRest2) _shn2 _shp2 u2 (var2 ::$ vars2, ix2) =
let ub = fromSNat' m + fromSNat' m2 - 1
var' = reboundsVarName (0, ub) var
astVarTrue = astVar $ reboundsVarName (fromSNat' m, ub) var
f i2 = substituteAst (astVarTrue - AstConcreteK (fromSNat' m)) var2
$ foldr (\(v, v2) -> substituteAst (astVar v) v2)
i2
(ixsZip (unAstVarListS vars) (unAstVarListS vars2))
g i i2 = astCondK (AstConcreteK (fromSNat' m) <=. astVar var')
(f i2) i -- i has, effectively, varFalse in it
ix3 = ixsZipWith g ix ix2
in if any isCond ix3
then Nothing -- a risk of nested conditional build-up
else Just $ astScatterS (snatPlus m m2 :$$ shmRest) shn shp
(astAppendS u u2) (var' ::$ vars, ix3)
-- Checks whether it's a conditional or something easily obtained
-- from conditionals via rewrite rules.
isCond :: AstInt AstMethodLet -> Bool
isCond (Ast.AstLet _var _u v) = isCond v
isCond Ast.AstCondK{} = True
isCond (AstPlusK _ Ast.AstCondK{}) = True
isCond (AstPlusK Ast.AstCondK{} _) = True
isCond (AstTimesK _ Ast.AstCondK{}) = True
isCond (AstTimesK Ast.AstCondK{} _) = True
isCond _ = False
-- Just as with astTimesK, factors are flattened and a constant comes first.
-- Also, negation comes first.
astTimesS :: (NumScalar r, KnownSpan s)
=> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
astTimesS = \cases
u v | FTKS (snat :$$ _) _ <- ftkAst u
, Just u0 <- unRepl1 u
, Just v0 <- unRepl1 v ->
astReplicateS (snat :$$ ZSS) $ astTimesS u0 v0
(Ast.AstPrimalPart u) (Ast.AstPrimalPart v) -> primalPart $ astTimesS u v
(Ast.AstPlainPart @_ @s1 u) (Ast.AstPlainPart @_ @s2 v)
| Just Refl <- testEquality (knownSpan @s1) (knownSpan @s2) ->
plainPart $ astTimesS u v
(Ast.AstFromPrimal u) (Ast.AstFromPrimal v) -> fromPrimal $ astTimesS u v
u@Ast.AstFromDual{} Ast.AstFromDual{} ->
fromPlain $ AstConcreteS $ defTargetRep $ ftkAst u
(Ast.AstFromPlain u) (Ast.AstFromPlain v) -> fromPlain $ astTimesS u v
(AstTimesS u w) v -> astTimesS u (astTimesS w v) -- flattening
u _ | Just 0 <- unReplC u -> u
_ v | Just 0 <- unReplC v -> v
u v | Just 1 <- unReplC u -> v
u v | Just 1 <- unReplC v -> u
u v | Just (-1) <- unReplC u -> negate v
u v | Just (-1) <- unReplC v -> negate u
u v | Just u0 <- unAstS u
, Just v0 <- unAstS v -> fromPlain $ AstConcreteS (u0 * v0)
u v | Just v0 <- unAstS v -> AstTimesS (fromPlain $ AstConcreteS v0) u
u (AstTimesS v w) | Just v0 <- unAstS v ->
astTimesS (fromPlain $ AstConcreteS v0) u * w
(Ast.AstN1S NegateOp u) (Ast.AstN1S NegateOp v) -> astTimesS u v
(Ast.AstN1S NegateOp u) (AstTimesS (Ast.AstN1S NegateOp v) w) ->
astTimesS u v * w
u (Ast.AstN1S NegateOp v) -> astTimesS (negate u) v
u (AstTimesS (Ast.AstN1S NegateOp v) w) -> astTimesS (negate u) v * w
(Ast.AstR2S DivideOp u1 u2) v -> astR2S DivideOp (u1 * v) u2
u (Ast.AstR2S DivideOp v1 v2) -> astR2S DivideOp (u * v1) v2
(Ast.AstR1S RecipOp u) v -> astR2S DivideOp v u
u (Ast.AstR1S RecipOp v) -> astR2S DivideOp u v
(Ast.AstScatterS shm shn shp v (vars, ix)) u
| Just w <- unRepl u, FTKS shv _ <- ftkAst v ->
Ast.AstScatterS shm shn shp (v `astTimesS` astReplicateK shv (kfromS w))
(vars, ix)
(Ast.AstScatterS shm shn shp v (vars, ix)) (AstTimesS u a)
| Just w <- unRepl u, FTKS shv _ <- ftkAst v ->
astTimesS
(Ast.AstScatterS shm shn shp (v `astTimesS` astReplicateK
shv (kfromS w))
(vars, ix)) a
u (Ast.AstScatterS shm shn shp v (vars, ix))
| Just w <- unRepl u, FTKS shv _ <- ftkAst v ->
Ast.AstScatterS shm shn shp (astReplicateK shv (kfromS w) `astTimesS` v)
(vars, ix)
u (AstTimesS (Ast.AstScatterS shm shn shp v (vars, ix)) a)
| Just w <- unRepl u, FTKS shv _ <- ftkAst v ->
astTimesS
(Ast.AstScatterS shm shn shp (astReplicateK shv (kfromS w)
`astTimesS` v)
(vars, ix)) a
(Ast.AstGatherS shm shn shp v (vars, ix)) u
| Just w <- unRepl u, FTKS shv _ <- ftkAst v ->
Ast.AstGatherS shm shn shp (v `astTimesS` astReplicateK shv (kfromS w))
(vars, ix)
(Ast.AstGatherS shm shn shp v (vars, ix)) (AstTimesS u a)
| Just w <- unRepl u, FTKS shv _ <- ftkAst v ->
astTimesS
(Ast.AstGatherS shm shn shp (v `astTimesS` astReplicateK shv (kfromS w))
(vars, ix)) a
u (Ast.AstGatherS shm shn shp v (vars, ix))
| Just w <- unRepl u, FTKS shv _ <- ftkAst v ->
Ast.AstGatherS shm shn shp (astReplicateK shv (kfromS w) `astTimesS` v)
(vars, ix)
u (AstTimesS (Ast.AstGatherS shm shn shp v (vars, ix)) a)
| Just w <- unRepl u, FTKS shv _ <- ftkAst v ->
astTimesS
(Ast.AstGatherS shm shn shp (astReplicateK shv (kfromS w) `astTimesS` v)
(vars, ix)) a
(AstConvUpSFromK u) (AstConvUpSFromK v) -> sfromK $ astTimesK u v
u (AstConvUpSFromK v) -> sfromK $ astTimesK (kfromS u) v
(AstConvUpSFromK u) v -> sfromK $ astTimesK u (kfromS v)
-- This breaks sharing, because although u is concrete and so doesn't
-- have to be shared, the multiplication is not shared --- we end up
-- with one addition and two multiplications, not one.
-- u@AstConcreteS{} * AstPlusS v w -> AstPlusS (astTimesS u v) (u * w)
u v -> AstTimesS u v
astN1S :: (NumScalar r, KnownSpan s)
=> OpCodeNum1 -> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
astN1S opCode t = case t of
_ | FTKS (snat :$$ _) _ <- ftkAst t
, Just t0 <- unRepl1 t ->
astReplicateS (snat :$$ ZSS) (astN1S opCode t0)
-- This increases work:
-- Ast.AstCond b n k -> astCond b (astN1S opCode n) (astN1S opCode k)
Ast.AstLet var n k -> astLet var n (astN1S opCode k)
Ast.AstPrimalPart n -> primalPart (astN1S opCode n)
Ast.AstDualPart n -> dualPart (astN1S opCode n)
Ast.AstPlainPart n -> plainPart (astN1S opCode n)
Ast.AstFromPrimal n -> fromPrimal (astN1S opCode n)
Ast.AstFromDual n -> fromDual (astN1S opCode n)
Ast.AstFromPlain n -> fromPlain (astN1S opCode n)
AstConvUpSFromK n -> sfromK $ astN1K opCode n
_ -> case (opCode, t) of
(NegateOp, AstConcreteS n) -> AstConcreteS (negate n)
(NegateOp, AstPlusS u v) -> AstPlusS (negate u) (negate v)
(NegateOp, AstTimesS u v) -> astTimesS (negate u) v
(NegateOp, Ast.AstN1S NegateOp u) -> u
(NegateOp, Ast.AstN1S SignumOp u) -> astN1S SignumOp (negate u)
(NegateOp, Ast.AstI2S QuotOp u v) -> astI2S QuotOp (negate u) v
-- v is likely positive and let's keep it so
(NegateOp, Ast.AstI2S RemOp u v) -> astI2S RemOp (negate u) v
(NegateOp, Ast.AstScatterS shm shn shp v (vars, ix)) ->
astScatterS shm shn shp (negate v) (vars, ix)
(NegateOp, Ast.AstGatherS shm shn shp v (vars, ix)) ->
astGatherS shm shn shp (negate v) (vars, ix)
(AbsOp, AstConcreteS n) -> AstConcreteS (abs n)
(AbsOp, Ast.AstN1S AbsOp u) -> astN1S AbsOp u
(AbsOp, Ast.AstN1S NegateOp u) -> astN1S AbsOp u
(SignumOp, AstConcreteS n) -> AstConcreteS (signum n)
(SignumOp, Ast.AstN1S SignumOp u) -> astN1S SignumOp u
_ -> Ast.AstN1S opCode t
astR1S :: (NumScalar r, Differentiable r, KnownSpan s)
=> OpCode1 -> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
astR1S opCode = \case
t | FTKS (snat :$$ _) _ <- ftkAst t
, Just t0 <- unRepl1 t ->
astReplicateS (snat :$$ ZSS) (astR1S opCode t0)
-- This increases work:
-- Ast.AstCond b n k -> astCond b (astR1S opCode n) (astR1S opCode k)
Ast.AstLet var n k -> astLet var n (astR1S opCode k)
Ast.AstPrimalPart u -> primalPart $ astR1S opCode u
Ast.AstPlainPart u -> plainPart $ astR1S opCode u
Ast.AstFromPrimal u -> fromPrimal $ astR1S opCode u
Ast.AstFromPlain u -> fromPlain $ astR1S opCode u
AstConvUpSFromK u -> sfromK $ astR1K opCode u
AstConcreteS u -> case opCode of
RecipOp -> AstConcreteS $ recip u
ExpOp -> AstConcreteS $ exp u
LogOp -> AstConcreteS $ log u
SqrtOp -> AstConcreteS $ sqrt u
SinOp -> AstConcreteS $ sin u
CosOp -> AstConcreteS $ cos u
TanOp -> AstConcreteS $ tan u
AsinOp -> AstConcreteS $ asin u
AcosOp -> AstConcreteS $ acos u
AtanOp -> AstConcreteS $ atan u
SinhOp -> AstConcreteS $ sinh u
CoshOp -> AstConcreteS $ cosh u
TanhOp -> AstConcreteS $ tanh u
AsinhOp -> AstConcreteS $ asinh u
AcoshOp -> AstConcreteS $ acosh u
AtanhOp -> AstConcreteS $ atanh u
u -> Ast.AstR1S opCode u
astR2S :: (NumScalar r, Differentiable r, KnownSpan s)
=> OpCode2
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
astR2S opCode = \cases
u v | FTKS (snat :$$ _) _ <- ftkAst u
, Just u0 <- unRepl1 u
, Just v0 <- unRepl1 v ->
astReplicateS (snat :$$ ZSS) (astR2S opCode u0 v0)
(Ast.AstPrimalPart u) (Ast.AstPrimalPart v) -> primalPart $ astR2S opCode u v
(Ast.AstPlainPart @_ @s1 u) (Ast.AstPlainPart @_ @s2 v)
| Just Refl <- testEquality (knownSpan @s1) (knownSpan @s2) ->
plainPart $ astR2S opCode u v
(Ast.AstFromPrimal u) (Ast.AstFromPrimal v) -> fromPrimal $ astR2S opCode u v
(Ast.AstFromPlain u) (Ast.AstFromPlain v) -> fromPlain $ astR2S opCode u v
(AstConvUpSFromK n) (AstConvUpSFromK k) ->
sfromK $ astR2K opCode n k
n (AstConvUpSFromK k) ->
sfromK $ astR2K opCode (kfromS n) k
(AstConvUpSFromK n) k ->
sfromK $ astR2K opCode n (kfromS k)
(AstConcreteS u) (AstConcreteS v) -> case opCode of
DivideOp -> AstConcreteS $ u / v
PowerOp -> AstConcreteS $ u ** v
LogBaseOp -> AstConcreteS $ logBase u v
Atan2Op -> AstConcreteS $ atan2H u v
u v -> case opCode of
DivideOp -> case u of
_ | Just 0 <- unReplC u -> u
_ | Just 1 <- unReplC u -> recip v
_ | Just 1 <- unReplC v -> u
_ | Just 0 <- unReplC v -> u -- the partiality-removal hack
AstTimesS n t
| eqY n v -> t
_ -> Ast.AstR2S DivideOp u v
-- TODO: add other rules that are relatively numerically stable
_ -> Ast.AstR2S opCode u v
astI2S :: (NumScalar r, IntegralH r, Nested.IntElt r, KnownSpan s)
=> OpCodeIntegral2
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
astI2S opCode = \cases
u v | FTKS (snat :$$ _) _ <- ftkAst u
, Just u0 <- unRepl1 u
, Just v0 <- unRepl1 v ->
astReplicateS (snat :$$ ZSS) (astI2S opCode u0 v0)
(Ast.AstPrimalPart n) (Ast.AstPrimalPart k) ->
primalPart (astI2S opCode n k)
(Ast.AstPlainPart @_ @s1 n) (Ast.AstPlainPart @_ @s2 k)
| Just Refl <- testEquality (knownSpan @s1) (knownSpan @s2) ->
plainPart (astI2S opCode n k)
(Ast.AstFromPrimal n) (Ast.AstFromPrimal k) ->
fromPrimal (astI2S opCode n k)
(Ast.AstFromPlain n) (Ast.AstFromPlain k) ->
fromPlain (astI2S opCode n k)
(AstConvUpSFromK n) (AstConvUpSFromK k) ->
sfromK $ astI2K opCode n k
n (AstConvUpSFromK k) ->
sfromK $ astI2K opCode (kfromS n) k
(AstConvUpSFromK n) k ->
sfromK $ astI2K opCode n (kfromS k)
u v -> case opCode of
QuotOp -> case (u, v) of
_ | Just u0 <- unAstS u
, Just v0 <- unAstS v -> fromPlain $ AstConcreteS (quotH u0 v0)
_ | Just 0 <- unReplC u -> u
_ | Just 1 <- unReplC v -> u
_ | Just 0 <- unReplC v -> u -- the partiality-removal hack
(Ast.AstI2S QuotOp u0 v0, _) ->
astI2S QuotOp u0 (astTimesS v0 v)
(AstTimesS n t, _)
| eqY n v -> t
_ -> Ast.AstI2S QuotOp u v
RemOp -> case (u, v) of
_ | Just u0 <- unAstS u
, Just v0 <- unAstS v -> fromPlain $ AstConcreteS (remH u0 v0)
_ | Just 0 <- unReplC u -> u
_ | Just 1 <- unReplC v ->
fromPlain $ AstConcreteS $ defTargetRep $ ftkAst u
_ | Just 0 <- unReplC v -> -- the partiality-removal hack
fromPlain $ AstConcreteS $ defTargetRep $ ftkAst u
_ -> Ast.AstI2S RemOp u v
astConcreteS :: GoodScalar r
=> Concrete (TKS sh r)
-> AstTensor AstMethodLet PlainSpan (TKS sh r)
astConcreteS = AstConcreteS . unConcrete
-- We don't floor concrete tensors ASAP but wait until AstTraverse,
-- because the flooring may fuse or simplify away yet.
astFloorS :: forall r1 r2 sh.
(NumScalar r1, Differentiable r1, NumScalar r2, Integral r2)
=> AstTensor AstMethodLet PlainSpan (TKS sh r1)
-> AstTensor AstMethodLet PlainSpan (TKS sh r2)
astFloorS t = case t of
_ | FTKS (snat :$$ _) _ <- ftkAst t
, Just u <- unRepl1 t ->
astReplicateS (snat :$$ ZSS) (astFloorS u)
Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, v) ->
let !v2 = astFloorS v
in Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, v2)
Ast.AstBuild1 snat STKScalar (var, v) ->
let !v2 = astFloorK v
in Ast.AstBuild1 snat STKScalar (var, v2)
Ast.AstLet var u v -> astLet var u (astFloorS v)
Ast.AstFloorS v -> astFloorS v
Ast.AstFromIntegralS v -> astFromIntegralS v
Ast.AstCastS v -> astFloorS v
Ast.AstScatterS shm shn shp v (vars, ix) ->
astScatterS shm shn shp (astFloorS v) (vars, ix)
Ast.AstGatherS shm shn shp v (vars, ix) ->
astGatherS shm shn shp (astFloorS v) (vars, ix)
Ast.AstIotaS snat -> Ast.AstIotaS snat
Ast.AstReverseS v -> astReverseS (astFloorS v)
Ast.AstTransposeS perm v -> astTransposeS perm (astFloorS v)
Ast.AstReshapeS sh v -> astReshapeS sh (astFloorS v)
AstConvUpSFromK a -> sfromK (astFloorK a)
_ -> Ast.AstFloorS t
astFromIntegralS :: forall r1 r2 sh. (NumScalar r1, NumScalar r2, Integral r1)
=> AstTensor AstMethodLet PlainSpan (TKS sh r1)
-> AstTensor AstMethodLet PlainSpan (TKS sh r2)
astFromIntegralS t = case t of
_ | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) -> t
_ | FTKS (snat :$$ _) _ <- ftkAst t
, Just u <- unRepl1 t ->
astReplicateS (snat :$$ ZSS) (astFromIntegralS u)
-- Ast.AstFromVector snat (STKS sh STKScalar) l ->
-- astFromVector snat (STKS sh STKScalar) (V.map astFromIntegralS l)
-- Ast.AstFromVector snat STKScalar l ->
-- astFromVector snat STKScalar (V.map astFromIntegralK l)
Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, v) ->
let !v2 = astFromIntegralS v
in Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, v2)
Ast.AstBuild1 snat STKScalar (var, v) ->
let !v2 = astFromIntegralK v
in Ast.AstBuild1 snat STKScalar (var, v2)
Ast.AstLet var u v -> astLet var u (astFromIntegralS v)
Ast.AstN1S opCode u -> astN1S opCode (astFromIntegralS u)
Ast.AstFromIntegralS v -> astFromIntegralS v
-- Ast.AstIndexS shn v ix -> astIndexS shn (astFromIntegralS v) ix
-- increases work; also index goes into fromIntegral, so we'd loop
Ast.AstScatterS shm shn shp v (vars, ix) ->
astScatterS shm shn shp (astFromIntegralS v) (vars, ix)
Ast.AstGatherS shm shn shp v (vars, ix) ->
astGatherS shm shn shp (astFromIntegralS v) (vars, ix)
Ast.AstIotaS snat -> Ast.AstIotaS snat
-- Ast.AstSliceS i n k v -> astSliceS i n k (astFromIntegralS v)
Ast.AstReverseS v -> astReverseS (astFromIntegralS v)
Ast.AstTransposeS perm v -> astTransposeS perm (astFromIntegralS v)
Ast.AstReshapeS sh v -> astReshapeS sh (astFromIntegralS v)
AstConvUpSFromK a -> sfromK (astFromIntegralK a)
_ -> Ast.AstFromIntegralS t
astCastS :: forall r1 r2 s sh.
( NumScalar r1, Differentiable r1
, NumScalar r2, Differentiable r2, KnownSpan s )
=> AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
astCastS t = case t of
_ | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) -> t
_ | FTKS (snat :$$ _) _ <- ftkAst t
, Just u <- unRepl1 t ->
astReplicateS (snat :$$ ZSS) (astCastS u)
-- Ast.AstFromVector snat (STKS sh STKScalar) l ->
-- astFromVector snat (STKS sh STKScalar) (V.map astCastS l)
-- Ast.AstFromVector snat STKScalar l ->
-- astFromVector snat STKScalar (V.map astCastK l)
{- This (and other similar rules) is bad, because it changes semantics
and also impacts performance negatively (a is larger than sum a):
Ast.AstSumS shm a -> astSumS shm (astCastS a) -}
Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, v) ->
let !v2 = astCastS v
in Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, v2)
Ast.AstBuild1 snat STKScalar (var, v) ->
let !v2 = astCastK v
in Ast.AstBuild1 snat STKScalar (var, v2)
Ast.AstLet var u v -> astLet var u (astCastS v)
-- These (rarely) loop if ast*Part is used instead of these versions.
Ast.AstPrimalPart a -> primalPart $ astCastS a
Ast.AstDualPart a -> dualPart $ astCastS a
Ast.AstPlainPart a -> plainPart $ astCastS a
Ast.AstFromPrimal v -> fromPrimal $ astCastS v
Ast.AstFromDual v -> fromDual $ astCastS v
Ast.AstFromPlain v -> fromPlain $ astCastS v
Ast.AstN1S opCode u -> astN1S opCode (astCastS u)
Ast.AstR1S opCode u -> astR1S opCode (astCastS u)
Ast.AstFromIntegralS v -> astFromIntegralS v
Ast.AstCastS v -> astCastS v
-- Ast.AstIndexS shn v ix -> astIndexS shn (astCastS v) ix
-- increases work; also index goes into astCastS, so we'd loop
Ast.AstScatterS shm shn shp v (vars, ix) ->
astScatterS shm shn shp (astCastS v) (vars, ix)
Ast.AstGatherS shm shn shp v (vars, ix) ->
astGatherS shm shn shp (astCastS v) (vars, ix)
Ast.AstIotaS snat -> Ast.AstIotaS snat
-- Ast.AstSliceS i n k v -> astSliceS i n k (astCastS v)
Ast.AstReverseS v -> astReverseS (astCastS v)
Ast.AstTransposeS perm v -> astTransposeS perm (astCastS v)
Ast.AstReshapeS sh v -> astReshapeS sh (astCastS v)
AstConvUpSFromK a -> sfromK (astCastK a)
_ -> Ast.AstCastS t
astIndexS
:: forall shm shn s x. KnownSpan s
=> ShS shn
-> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x) -> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKS2 shn x)
astIndexS = astIndexKnobsS defaultKnobs
astIndexKnobsS
:: forall shm shn s x. KnownSpan s
=> SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKS2 shn x)
astIndexKnobsS _ _ v0 ZIS = v0
astIndexKnobsS _ shn v0 (i1 :.$ _)
| Just (lb, ub) <- intBounds i1
, FTKS (snat :$$ _) x <- ftkAst v0
, ub < 0 || lb >= fromSNat' snat =
let ftk = FTKS shn x
in fromPlain $ astConcrete ftk (tdefTarget ftk)
astIndexKnobsS knobs shn u0 (Ast.AstCondK b v w :.$ rest0)
| knobPhase knobs `notElem` [PhaseUnspecified, PhaseVectorization] =
-- don't undo vectorization tweaks
astLetFun u0 $ \u ->
shareIx rest0 $ \rest ->
astCondS b (astIndexKnobsS knobs shn u (v :.$ rest))
(astIndexKnobsS knobs shn u (w :.$ rest))
astIndexKnobsS knobs shn v0 ix@(i1 :.$ rest1)
| FTKS shmshn x <- ftkAst v0
, (:$$) @in1 @shm1 SNat _ <- shsTakeIx @shm @shn Proxy ix shmshn
, let tryRecursing :: AstTensor AstMethodLet s (TKS2 (shm1 ++ shn) x)
-> AstTensor AstMethodLet s (TKS2 shn x)
tryRecursing u =
let ftk = FTKS shn x
defArr = fromPlain $ astConcrete ftk (tdefTarget ftk)
in case 0 <=. i1 &&* i1 <=. valueOf @in1 - 1 of
AstConcreteK b -> if b then astIndex shn u rest1 else defArr
_ -> Ast.AstIndexS shn v0 ix = case v0 of
Ast.AstProject1{} -> Ast.AstIndexS shn v0 ix
Ast.AstProject2{} -> Ast.AstIndexS shn v0 ix
Ast.AstApply{} -> Ast.AstIndexS shn v0 ix
Ast.AstVar{} -> Ast.AstIndexS shn v0 ix
{- This is wrong: in a counterfactual case, astLet assigns OOB i to var2,
violating the invariant about variables bounds:
Ast.AstBuild1 (SNat @k) STKS{} (var2, v) ->
let ftk = FTKS shn x
defArr = fromPlain $ astConcrete ftk (tdefTarget ftk)
in astLetFun i1 $ \i ->
astCondInitial (0 <=. i &&* i <=. valueOf @k - 1)
(astIndex shn (astLet var2 i v) rest1)
defArr -}
Ast.AstBuild1 _ STKScalar (var2, v) | ZIS <- rest1 ->
tryRecursing $ sfromK $ astLet var2 i1 v
Ast.AstBuild1 _ STKS{} (var2, v) ->
tryRecursing $ astLet var2 i1 v
Ast.AstLet var u v -> astLetRefresh var u v $ \v' -> astIndex shn v' ix
Ast.AstPrimalPart{} -> Ast.AstIndexS shn v0 ix -- must be a NF
Ast.AstDualPart{} -> Ast.AstIndexS shn v0 ix
Ast.AstPlainPart{} -> Ast.AstIndexS shn v0 ix
Ast.AstFromPrimal v -> fromPrimal $ astIndex shn v ix
Ast.AstFromDual v -> fromDual $ astIndex shn v ix
Ast.AstFromPlain v -> fromPlain $ astIndex shn v ix
AstPlusS u v ->
shareIx ix $ \ !ix2 -> astPlusS (astIndex shn u ix2)
(astIndex shn v ix2)
AstTimesS u v ->
shareIx ix $ \ !ix2 -> astTimesS (astIndex shn u ix2)
(astIndex shn v ix2)
Ast.AstN1S opCode u -> astN1S opCode (astIndex shn u ix)
Ast.AstR1S opCode u -> astR1S opCode (astIndex shn u ix)
Ast.AstR2S opCode u v ->
shareIx ix $ \ !ix2 -> astR2S opCode (astIndex shn u ix2)
(astIndex shn v ix2)
Ast.AstI2S opCode u v ->
shareIx ix $ \ !ix2 -> astI2S opCode (astIndex shn u ix2)
(astIndex shn v ix2)
AstConcreteS a | AstConcreteK i <- i1 ->
let u = withKnownShS (shsTail shmshn) $
tsindex (Concrete a) (Concrete i :.$ ZIS)
in astIndex shn (astConcreteS u) rest1
AstConcreteS{} -> case unRepl1 v0 of
Just u -> tryRecursing u
_ -> Ast.AstIndexS shn v0 ix
Ast.AstFloorS v -> astFloorS $ astIndex shn v ix
Ast.AstFromIntegralS v -> astFromIntegralS $ astIndex shn v ix
Ast.AstCastS t -> astCastS $ astIndex shn t ix
Ast.AstArgMinS @n1 @shz v | FTKS nsh _ <- ftkAst v
, nl@(SNat @nl) <- shsLast nsh ->
let shnl = shn `shsAppend` (nl :$$ ZSS)
in gcastWith (unsafeCoerceRefl
:: Head (shn ++ '[nl]) ': Tail (shn ++ '[nl])
:~: shn ++ '[nl]) $
gcastWith (unsafeCoerceRefl :: Init (shn ++ '[nl]) :~: shn) $
gcastWith (unsafeCoerceRefl :: shm ++ (shn ++ '[nl]) :~: n1 ': shz) $
Ast.AstArgMinS @(Head (shn ++ '[nl])) @(Tail (shn ++ '[nl]))
$ astIndex shnl v ix
Ast.AstArgMaxS @n1 @shz v | FTKS nsh _ <- ftkAst v
, nl@(SNat @nl) <- shsLast nsh ->
let shnl = shn `shsAppend` (nl :$$ ZSS)
in gcastWith (unsafeCoerceRefl
:: Head (shn ++ '[nl]) ': Tail (shn ++ '[nl])
:~: shn ++ '[nl]) $
gcastWith (unsafeCoerceRefl :: Init (shn ++ '[nl]) :~: shn) $
gcastWith (unsafeCoerceRefl :: shm ++ (shn ++ '[nl]) :~: n1 ': shz) $
Ast.AstArgMaxS @(Head (shn ++ '[nl])) @(Tail (shn ++ '[nl]))
$ astIndex shnl v ix
Ast.AstIndexS @shm4 _ v ix2
| Refl <- lemAppAssoc (Proxy @shm4) (Proxy @shm) (Proxy @shn) ->
astIndex shn v (ix2 `ixsAppend` ix)
Ast.AstCondS b v w ->
shareIx ix $ \ !ix2 -> astCondS b (astIndex shn v ix2)
(astIndex shn w ix2)
Ast.AstFromVectorK (_ :$$ ZSS) l | AstConcreteK i <- i1, ZIS <- rest1 ->
tryRecursing $ sfromK (l V.! i)
Ast.AstFromVectorK{} -> -- normal form
Ast.AstIndexS shn v0 ix
Ast.AstFromVectorS (_ :$$ ZSS) l | AstConcreteK i <- i1 ->
tryRecursing (l V.! i)
Ast.AstFromVectorS @shm2 @shn2 shm2 l | shsLength shm2 < length ix
, SNat <- shsRank shm2 ->
gcastWith (unsafeCoerceRefl :: shn2 :~: Drop (Rank shm2) shm ++ shn) $
gcastWith (unsafeCoerceRefl :: Take (Rank shm2) shm :~: shm2) $
shareIx ix $ \ !ix2 ->
let ixRest2 = ixsDrop @(Rank shm2) ix2
in Ast.AstIndexS shn
(astFromVectorS @shm2 @shn shm2
$ V.map (\a -> astIndex shn a ixRest2) l)
(ixsTake @(Rank shm2) ix2)
Ast.AstFromVectorS{} -> -- normal form (shsLength shm2 >= length ix)
Ast.AstIndexS shn v0 ix
Ast.AstSumS @shm2 shm2 v ->
let perm1 = backpermCycleN (shsLength shm2) (shsLength shm2 + length ix)
in Permutation.permFromListCont perm1 $ \(perm :: Permutation.Perm perm) ->
gcastWith (unsafeCoerceRefl
:: (Rank perm <=? Rank (shm2 ++ (shm ++ shn))) :~: True) $
gcastWith (unsafeCoerceRefl
:: Permutation.PermutePrefix perm (shm2 ++ (shm ++ shn))
:~: shm ++ (shm2 ++ shn)) $
fromMaybe (error "astIndexKnobsS: impossible non-permutation")
$ Permutation.permCheckPermutation perm
$ astSumS shm2
$ astIndex (shm2 `shsAppend` shn) (astTransposeS perm v) ix
Ast.AstScatterS shm7 shn7 shp7 v (vars, i5 :.$ ix2) | eqY i5 i1 ->
astIndex shn (astScatter shm7 shn7 (shsTail shp7) v (vars, ix2)) rest1
Ast.AstScatterS _ _ _ _ (_, AstConcreteK{} :.$ _)
| AstConcreteK{} <- i1 -> -- from above we know i5 /= i1
let ftk = FTKS shn x
in fromPlain $ astConcrete ftk (tdefTarget ftk)
Ast.AstScatterS{} -> -- normal form
Ast.AstIndexS shn v0 ix
Ast.AstReplicateK (_ :$$ shmRest) v ->
tryRecursing $ astReplicateK shmRest v
Ast.AstReplicateS ZSS v -> astIndex shn v ix
Ast.AstReplicateS (_ :$$ shmRest) v ->
tryRecursing $ astReplicateS shmRest v
-- This is not a possible normal form, but pattern needs to be exhaustive.
Ast.AstGatherS @_ @_ @shp' _ _ _ v (ZS, ix2)
| Refl <- lemAppAssoc (Proxy @shp') (Proxy @(in1 ': shm1)) (Proxy @shn) ->
astIndex @(shp' ++ shm) @shn shn v (ix2 `ixsAppend` ix)
Ast.AstGatherS @_ @shn' (_ :$$ (shm71 :: ShS shm71)) shn' shp'
v (var2 ::$ vars, ix2) ->
gcastWith (unsafeCoerceRefl :: shm71 ++ shn' :~: shm1 ++ shn) $
let ftk = FTKS shn x
defArr = fromPlain $ astConcrete ftk (tdefTarget ftk)
w :: AstTensor AstMethodLet s (TKS2 (shm1 ++ shn) x)
w = astGather shm71 shn' shp' v (vars, ix2)
u = astLet var2 i1 $ astIndex @shm1 @shn shn w rest1
-- This let makes it impossible to use astCond when i1 is OOB.
-- No astLetRefresh is needed, because the smart constructor
-- for AstGatherS ensures var2 doesn't occur in v, even bound.
in case 0 <=. i1 &&* i1 <=. valueOf @in1 - 1 of
AstConcreteK b -> if b then u else defArr
_ -> Ast.AstIndexS shn v0 ix
Ast.AstIotaS _ -> case shn of
ZSS ->
let ftk = FTKS ZSS x
defArr = fromPlain $ astConcrete ftk (tdefTarget ftk)
in astLetFun i1 $ \i ->
astCondSInitial (0 <=. i &&* i <=. valueOf @in1 - 1)
(astFromIntegralS $ sfromK i)
defArr
_ -> error "astIndexKnobsS: shn not []"
Ast.AstAppendS u v | FTKS (SNat @m :$$ _) _ <- ftkAst u ->
if knobPhase knobs == PhaseExpansion then
astLetFun i1 $ \i ->
shareIx rest1 $ \ !rest2 ->
let ulen = valueOf @m
ix1 = i :.$ rest2
ix2 = i - ulen :.$ rest2
in case ulen <=. i of
AstConcreteK b -> if b then astIndex shn v ix2 else astIndex shn u ix1
bExpr ->
-- This results in a larger term, so we consider this late.
astCondSInitial bExpr (astIndex shn v ix2) (astIndex shn u ix1)
else
-- We don't need to share i1 and rest1 here and, in fact, we can't
-- because we'd loop via building AstLet around the resulting AstIndexS.
let ulen = valueOf @m
ix1 = i1 :.$ rest1
ix2 = i1 - ulen :.$ rest1
in case ulen <=. i1 of
AstConcreteK b -> if b then astIndex shn v ix2 else astIndex shn u ix1
_ -> Ast.AstIndexS shn v0 ix
Ast.AstSliceS i@(SNat @i) (SNat @n) k@SNat v ->
astLetFun i1 $ \iShared ->
let ftk = FTKS shn x
defArr = fromPlain $ astConcrete ftk (tdefTarget ftk)
b = (if fromSNat' i == 0 then true else 0 <=. iShared)
&&* (if fromSNat' k == 0 then true else iShared <=. valueOf @n - 1)
ii = valueOf @i + iShared
in astCondSInitial b (astIndex shn v (ii :.$ rest1)) defArr
Ast.AstReverseS v ->
let iRev = valueOf @in1 - 1 - i1
in astIndex shn v (iRev :.$ rest1)
Ast.AstTransposeS @_ @sh2 perm v
| fromSNat' (Permutation.permRank perm) <= length ix ->
-- TODO: remake once there's an S version of permInverse:
permInverse perm $ \(invperm :: Permutation.Perm invperm) _ ->
let ix2 :: AstIxS AstMethodLet (Permutation.PermutePrefix invperm shm)
ix2 = ixsPermutePrefix invperm ix
in gcastWith (unsafeCoerceRefl
:: sh2 :~: Permutation.PermutePrefix invperm shm ++ shn) $
astIndex shn v ix2
-- This results in a larger term, so we consider this late.
Ast.AstTransposeS @perm perm v
| knobPhase knobs `elem` [PhaseVectorization, PhaseExpansion] ->
astIndex shn (astTransposeAsGatherS @perm (deVect knobs) perm v) ix
Ast.AstTransposeS{} -> Ast.AstIndexS shn v0 ix
-- This results in a larger term, so we only permit it for the flat case.
Ast.AstReshapeS sh v | shsLength sh <= 1 -> -- this is flatten
astIndex shn (astReshapeAsGatherS (deVect knobs) sh v) ix
Ast.AstReshapeS{} -> Ast.AstIndexS shn v0 ix
-- Conversions to shaped need to stay down, so this is NF.
Ast.AstConvert{} -> Ast.AstIndexS shn v0 ix
Ast.AstDot1InS @_ @n _ n u v
| Refl <- lemAppAssoc (Proxy @shm) (Proxy @shn) (Proxy @'[n]) ->
astDot1InS shn n (astIndex (shn `shsAppend` (n :$$ ZSS)) u ix)
(astIndex (shn `shsAppend` (n :$$ ZSS)) v ix)
Ast.AstMatmul2S{} -> Ast.AstIndexS shn v0 ix -- TODO
Ast.AstBoolNotS a -> Ast.AstBoolNotS $ astIndex shn a ix
Ast.AstBoolAndS b1 b2 -> Ast.AstBoolAndS (astIndex shn b1 ix)
(astIndex shn b2 ix)
Ast.AstLeqS @_ @sh _shb sh arg1 arg2
| Refl <- lemAppAssoc (Proxy @shm) (Proxy @shn) (Proxy @sh) ->
Ast.AstLeqS shn sh (astIndex (shn `shsAppend` sh) arg1 ix)
(astIndex (shn `shsAppend` sh) arg2 ix)
where
astIndex
:: forall shm' shn' s' x'. KnownSpan s'
=> ShS shn'
-> AstTensor AstMethodLet s' (TKS2 (shm' ++ shn') x')
-> AstIxS AstMethodLet shm'
-> AstTensor AstMethodLet s' (TKS2 shn' x')
astIndex = astIndexKnobsS (deVect knobs)
astGather
:: forall shm' shn' shp' x'.
ShS shm' -> ShS shn' -> ShS shp'
-> AstTensor AstMethodLet s (TKS2 (shp' ++ shn') x')
-> (AstVarListS shm', AstIxS AstMethodLet shp')
-> AstTensor AstMethodLet s (TKS2 (shm' ++ shn') x')
astGather = astGatherKnobsS (deVect knobs)
astScatter
:: forall shm' shn' shp' x'. TKAllNum x'
=> ShS shm' -> ShS shn' -> ShS shp'
-> AstTensor AstMethodLet s (TKS2 (shm' ++ shn') x')
-> (AstVarListS shm', AstIxS AstMethodLet shp')
-> AstTensor AstMethodLet s (TKS2 (shp' ++ shn') x')
astScatter = astScatterKnobsS (deVect knobs)
-- TODO: compared to tletIx, it adds many lets, not one, but does not
-- create other (and non-simplified!) big terms and also uses astIsSmall,
-- so it's probably more efficient. Use this instead of tletIx
-- or design something even better.
--
-- See https://github.com/Mikolaj/horde-ad/issues/119.
shareIx :: forall s shm y. KnownSpan s
=> AstIxS AstMethodLet shm
-> (AstIxS AstMethodLet shm -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
{-# NOINLINE shareIx #-}
shareIx ix f = unsafePerformIO $ do
let shareI :: AstInt AstMethodLet
-> IO ( Maybe (IntVarName, AstInt AstMethodLet)
, AstInt AstMethodLet )
shareI i | astIsSmall True i = return (Nothing, i)
shareI i =
-- i can be OOB, so we can't use shape to determine its bounds
funToAstIntMaybeIO (intBounds i) $ \ (!varFresh, !astVarFresh) ->
(Just (varFresh, i), astVarFresh)
(bindings, ix2) <- mapAndUnzipM shareI (Foldable.toList ix)
return $! foldl' (\v (var, u) -> astLet var u v)
(f $ ixsFromIxS ix ix2)
(catMaybes bindings)
astCondKInitial :: forall r s. (GoodScalar r, KnownSpan s)
=> AstBool AstMethodLet
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
astCondKInitial b v w = case b of
AstLeqInt (AstConcreteK j) (AstIntVar var)
| Just (lb, ub) <- varNameToBounds var ->
let varTrue = reboundsVarName (max lb j, ub) var
varFalse = reboundsVarName (lb, min ub (j - 1)) var
in astCondK b (substituteAst (astVar varTrue) var v)
(substituteAst (astVar varFalse) var w)
_ -> astCondK b v w
astCondK :: forall r s. (GoodScalar r, KnownSpan s)
=> AstBool AstMethodLet
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
astCondK b v w | Just b0 <- unAstK b = if b0 then v else w
astCondK _b u v | eqY u v = u
astCondK (Ast.AstBoolNotK b) v w = astCondK b w v
astCondK (Ast.AstLet var n b) u v =
astLetRefresh var n b $ \b' -> astCondKInitial b' u v
astCondK b (Ast.AstLet var n u) v =
astLetRefresh var n u $ \u' -> astCondK b u' v
astCondK b u (Ast.AstLet var n v) =
astLetRefresh var n v $ \v' -> astCondK b u v'
astCondK b (Ast.AstFromPrimal v) (Ast.AstFromPrimal w) =
fromPrimal $ astCondK b v w
astCondK b (Ast.AstFromDual v) (Ast.AstFromDual w) =
fromDual $ astCondK b v w
astCondK b (Ast.AstFromPlain v) (Ast.AstFromPlain w) =
fromPlain $ astCondK b v w
astCondK b (AstPlusK u1 u2) (AstPlusK v1 v2) | eqY u1 v1 =
astPlusK u1 (astCondK b u2 v2)
astCondK b (AstPlusK u1 u2) v1 | eqY u1 v1 =
astPlusK u1 (astCondK b u2 0)
astCondK b u1 (AstPlusK v1 v2) | eqY u1 v1 =
astPlusK u1 (astCondK b 0 v2)
astCondK b (AstPlusK u1 u2) (AstPlusK v1 v2) | eqY u2 v2 =
astPlusK (astCondK b u1 v1) u2
astCondK b (AstPlusK u1 u2) v2 | eqY u2 v2 =
astPlusK (astCondK b u1 0) u2
astCondK b u2 (AstPlusK v1 v2) | eqY u2 v2 =
astPlusK (astCondK b 0 v1) u2
astCondK b (AstTimesK u1 u2) (AstTimesK v1 v2) | eqY u1 v1 =
astTimesK u1 (astCondK b u2 v2)
astCondK b (AstTimesK u1 u2) v1 | eqY u1 v1 =
astTimesK u1 (astCondK b u2 0)
astCondK b u1 (AstTimesK v1 v2) | eqY u1 v1 =
astTimesK u1 (astCondK b 0 v2)
astCondK b (AstTimesK u1 u2) (AstTimesK v1 v2) | eqY u2 v2 =
astTimesK (astCondK b u1 v1) u2
astCondK b u2 (AstTimesK v1 v2) | eqY u2 v2 =
astTimesK (astCondK b 0 v1) u2
astCondK b (AstTimesK u1 u2) v2 | eqY u2 v2 =
astTimesK (astCondK b u1 0) u2
-- These two normal forms prevent simple cases of an exponential blow-up
-- when substituting into nested conditionals in the rule that follows them.
-- For now they don't seem to be needed, though.
-- astCondK b v@Ast.AstCond{} w = Ast.AstCond b v w
-- astCondK b v w@Ast.AstCond{} = Ast.AstCond b v w
astCondK (AstLeqInt (AstConcreteK j) var0@(AstIntVar var)) v w
| Just Refl <- testEquality (typeRep @r) (typeRep @Int)
, Just (lb, ub) <- varNameToBounds var
, let varTrue = reboundsVarName (max lb j, ub) var
varFalse = reboundsVarName (lb, min ub (j - 1)) var
, AstConcreteK dv <- v - fromPlain (astVar varTrue)
, AstConcreteK dw <- w - fromPlain (astVar varFalse)
, dv == dw =
fromPlain $ AstConcreteK dv + var0
astCondK b@(AstLeqInt AstConcreteK{}
(Ast.AstN1K NegateOp AstIntVar{})) v w =
astCondKInitial (astBoolNotK b) w v -- eliminate the NegateOp
{- TODO: Investigate how much stronger this rule is than the above (it's slower)
astCondK (AstLeqInt (AstConcreteK k) var0@(AstIntVar var)) v w
| FTKScalar @r <- ftkAst v
, Just Refl <- testEquality (typeRep @r) (typeRep @Int)
, Just (lb, ub) <- varNameToBounds var
, let varTrue = reboundsVarName (max lb k, ub) var
varFalse = reboundsVarName (lb, min ub (k - 1)) var
d = v - fromPlain (astVar varTrue)
, eqY w (substituteAst (astVar varFalse) varTrue d
+ fromPlain (astVar varFalse)) =
substituteAst var0 varTrue d + fromPlain var0 -}
{- Disabled, because scatters and gathers simplify better with conditionals
than with quotH, so the rules need to be restricted at least.
-- TODO: prove nothing wrong comes out of dividing by zero here:
astCondK (AstLeqInt (AstConcreteK k) var0@(AstIntVar var)) v w
| FTKScalar @r <- ftkAst v
, Just Refl <- testEquality (typeRep @r) (typeRep @Int)
, Just (lb, ub) <- varNameToBounds var
, let varTrue = reboundsVarName (max lb k, ub) var
varFalse = reboundsVarName (lb, min ub (k - 1)) var
d = fromPlain (AstConcreteK k) -- a simplified guess
`quotH` substituteAst (AstConcreteK k) var v
, eqY (fromPlain (astVar varTrue) `quotH` d)
(substituteAst (astVar varTrue) var v)
, eqY (fromPlain (astVar varFalse) `quotH` d)
(substituteAst (astVar varFalse) var w) =
fromPlain var0 `quotH` d
astCondK (AstLeqInt (AstConcreteK k) var0@(AstIntVar var)) v w
| FTKScalar @r <- ftkAst v
, Just Refl <- testEquality (typeRep @r) (typeRep @Int)
, Just (lb, ub) <- varNameToBounds var
, let varTrue = reboundsVarName (max lb k, ub) var
varFalse = reboundsVarName (lb, min ub (k - 1)) var
d = fromPlain (AstConcreteK (k - 1)) -- the only change vs above
`quotH` substituteAst (AstConcreteK (k - 1)) var w
, eqY (fromPlain (astVar varTrue) `quotH` d)
(substituteAst (astVar varTrue) var v)
, eqY (fromPlain (astVar varFalse) `quotH` d)
(substituteAst (astVar varFalse) var w) =
fromPlain var0 `quotH` d -}
astCondK b v w = Ast.AstCondK b v w
astCondSInitial :: forall sh x s. KnownSpan s
=> AstBool AstMethodLet
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astCondSInitial b v w = case b of
AstLeqInt (AstConcreteK j) (AstIntVar var)
| Just (lb, ub) <- varNameToBounds var ->
let varTrue = reboundsVarName (max lb j, ub) var
varFalse = reboundsVarName (lb, min ub (j - 1)) var
in astCondS b (substituteAst (astVar varTrue) var v)
(substituteAst (astVar varFalse) var w)
_ -> astCondS b v w
astCondS :: forall sh x s. KnownSpan s
=> AstBool AstMethodLet
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astCondS b v w | Just b0 <- unAstK b = if b0 then v else w
astCondS _b u v | eqY u v = u
astCondS b v w | FTKS (snat :$$ _) _ <- ftkAst v
, Just v1 <- unRepl1 v
, Just w1 <- unRepl1 w =
astReplicateS (snat :$$ ZSS) (astCondS b v1 w1)
astCondS (Ast.AstBoolNotK b) v w = astCondS b w v
astCondS (Ast.AstLet var n b) u v =
astLetRefresh var n b $ \b' -> astCondS b' u v
astCondS b (Ast.AstLet var n u) v =
astLetRefresh var n u $ \u' -> astCondS b u' v
astCondS b u (Ast.AstLet var n v) =
astLetRefresh var n v $ \v' -> astCondS b u v'
astCondS b (Ast.AstFromPrimal v) (Ast.AstFromPrimal w) =
fromPrimal $ astCondS b v w
astCondS b (Ast.AstFromDual v) (Ast.AstFromDual w) =
fromDual $ astCondS b v w
astCondS b (Ast.AstFromPlain v) (Ast.AstFromPlain w) =
fromPlain $ astCondS b v w
astCondS _b (Ast.AstIotaS n) Ast.AstIotaS{} = Ast.AstIotaS n
astCondS b v w | FTKS ZSS FTKScalar <- ftkAst v =
sfromK (astCondK b (kfromS v) (kfromS w))
astCondS b v w = Ast.AstCondS b v w
astFromVectorK :: forall shm r s. (GoodScalar r, KnownSpan s)
=> ShS shm
-> Data.Vector.Vector (AstTensor AstMethodLet s (TKScalar r))
-> AstTensor AstMethodLet s (TKS shm r)
astFromVectorK ZSS l = assert (V.length l == 1) $ sfromK $ l V.! 0
astFromVectorK _ l | V.null l = error "astFromVectorK: empty vector"
astFromVectorK shm l | SNat' @1 <- shsProduct shm =
astReplicateK shm (l V.! 0)
astFromVectorK shm l = fromMaybe (Ast.AstFromVectorK shm l) $
-- This disables some rules, e.g., indexing or summing of fromVector
-- of concrete arrays, but allocating an extra array of the same size
-- as the fromVector is not a big deal and early rules are better
-- then the same rules in contraction phase.
(let unConc :: AstTensor AstMethodLet s (TKScalar r)
-> Maybe (Concrete (TKScalar r))
unConc t = Concrete <$> unAstK t
in case V.mapM unConc l of
Just l4 -> Just $ fromPlain $ astConcreteS (tsfromVectorLinear shm l4)
Nothing -> Nothing)
`mplus`
(let unFromPrimal :: AstTensor AstMethodLet s2 y
-> Maybe (AstTensor AstMethodLet (PrimalStepSpan s2) y)
unFromPrimal (Ast.AstFromPrimal t) = Just t
unFromPrimal _ = Nothing
in case V.mapM unFromPrimal l of
Just l2 -> Just $ fromPrimal $ astFromVectorK shm l2
Nothing -> Nothing)
`mplus`
(case knownSpan @s of
SFullSpan ->
let unFromDual :: AstTensor AstMethodLet FullSpan y
-> Maybe (AstTensor AstMethodLet DualSpan y)
unFromDual (Ast.AstFromDual t) = Just t
unFromDual _ = Nothing
in case V.mapM unFromDual l of
Just l2 -> Just $ fromDual $ astFromVectorK shm l2
Nothing -> Nothing
_ -> Nothing)
`mplus`
(let unFromPlain :: AstTensor AstMethodLet s2 y
-> Maybe (AstTensor AstMethodLet PlainSpan y)
unFromPlain (Ast.AstFromPlain t) = Just t
unFromPlain _ = Nothing
in case V.mapM unFromPlain l of
Just l2 -> Just $ fromPlain $ astFromVectorK shm l2
Nothing -> Nothing)
astFromVectorS :: forall shm shn x s. KnownSpan s
=> ShS shm
-> Data.Vector.Vector (AstTensor AstMethodLet s (TKS2 shn x))
-> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)
astFromVectorS ZSS l = assert (V.length l == 1) $ l V.! 0
astFromVectorS _ l | V.null l = error "astFromVectorS: empty vector"
astFromVectorS shm l | SNat' @1 <- shsProduct shm =
astReplicateS shm (l V.! 0)
astFromVectorS shm l = fromMaybe (Ast.AstFromVectorS shm l) $
(let unFromPrimal :: AstTensor AstMethodLet s2 y
-> Maybe (AstTensor AstMethodLet (PrimalStepSpan s2) y)
unFromPrimal (Ast.AstFromPrimal t) = Just t
unFromPrimal _ = Nothing
in case V.mapM unFromPrimal l of
Just l2 -> Just $ fromPrimal $ astFromVectorS shm l2
Nothing -> Nothing)
`mplus`
(case knownSpan @s of
SFullSpan ->
let unFromDual :: AstTensor AstMethodLet FullSpan y
-> Maybe (AstTensor AstMethodLet DualSpan y)
unFromDual (Ast.AstFromDual t) = Just t
unFromDual _ = Nothing
in case V.mapM unFromDual l of
Just l2 -> Just $ fromDual $ astFromVectorS shm l2
Nothing -> Nothing
_ -> Nothing)
`mplus`
(let unFromPlain :: AstTensor AstMethodLet s2 y
-> Maybe (AstTensor AstMethodLet PlainSpan y)
unFromPlain (Ast.AstFromPlain t) = Just t
unFromPlain _ = Nothing
in case V.mapM unFromPlain l of
Just l2 -> Just $ fromPlain $ astFromVectorS shm l2
Nothing -> Nothing)
`mplus`
(case ftkAst (l V.! 0) of
FTKS ZSS FTKScalar | Refl <- lemAppNil @shm ->
Just $ astFromVectorK shm $ V.map kfromS l
FTKS shn FTKScalar ->
-- This disables some rules, e.g., indexing or summing of fromVector
-- of concrete arrays, but allocating an extra array of the same size
-- as the fromVector is not a big deal and early rules are better
-- then the same rules in contraction phase.
let unConc :: AstTensor AstMethodLet s (TKS shn r)
-> Maybe (Concrete (TKS shn r))
unConc t = Concrete <$> unAstS t
in case V.mapM unConc l of
Just l4 ->
withKnownShS shn $
Just $ fromPlain $ astConcreteS (tsfromVectorN shm l4)
Nothing -> Nothing
FTKS{} -> Nothing)
astSumK :: forall shm r s. (NumScalar r, KnownSpan s)
=> AstTensor AstMethodLet s (TKS shm r)
-> AstTensor AstMethodLet s (TKScalar r)
astSumK t0 | FTKS shm FTKScalar <- ftkAst t0 = case t0 of
_ | SZ <- shsProduct shm -> 0
_ | SNat' @1 <- shsProduct shm ->
astIndexK t0 (AstConcreteK <$> ixsZero shm)
Ast.AstLet var u v -> astLet var u (astSumK v)
Ast.AstFromPrimal u -> fromPrimal $ astSumK u
Ast.AstFromDual u -> fromDual $ astSumK u
Ast.AstFromPlain u -> fromPlain $ astSumK u
Ast.AstN1S NegateOp u -> negate $ astSumK u
AstTimesS t1 t2
| Just u <- unRepl t1 ->
astTimesK (kfromS u) (astSumK t2)
AstConcreteS v ->
withKnownShS (Nested.sshape v) $
astConcreteK $ tssum0 (Concrete v)
Ast.AstFromVectorK _ l | V.length l < 100 ->
foldr1 astPlusK l
Ast.AstFromVectorS _ l | V.length l < 100 ->
astSumK $ foldr1 astPlusS l
Ast.AstSumS _ u -> astSumK u
Ast.AstScatterS shm1@(k2 :$$ _) shn1 shp1 v (vars, i1 :.$ rest)
| -- This boolean term may have free variables that act as universally
-- quantified.
AstConcreteK True <- 0 <=. i1
&&* i1 <=. AstConcreteK (fromSNat' k2 - 1) ->
astSumK $ astScatterS shm1 shn1 (shsTail shp1) v (vars, rest)
Ast.AstReplicateK _ u ->
astTimesK u
(fromPlain $ AstConcreteK $ fromIntegral $ shsSize shm)
Ast.AstReplicateS shm2 u ->
astTimesK (astSumK u)
(fromPlain $ AstConcreteK $ fromIntegral $ shsSize shm2)
_ | ZSS <- shm -> kfromS t0
_ | snat :$$ _ <- shm
, Just u <- unRepl1 t0 ->
astTimesK (astSumK u)
(fromPlain $ AstConcreteK $ fromIntegral $ fromSNat' snat)
-- TODO: if ever needed, we may also add case
-- Ast.AstCond AstReplicate AstReplicate
-- or use unRepl and unRepl1 to handle it.
Ast.AstIotaS n ->
AstConcreteK $ fromIntegral $ fromSNat' n * (fromSNat' n - 1) `div` 2
Ast.AstAppendS u v -> astSumK u `astPlusK` astSumK v
Ast.AstReverseS u -> astSumK u
Ast.AstTransposeS _ u -> astSumK u
Ast.AstReshapeS _ u -> astSumK u
_ -> Ast.AstSumK t0
astSumS :: forall shm shn x s. (KnownSpan s, TKAllNum x)
=> ShS shm -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)
-> AstTensor AstMethodLet s (TKS2 shn x)
astSumS ZSS t0 = t0
astSumS shm@(snat :$$ _) t0 | FTKS shmshn x <- ftkAst t0
, SNat <- shsRank shm
, let shn = shsDrop @(Rank shm) shmshn =
gcastWith (unsafeCoerceRefl
:: Drop (Rank shm) (shm ++ shn)
:~: shn) $
case t0 of
_ | SZ <- shsProduct shmshn ->
let ftk = FTKS shn x
in fromPlain $ astConcrete ftk (tdefTarget ftk)
_ | SNat' @1 <- shsProduct shm ->
astIndexS shn t0 (AstConcreteK <$> ixsZero shm)
-- astReshape slows down the CNNO test
_ | ZSS <- shn
, FTKScalar @r <- x
, Dict0 <- numFromTKAllNum (Proxy @r) ->
sfromK $ astSumK t0
-- This keeps tensors alive for longer, but it enables new simplifications,
-- while hiding a sum inside let not often prevents other simplifications,
-- because there are few redexes with sum but not at the top.
Ast.AstLet var u v -> astLet var u (astSumS shm v)
Ast.AstFromPrimal v -> fromPrimal $ astSumS shm v
Ast.AstFromDual v -> fromDual $ astSumS shm v
Ast.AstFromPlain v -> fromPlain $ astSumS shm v
Ast.AstN1S NegateOp u -> negate $ astSumS shm u
-- This exchanges a multiplication at rank n+m for one
-- at rank n, which should be faster. We choose t1, since it's likely to be
-- concrete and so it's easier to see if it's replicated and also t2 is
-- likely to be AstTimes, to be processed recusively. This rule may distrupt
-- recognizing matmul, sdot1In, etc., but it sometimes leads to a smaller
-- and better sdot1In.
AstTimesS t1 t2
| Just u <- unReplN shm t1
, FTKScalar @r <- x
, Dict0 <- numFromTKAllNum (Proxy @r) ->
astTimesS u (astSumS shm t2)
AstConcreteS t ->
withKnownShS shm $
withKnownShS shn $
astConcreteS $ tssumN @_ @shm $ Concrete t
Ast.AstFromVectorS (_ :$$ ZSS) l
| V.length l < 100
, FTKScalar @r <- x
, Dict0 <- numFromTKAllNum (Proxy @r) ->
astSumS (shsTail shm) $ foldr1 astPlusS l
Ast.AstSumS @shm2 shm2 u
| Refl <- lemAppAssoc (Proxy @shm2) (Proxy @shm) (Proxy @shn) ->
astSumS (shm2 `shsAppend` shm) u
Ast.AstScatterS shm1@(k2 :$$ _) shn1 shp1 v (vars, i1 :.$ rest)
| -- This boolean term may have free variables that act as universally
-- quantified.
AstConcreteK True <- 0 <=. i1
&&* i1 <=. AstConcreteK (fromSNat' k2 - 1) ->
astSumS (shsTail shm)
$ astScatterS shm1 shn1 (shsTail shp1) v (vars, rest)
Ast.AstReplicateK _ u | FTKScalar @r <- x
, Dict0 <- numFromTKAllNum (Proxy @r) ->
astReplicateK shn
$ astTimesK u (fromPlain $ AstConcreteK $ fromIntegral $ shsSize shm)
_ | Just u <- unRepl1 t0
, FTKScalar @r <- x
, Dict0 <- numFromTKAllNum (Proxy @r) ->
astTimesS (astSumS (shsTail shm) u)
(astReplicateK shn $ fromPlain $ AstConcreteK $ fromIntegral
$ fromSNat' snat)
Ast.AstAppendS u v
| FTKScalar @r <- x
, Dict0 <- numFromTKAllNum (Proxy @r)
, FTKS (m :$$ _) _ <- ftkAst u
, FTKS (n :$$ _) _ <- ftkAst v ->
astSumS (m :$$ shsTail shm) u `astPlusS` astSumS (n :$$ shsTail shm) v
Ast.AstReverseS v -> astSumS shm v
-- See the analogous astSliceS rule.
Ast.AstTransposeS perm t
| SNat <- Permutation.permRank perm
, LTI <- cmpNat (Permutation.permRank perm) (shsRank shm)
, FTKS @shkshn shkshn _ <- ftkAst t ->
gcastWith (unsafeCoerceRefl :: Drop (Rank shm) shkshn :~: shn) $
gcastWith (unsafeCoerceRefl ::
Take (Rank shm) shkshn ++ Drop (Rank shm) shkshn :~: shkshn) $
astSumS (shsTake @(Rank shm) shkshn) t
Ast.AstTransposeS perm t
| SNat <- Permutation.permRank perm
, EQI <- cmpNat (Permutation.permRank perm) (shsRank shm)
, FTKS @shkshn shkshn _ <- ftkAst t ->
gcastWith (unsafeCoerceRefl :: Drop (Rank shm) shkshn :~: shn) $
gcastWith (unsafeCoerceRefl ::
Take (Rank shm) shkshn ++ Drop (Rank shm) shkshn :~: shkshn) $
astSumS (shsTake @(Rank shm) shkshn) t
_ -> Ast.AstSumS shm t0
-- TODO: fuse scatters, scatter and sum, and perhaps more (fromList?)
astScatterS :: forall shm shn shp x s. (KnownSpan s, TKAllNum x)
=> ShS shm -> ShS shn -> ShS shp
-> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 (shp ++ shn) x)
astScatterS = astScatterKnobsS defaultKnobs
astScatterKnobsS :: forall shm shn shp x s. (KnownSpan s, TKAllNum x)
=> SimplifyKnobs
-> ShS shm -> ShS shn -> ShS shp
-> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 (shp ++ shn) x)
astScatterKnobsS _ _ _ _ v0 (!vars0, !_ix0)
| Foldable.any (`varNameInAst` v0) vars0 =
error $ "astScatterKnobsS: scatter vars in v0: " ++ show (vars0, v0)
astScatterKnobsS _ shm _ _ v0 (_, ZIS) = astSumS shm v0
astScatterKnobsS _ _ shn shp@(SZ :$$ _) v0 _
| FTKS _ x <- ftkAst v0 =
let ftk = FTKS (shp `shsAppend` shn) x
in fromPlain $ astConcrete ftk (tdefTarget ftk)
astScatterKnobsS _ _ shn shp@(k :$$ _) v0 (_, i1 :.$ _)
| Just (lb, ub) <- intBounds i1
, let FTKS _ x = ftkAst v0
, ub < 0 || lb >= fromSNat' k =
let ftk = FTKS (shp `shsAppend` shn) x
in fromPlain $ astConcrete ftk (tdefTarget ftk)
astScatterKnobsS knobs shm shn shp v0 (vars0@(_ ::$ _), ix0@(_ :.$ _))
| let ixInit = ixsInit ix0
varInit = AstVarListS $ ixsInit $ unAstVarListS vars0
varLast = ixsLast $ unAstVarListS vars0
, AstIntVar ixvarLast <- ixsLast ix0
, ixvarLast == varLast
, not (varLast `varNameInIxS` ixInit)
, kLast@SNat <- shsLast shm
, Just Refl <- testEquality kLast (shsLast shp) =
gcastWith (unsafeCoerceRefl
:: Init shp ++ (Last shm ': shn) :~: shp ++ shn) $
gcastWith (unsafeCoerceRefl
:: Init shm ++ (Last shm ': shn) :~: shm ++ shn) $
astScatterKnobsS knobs (shsInit shm) (kLast :$$ shn) (shsInit shp)
v0 (varInit, ixInit)
astScatterKnobsS knobs shm@(m :$$ _) shn (p@(SNat @p) :$$ ZSS) v0
( varm ::$ mrest
, Ast.AstI2K QuotOp (AstIntVar varm') (AstConcreteK i0) :.$ ZIS )
| varm == varm'
, i0 > 0 -- ensured by other rules; makes this easier to reason about
, FTKS _ x <- ftkAst v0
, let k0 = (fromSNat' m + i0 - 1) `quot` i0
, fromSNat' p >= k0 = -- most likely, because not OOB
withSNat i0 $ \i ->
withSNat k0 $ \k@(SNat @k) ->
withSNat (fromSNat' p - k0) $ \z2@(SNat @z2) ->
let ftk = FTKS (z2 :$$ shn) x
in gcastWith (unsafeCoerceRefl :: k + z2 :~: p) $
(`astAppendS` fromPlain (astConcrete ftk (tdefTarget ftk)))
$ astSumS (i :$$ ZSS)
$ astReshapeS (i :$$ k :$$ shn)
-- This looks like we are creating a larger outcome tensor than before,
-- but the following reduces to a nested sum of v0 with some transposes
-- and appends, so we likely create something much smaller than v0
-- despite the outermost dimension not shrinking as it did before.
$ astScatterKnobsS knobs shm shn (snatMul k i :$$ ZSS)
v0 (varm ::$ mrest, astVar varm :.$ ZIS)
-- this gather may still index out of bounds, which is fine
astScatterKnobsS knobs shm shn shp@(SNat' @1 :$$ _)
v (vars, AstConcreteK _ :.$ rest) = -- if OOB, covered above
astReplicateS (SNat @1 :$$ ZSS)
$ astScatterKnobsS knobs shm shn (shsTail shp) v (vars, rest)
astScatterKnobsS knobs shm shn shp (Ast.AstLet var u v) (vars, ix) =
astLetRefresh var u v
$ \v' -> astScatterKnobsS knobs shm shn shp v' (vars, ix)
astScatterKnobsS knobs shm shn shp (Ast.AstFromPrimal v) (vars, ix) =
fromPrimal $ astScatterKnobsS knobs shm shn shp v (vars, ix)
astScatterKnobsS knobs shm shn shp (Ast.AstFromDual v) (vars, ix) =
fromDual $ astScatterKnobsS knobs shm shn shp v (vars, ix)
astScatterKnobsS knobs shm shn shp (Ast.AstFromPlain v) (vars, ix) =
fromPlain $ astScatterKnobsS knobs shm shn shp v (vars, ix)
astScatterKnobsS knobs ZSS _shn shp
(Ast.AstScatterS @_ @shn2 @shp2 shm2 shn2 shp2 v (vars, ix2))
(ZS, ix) = -- oneHot (scatter) fusion
gcastWith (unsafeCoerceRefl :: shp ++ shn :~: shp ++ shp2 ++ shn2) $
astScatterKnobsS knobs shm2 shn2 (shp `shsAppend` shp2)
v (vars, ix `ixsAppend` ix2)
astScatterKnobsS knobs (snat@(SNat' @1) :$$ ZSS) _shn shp
(Ast.AstReplicateS (_ :$$ ZSS)
(Ast.AstScatterS @_ @shn2 @shp2 shm2 shn2 shp2
v (vars, ix2)))
(var ::$ ZS, ix) = -- oneHot1 (scatter) fusion
gcastWith (unsafeCoerceRefl :: shp ++ shn :~: shp ++ shp2 ++ shn2) $
astScatterKnobsS knobs (snat :$$ shm2) shn2 (shp `shsAppend` shp2)
(Ast.AstReplicateS (snat :$$ ZSS) v)
(var ::$ vars, ix `ixsAppend` ix2)
astScatterKnobsS knobs (snat@(SNat' @1) :$$ ZSS) _shn shp
(Ast.AstReplicateS (_ :$$ ZSS) -- TODO: generalize
(Ast.AstScatterS @_ @shn2 @shp2 shm2 shn2 shp2
v (vars, ix2)))
(var ::$ ZS, ix) = -- oneHot1 (scatter) fusion
gcastWith (unsafeCoerceRefl :: shp ++ shn :~: shp ++ shp2 ++ shn2) $
astScatterKnobsS knobs (snat :$$ shm2) shn2 (shp `shsAppend` shp2)
(Ast.AstReplicateS (snat :$$ ZSS) v)
(var ::$ vars, ix `ixsAppend` ix2)
astScatterKnobsS knobs shm shn shp (Ast.AstTransposeS @perm @sh perm v)
(vars, ix) -- see the same case of astGatherKnobsS
| FTKS sh _ <- ftkAst v
, knobPhase knobs `elem` [PhaseVectorization, PhaseExpansion]
, let rankPerm = Permutation.permRank perm
, fromSNat' rankPerm <= length vars =
gcastWith (lemRankMapJust $ shsTakeLenPerm perm sh) $
gcastWith (unsafeCoerceRefl :: Rank (TakeLen perm sh) :~: Rank perm) $
permInverse perm
$ \(invperm :: Nested.Perm invperm) proof ->
case proof (ssxFromShX $ shxFromShS $ shsTakeLenPerm perm sh) of
Refl ->
gcastWith
(unsafeCoerceRefl
:: Permutation.PermutePrefix invperm shm ++ shn
:~: Permutation.PermutePrefix invperm (shm ++ shn)) $
gcastWith
(unsafeCoerceRefl
:: Permutation.Permute invperm (TakeLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh))
:~: Permutation.Permute invperm (Permutation.Permute perm (TakeLen perm sh))) $
gcastWith
(unsafeCoerceRefl
:: Permutation.Permute invperm (Permutation.Permute perm (TakeLen perm sh))
:~: TakeLen perm sh) $
gcastWith
(unsafeCoerceRefl
:: DropLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh)
:~: DropLen perm sh) $
gcastWith
(unsafeCoerceRefl
:: TakeLen perm sh ++ DropLen perm sh :~: sh) $
let invvars = AstVarListS
$ ixsPermutePrefix invperm $ unAstVarListS vars
in astScatterKnobsS knobs (shsPermutePrefix invperm shm) shn shp
v (invvars, ix)
astScatterKnobsS knobs
shm@(SNat @m :$$ (_ :: ShS shmTail))
shn
shp@(SNat @p :$$ (_ :: ShS shpTail))
v0
( varm ::$ mrest
, AstIntVar varp :.$ prest )
| knobPhase knobs `notElem` [PhaseVectorization, PhaseExpansion]
-- prevent a loop
, varm == varp
, not (varm `varNameInIxS` prest)
, FTKS _ x <- ftkAst v0 =
withSNat (min (valueOf @p) (valueOf @m)) $ \(SNat @m2) ->
gcastWith (unsafeCoerceRefl :: (m2 <=? p) :~: True) $
gcastWith (unsafeCoerceRefl :: (m2 <=? m) :~: True) $
Permutation.permFromListCont (permCycle
$ shsLength (shsTail shp) + 1)
$ \(permVars :: Permutation.Perm permVars) ->
Permutation.permFromListCont (backpermCycle
$ shsLength (shsTail shm) + 1)
$ \(permIx :: Permutation.Perm permIx) ->
gcastWith (unsafeCoerceRefl
:: m2 ': shpTail ++ shn
:~: Permutation.PermutePrefix
permVars (shpTail ++ (m2 ': shn))) $
gcastWith (unsafeCoerceRefl
:: shmTail ++ (m2 ': shn)
:~: Permutation.PermutePrefix
permIx (m2 ': shmTail ++ shn)) $
gcastWith (unsafeCoerceRefl
:: (Rank permVars <=? Rank (shpTail ++ (m2 ': shn)))
:~: True) $
gcastWith (unsafeCoerceRefl
:: (Rank permIx <=? Rank (m2 ': shmTail ++ shn)) :~: True) $
fromMaybe (error "astScatterKnobsS: impossible non-permutation")
$ Permutation.permCheckPermutation permVars
$ fromMaybe (error "astScatterKnobsS: impossible non-permutation")
$ Permutation.permCheckPermutation permIx
$ let v2 = astTransposeS permIx
$ astSliceS (SNat @0) (SNat @m2) (SNat @(m - m2)) v0
u = astScatterKnobsS
knobs (shsTail shm) (SNat @m2 :$$ shn) (shsTail shp)
v2 (mrest, prest)
ftk = FTKS (SNat @(p - m2) :$$ shsTail shp `shsAppend` shn) x
in astTransposeS permVars u
`astAppendS`
fromPlain (astConcrete ftk (tdefTarget ftk))
astScatterKnobsS knobs shm@(SNat' @1 :$$ ZSS) shn shp
v@Ast.AstReplicateS{} (vars, ix)
| knobPhase knobs /= PhaseContraction =
Ast.AstScatterS shm shn shp v (vars, ix) -- oneHot1 NF
-- The above normal form that prevents the use of the rule below
-- and the subsequent rule is to keep one-hots and similar in forms
-- easier for fusion of addition of scatters to operate on.
astScatterKnobsS knobs (snat :$$ shm2) shn shp v (var ::$ vars, ix)
| not $ var `varNameInIxS` ix = -- eliminates oneHot1
astScatterKnobsS knobs shm2 shn shp
(astSumS (snat :$$ ZSS) v) (vars, ix)
astScatterKnobsS knobs ZSS shn shp u (ZS, ix) -- oneHot1 intro
| knobPhase knobs /= PhaseContraction =
let shm = SNat @1 :$$ ZSS
a = Ast.AstReplicateS (SNat @1 :$$ ZSS) u
in funToVarsIxS shm $ \(var ::$ ZS) _ ->
Ast.AstScatterS shm shn shp a (var ::$ ZS, ix)
astScatterKnobsS knobs shm shn shp@(p@(SNat @p) :$$ _) v0
(vars, AstPlusK (AstConcreteK i0) i1 :.$ prest)
| knobPhase knobs `elem` [PhaseExpansion, PhaseContraction]
-- give time for additions of scatters to fuse; TODO: harden fusion
-- or keep as normal forms more terms than one-hots, until contraction
, Just (lb, _) <- intBounds i1
, lb >= 0 -- ensured by the OOB rule above: fromSNat' p >= i0
|| i0 <= 0
, FTKS _ x <- ftkAst v0 =
if i0 >= 0 then
withSNat i0 $ \i@(SNat @i) ->
withSNat (fromSNat' p - i0) $ \k@(SNat @k) ->
gcastWith (unsafeCoerceRefl :: i + k :~: p) $
let ftk = FTKS (i :$$ shsTail shp `shsAppend` shn) x
in astAppendS
(fromPlain (astConcrete ftk (tdefTarget ftk)))
(astScatterKnobsS knobs shm shn (k :$$ shsTail shp)
v0 (vars, i1 :.$ prest))
-- this gather may still index out of bounds, which is fine
else
withSNat (- i0) $ \i ->
astSliceS i p SZ
$ astScatterKnobsS knobs shm shn (snatPlus p i :$$ shsTail shp)
v0 (vars, i1 :.$ prest)
-- this gather may still index out of bounds, which is fine
astScatterKnobsS knobs shm shn shp@(p@(SNat @p) :$$ _) v0
(vars, Ast.AstLet varN uN (AstPlusK (AstConcreteK i0) i1) :.$ prest)
| knobPhase knobs `elem` [PhaseExpansion, PhaseContraction]
-- give time for additions of scatters to fuse; TODO
, Just (lb, _) <- intBounds i1
, lb >= 0 -- ensured by the OOB rule above: fromSNat' p >= i0
|| i0 <= 0
, FTKS _ x <- ftkAst v0 =
if i0 >= 0 then
withSNat i0 $ \i@(SNat @i) ->
withSNat (fromSNat' p - i0) $ \k@(SNat @k) ->
gcastWith (unsafeCoerceRefl :: i + k :~: p) $
let ftk = FTKS (i :$$ shsTail shp `shsAppend` shn) x
in astAppendS
(fromPlain (astConcrete ftk (tdefTarget ftk)))
(astScatterKnobsS knobs shm shn (k :$$ shsTail shp)
v0 (vars, astLet varN uN i1 :.$ prest))
-- this gather may still index out of bounds, which is fine
else
withSNat (- i0) $ \i ->
astSliceS i p SZ
$ astScatterKnobsS knobs shm shn (snatPlus p i :$$ shsTail shp)
v0 (vars, astLet varN uN i1 :.$ prest)
-- this gather may still index out of bounds, which is fine
-- These rules are questionable, because at worst, we allocate twice.
-- However, often the scatters or one of them would simplify away.
-- Note these don't revert the scatter addtion fusion in rules for astPlusS,
-- because the latter don't fire if the conditional can't be simplified away.
astScatterKnobsS knobs shm@(SNat @m :$$ _) shn shp u0
( vars@(varm ::$ mrest)
, Ast.AstCondK (AstLeqInt (AstConcreteK j0) (AstIntVar varp)) v w
:.$ prest )
| varm == varp
, j0 <= 0 || j0 >= valueOf @m || ixIsSmall prest
, FTKS _ (FTKScalar @r) <- ftkAst u0
, Dict0 <- numFromTKAllNum (Proxy @r) =
if | j0 <= 0 ->
astScatterKnobsS knobs shm shn shp
u0 (vars, v :.$ prest)
| j0 >= valueOf @m ->
astScatterKnobsS knobs shm shn shp
u0 (vars, w :.$ prest)
| otherwise ->
withSNat j0 $ \j@(SNat @j) ->
withSNat (valueOf @m - j0) $ \msj@(SNat @msj) ->
gcastWith (unsafeCoerceRefl :: msj + j :~: m) $
astLetFun u0 $ \u ->
let u2 = astSliceS SZ j msj u
u3 = astSliceS j msj SZ u
varm2 = reboundsVarName (0, j0 - 1) varm
varm3 = reboundsVarName (0, valueOf @msj - 1) varm
in astScatterKnobsS knobs (j :$$ shsTail shm) shn shp u2
( varm2 ::$ mrest
, substituteAstIxS (astVar varm2)
varm (w :.$ prest) )
`astPlusS`
astScatterKnobsS knobs (msj :$$ shsTail shm) shn shp u3
( varm3 ::$ mrest
, substituteAstIxS (AstConcreteK j0 + astVar varm3)
varm (v :.$ prest) )
astScatterKnobsS knobs shm@(SNat @m :$$ _) shn shp u0
( vars@(varm ::$ mrest)
, Ast.AstLet varN uN
(Ast.AstCondK (AstLeqInt (AstConcreteK j0) (AstIntVar varp)) v w)
:.$ prest )
| varm == varp
, let varIn1 = varN `varNameInAst` v
varIn2 = varN `varNameInAst` w
, j0 <= 0 || j0 >= valueOf @m
|| ixIsSmall prest && not (varIn1 && varIn2)
, FTKS _ (FTKScalar @r) <- ftkAst u0
, Dict0 <- numFromTKAllNum (Proxy @r) =
if | j0 <= 0 ->
astScatterKnobsS knobs shm shn shp
u0 (vars, astLet varN uN v :.$ prest)
| j0 >= valueOf @m ->
astScatterKnobsS knobs shm shn shp
u0 (vars, astLet varN uN w :.$ prest)
| otherwise ->
withSNat j0 $ \j@(SNat @j) ->
withSNat (valueOf @m - j0) $ \msj@(SNat @msj) ->
gcastWith (unsafeCoerceRefl :: msj + j :~: m) $
astLetFun u0 $ \u ->
let u2 = astSliceS SZ j msj u
u3 = astSliceS j msj SZ u
varm2 = reboundsVarName (0, j0 - 1) varm
varm3 = reboundsVarName (0, valueOf @msj - 1) varm
in astScatterKnobsS knobs (j :$$ shsTail shm) shn shp u2
( varm2 ::$ mrest
, substituteAstIxS (astVar varm2)
varm ((if varIn2
then astLet varN uN w
else w) :.$ prest) )
`astPlusS`
astScatterKnobsS knobs (msj :$$ shsTail shm) shn shp u3
( varm3 ::$ mrest
, substituteAstIxS (AstConcreteK j0 + astVar varm3)
varm ((if varIn1
then astLet varN uN v
else v) :.$ prest) )
astScatterKnobsS _ shm shn shp v (vars, ix) =
Ast.AstScatterS shm shn shp v (vars, ix)
astReplicateK :: forall shm r s. (GoodScalar r, KnownSpan s)
=> ShS shm -> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS shm r)
astReplicateK ZSS t0 = sfromK t0
astReplicateK shm t0 = case t0 of
Ast.AstLet var t v -> astLet var t (astReplicateK shm v)
Ast.AstFromPrimal v -> fromPrimal $ astReplicateK shm v
Ast.AstFromDual v -> fromDual $ astReplicateK shm v
Ast.AstFromPlain v -> fromPlain $ astReplicateK shm v
AstConcreteK t | shsSize shm < 100 -> -- tough trade-offs here
astConcreteS $ tsreplicate0N shm $ Concrete t
_ -> Ast.AstReplicateK shm t0
astReplicateS :: forall shm shn x s. KnownSpan s
=> ShS shm -> AstTensor AstMethodLet s (TKS2 shn x)
-> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)
astReplicateS ZSS t0 = t0
astReplicateS shm t0 = case t0 of
-- TODO: This rules is, in principle, very good, because it permits many other
-- rules to fire. However, one of these other rules is indexing of transpose
-- that in some cases complicates terms and degrades performance considerably.
-- Ast.AstLet var t v -> astLet var t (astReplicateS shm v)
Ast.AstFromPrimal v -> fromPrimal $ astReplicateS shm v
Ast.AstFromDual v -> fromDual $ astReplicateS shm v
Ast.AstFromPlain v -> fromPlain $ astReplicateS shm v
AstConcreteS t | shsSize shm < 100 -> -- tough trade-offs here
let shn = Nested.sshape t
in withKnownShS shn $
astConcreteS $ tsreplicateN shm $ Concrete t
-- revisit the trade-offs once we compile instead of interpreting
-- and so building big blobby concrete arrays is cheap
Ast.AstReplicateK shm2 v ->
astReplicateK (shm `shsAppend` shm2) v
Ast.AstReplicateS @shm2 @shn2 shm2 v
| Refl <- lemAppAssoc (Proxy @shm) (Proxy @shm2) (Proxy @shn2) ->
astReplicateS (shm `shsAppend` shm2) v
_ | FTKS ZSS FTKScalar <- ftkAst t0
, Refl <- lemAppNil @shm -> Ast.AstReplicateK shm (kfromS t0)
_ -> Ast.AstReplicateS shm t0
-- TODO: maybe add a rule and then generalize:
-- replicate n1 (str (replicate n2 u))
-- ~> transpose [0, 2, 1] (replicate n1 (replicate n2 u))
-- but the reverse rule is already in astTransposeS
flipCompare :: forall (a :: Nat) b. Compare a b ~ GT
=> Proxy a -> Proxy b -> Compare b a :~: LT
flipCompare _ _ = unsafeCoerceRefl
astGatherS
:: forall shm shn shp r s. KnownSpan s
=> ShS shm -> ShS shn -> ShS shp
-> AstTensor AstMethodLet s (TKS2 (shp ++ shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 (shm ++ shn) r)
astGatherS = astGatherKnobsS defaultKnobs
-- Assumption: vars0 don't not occur in v0. The assumption only holds
-- when newly generated variables are fresh, which is the case as long
-- as resetVarCounter is not used. The assumption makes it easier to spot
-- bugs or corruption, hence we assert it in the code below.
astGatherKnobsS
:: forall shm shn shp x s. KnownSpan s
=> SimplifyKnobs
-> ShS shm -> ShS shn -> ShS shp
-> AstTensor AstMethodLet s (TKS2 (shp ++ shn) x)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)
astGatherKnobsS _ _ _ _ v0 (!vars0, !_ix0)
| Foldable.any (`varNameInAst` v0) vars0 =
error $ "astGatherKnobsS: gather vars in v0: " ++ show (vars0, v0)
astGatherKnobsS knobs _ shn _ v0 (ZS, ix0) = astIndexKnobsS knobs shn v0 ix0
astGatherKnobsS _ shm _ _ v0 (_, ZIS) = astReplicateS shm v0
astGatherKnobsS _ shm@(SZ :$$ _) shn _ v0 _
| FTKS _ x <- ftkAst v0 =
let ftk = FTKS (shm `shsAppend` shn) x
in fromPlain $ astConcrete ftk (tdefTarget ftk)
astGatherKnobsS _ shm shn _shp v0 (_, i1 :.$ _)
| Just (lb, ub) <- intBounds i1
, FTKS (snat :$$ _) x <- ftkAst v0
, ub < 0 || lb >= fromSNat' snat =
let ftk = FTKS (shm `shsAppend` shn) x
in fromPlain $ astConcrete ftk (tdefTarget ftk)
astGatherKnobsS knobs shm@(k :$$ _) shn shp v0 (var1 ::$ vars1, ix0)
| not (var1 `varNameInIxS` ix0) =
astReplicateS (k :$$ ZSS)
(astGatherKnobsS knobs (shsTail shm) shn shp v0 (vars1, ix0))
astGatherKnobsS knobs shm shn shp v0 (vars0@(_ ::$ _), ix0@(_ :.$ _))
| let ixInit = ixsInit ix0
varInit = AstVarListS $ ixsInit $ unAstVarListS vars0
varLast = ixsLast $ unAstVarListS vars0
, AstIntVar ixvarLast <- ixsLast ix0
, ixvarLast == varLast
, not (varLast `varNameInIxS` ixInit)
, kLast@SNat <- shsLast shm
, Just Refl <- testEquality kLast (shsLast shp) =
gcastWith (unsafeCoerceRefl
:: Init shp ++ (Last shm ': shn) :~: shp ++ shn) $
gcastWith (unsafeCoerceRefl
:: Init shm ++ (Last shm ': shn) :~: shm ++ shn) $
astGatherKnobsS knobs (shsInit shm) (kLast :$$ shn) (shsInit shp)
v0 (varInit, ixInit)
astGatherKnobsS
knobs shm@(m :$$ _) shn shp v0
( vars@(varm ::$ _)
, Ast.AstCondK
(Ast.AstBoolAndK a@(AstLeqInt (AstConcreteK j0) (AstIntVar varp)) b)
v w :.$ prest )
| varm == varp
, j0 <= 0 || j0 >= fromSNat' m || ixIsSmall prest =
let i = astLetFun w $ \wShared ->
astCondKInitial a (astCondKInitial b v wShared) wShared
in astGatherKnobsS knobs shm shn shp v0 (vars, i :.$ prest)
astGatherKnobsS
knobs shm@(m :$$ _) shn shp v0
( vars@(varm ::$ _)
, Ast.AstLet varN uN
(Ast.AstCondK
(Ast.AstBoolAndK a@(AstLeqInt (AstConcreteK j0) (AstIntVar varp)) b)
v w) :.$ prest )
| varm == varp
, let varIn1 = varN `varNameInAst` v
varIn2 = varN `varNameInAst` w
, j0 <= 0 || j0 >= fromSNat' m
|| ixIsSmall prest && not (varIn1 && varIn2) =
let i = astLetFun w $ \wShared ->
astCondKInitial a (astCondKInitial b v wShared) wShared
in astGatherKnobsS knobs shm shn shp v0 (vars, astLet varN uN i :.$ prest)
astGatherKnobsS
knobs shm@(m :$$ _) shn shp v0
( vars@(varm ::$ _)
, Ast.AstCondK
(Ast.AstBoolAndK a@(AstLeqInt (AstConcreteK j0)
(Ast.AstN1K NegateOp (AstIntVar varp))) b)
v w :.$ prest )
| varm == varp
, - j0 + 1 <= 0 || - j0 + 1 >= fromSNat' m || ixIsSmall prest =
let i = astLetFun w $ \wShared ->
astCondKInitial a (astCondKInitial b v wShared) wShared
in astGatherKnobsS knobs shm shn shp v0 (vars, i :.$ prest)
astGatherKnobsS
knobs shm@(m :$$ _) shn shp v0
( vars@(varm ::$ _)
, Ast.AstLet varN uN
(Ast.AstCondK
(Ast.AstBoolAndK
a@(AstLeqInt (AstConcreteK j0)
(Ast.AstN1K NegateOp (AstIntVar varp))) b)
v w) :.$ prest )
| varm == varp
, let varIn1 = varN `varNameInAst` v
varIn2 = varN `varNameInAst` w
, - j0 + 1 <= 0 || - j0 + 1 >= fromSNat' m
|| ixIsSmall prest && not (varIn1 && varIn2) =
let i = astLetFun w $ \wShared ->
astCondKInitial a (astCondKInitial b v wShared) wShared
in astGatherKnobsS knobs shm shn shp v0 (vars, astLet varN uN i :.$ prest)
astGatherKnobsS
knobs shm@(m :$$ _) shn shp v0
( vars@(varm ::$ _)
, Ast.AstCondK
(Ast.AstBoolAndK
a@(Ast.AstBoolNotK
(Ast.AstBoolAndK
(AstLeqInt (AstConcreteK j0) (AstIntVar varp)) _)) b)
v w :.$ prest )
| varm == varp
, j0 <= 0 || j0 >= fromSNat' m || ixIsSmall prest =
let i = astLetFun w $ \wShared ->
astCondKInitial a (astCondKInitial b v wShared) wShared
in astGatherKnobsS knobs shm shn shp v0 (vars, i :.$ prest)
astGatherKnobsS
knobs shm@(m :$$ _) shn shp v0
( vars@(varm ::$ _)
, Ast.AstLet varN uN
(Ast.AstCondK
(Ast.AstBoolAndK
a@(Ast.AstBoolNotK
(Ast.AstBoolAndK (AstLeqInt (AstConcreteK j0)
(AstIntVar varp)) _)) b)
v w) :.$ prest )
| varm == varp
, let varIn1 = varN `varNameInAst` v
varIn2 = varN `varNameInAst` w
, j0 <= 0 || j0 >= fromSNat' m
|| ixIsSmall prest && not (varIn1 && varIn2) =
let i = astLetFun w $ \wShared ->
astCondKInitial a (astCondKInitial b v wShared) wShared
in astGatherKnobsS knobs shm shn shp v0 (vars, astLet varN uN i :.$ prest)
astGatherKnobsS
knobs shm@(m :$$ _) shn shp v0
( vars@(varm ::$ _)
, Ast.AstCondK
(Ast.AstBoolAndK
a@(Ast.AstBoolNotK
(Ast.AstBoolAndK
(AstLeqInt (AstConcreteK j0)
(Ast.AstN1K NegateOp (AstIntVar varp))) _)) b)
v w :.$ prest )
| varm == varp
, - j0 + 1 <= 0 || - j0 + 1 >= fromSNat' m || ixIsSmall prest =
let i = astLetFun w $ \wShared ->
astCondKInitial a (astCondKInitial b v wShared) wShared
in astGatherKnobsS knobs shm shn shp v0 (vars, i :.$ prest)
astGatherKnobsS
knobs shm@(m :$$ _) shn shp v0
( vars@(varm ::$ _)
, Ast.AstLet varN uN
(Ast.AstCondK
(Ast.AstBoolAndK
a@(Ast.AstBoolNotK
(Ast.AstBoolAndK
(AstLeqInt (AstConcreteK j0)
(Ast.AstN1K NegateOp (AstIntVar varp))) _)) b)
v w) :.$ prest )
| varm == varp
, let varIn1 = varN `varNameInAst` v
varIn2 = varN `varNameInAst` w
, - j0 + 1 <= 0 || - j0 + 1 >= fromSNat' m
|| ixIsSmall prest && not (varIn1 && varIn2) =
let i = astLetFun w $ \wShared ->
astCondKInitial a (astCondKInitial b v wShared) wShared
in astGatherKnobsS knobs shm shn shp v0 (vars, astLet varN uN i :.$ prest)
-- Rules with AstConcreteK on the right hand side of AstPlusK are
-- not needed, thanks to the normal form of AstPlusK rewriting.
astGatherKnobsS knobs shm shn shp v0
(vars, AstPlusK (AstConcreteK i) i1 :.$ prest)
| Just (lb, ub) <- intBounds i1
, lb >= 0 -- if not, we may need to apply astReverse first
|| i <= 0
, FTKS (SNat @p :$$ _) x <- ftkAst v0 =
if i >= 0 then
withSNat i $ \(SNat @i) ->
withSNat (min (valueOf @p - i) (ub + 1)) $ \(SNat @k) ->
gcastWith (unsafeCoerceRefl :: (i + k <=? p) :~: True) $
let v2 = astSliceS (SNat @i) (SNat @k) (SNat @(p - (i + k))) v0
in astGatherKnobsS knobs shm shn (SNat @k :$$ shsTail shp)
v2 (vars, i1 :.$ prest)
-- this gather may still index out of bounds, which is fine
else
withSNat (- i) $ \(SNat @i) ->
let ftk = FTKS (SNat @i :$$ shsTail shp `shsAppend` shn) x
v2 = fromPlain (astConcrete ftk (tdefTarget ftk))
`astAppendS`
v0
in astGatherKnobsS knobs shm shn (SNat @(p + i) :$$ shsTail shp)
v2 (vars, i1 :.$ prest)
-- this gather may still index out of bounds, which is fine
astGatherKnobsS knobs shm shn shp v0
(vars, Ast.AstLet varN uN (AstPlusK (AstConcreteK i) i1) :.$ prest)
| Just (lb, ub) <- intBounds i1
, lb >= 0 -- if not, we may need to apply astReverse first
|| i <= 0
, FTKS (SNat @p :$$ _) x <- ftkAst v0 =
if i >= 0 then
withSNat i $ \(SNat @i) ->
withSNat (min (valueOf @p - i) (ub + 1)) $ \(SNat @k) ->
gcastWith (unsafeCoerceRefl :: (i + k <=? p) :~: True) $
let v2 = astSliceS (SNat @i) (SNat @k) (SNat @(p - (i + k))) v0
in astGatherKnobsS knobs shm shn (SNat @k :$$ shsTail shp)
v2 (vars, astLet varN uN i1 :.$ prest)
-- this gather may still index out of bounds, which is fine
else
withSNat (- i) $ \(SNat @i) ->
let ftk = FTKS (SNat @i :$$ shsTail shp `shsAppend` shn) x
v2 = fromPlain (astConcrete ftk (tdefTarget ftk))
`astAppendS`
v0
in astGatherKnobsS knobs shm shn (SNat @(p + i) :$$ shsTail shp)
v2 (vars, astLet varN uN i1 :.$ prest)
-- this gather may still index out of bounds, which is fine
astGatherKnobsS knobs (m@(SNat @m) :$$ (shmRest :: ShS shmRest)) shn shp v0
( varm ::$ mrest
, Ast.AstI2K QuotOp (AstIntVar varm') (AstConcreteK i0) :.$ prest )
| varm == varm'
, not (varm `varNameInIxS` prest)
, i0 > 0 -- ensured by other rules; makes this easier to reason about
, let k0 = (fromSNat' m + i0 - 1) `quot` i0
, SNat <- shsRank shmRest
, SNat <- shsRank shn =
withSNat i0 $ \i@(SNat @i) ->
withSNat k0 $ \k@(SNat @k) ->
withSNat (k0 * i0 - fromSNat' m) $ \z@(SNat @z) ->
let perm = Permutation.makePerm @'[1, 0]
varm2 = reboundsVarName (0, k0 - 1) varm
in gcastWith (unsafeCoerceRefl :: m + z :~: k * i) $
gcastWith (unsafeCoerceRefl
:: Rank ([i, k] ++ shmRest ++ shn)
:~: 2 + Rank shmRest + Rank shn) $
astSliceS SZ m z
$ astReshapeS (snatMul k i :$$ shmRest `shsAppend` shn)
$ astTransposeS perm
$ astReplicateS (i :$$ ZSS)
$ astGatherKnobsS knobs (k :$$ shmRest) shn shp
v0 (varm2 ::$ mrest, astVar varm2 :.$ prest)
-- this gather may still index out of bounds, which is fine
astGatherKnobsS knobs shm@(SNat @m :$$ _) shn shp u0
( vars@(varm ::$ mrest)
, Ast.AstCondK (AstLeqInt (AstConcreteK j) (AstIntVar varp)) v w
:.$ prest )
| varm == varp
, j <= 0 || j >= valueOf @m || ixIsSmall prest =
if | j <= 0 ->
astGatherKnobsS knobs shm shn shp
u0 (vars, v :.$ prest)
| j >= valueOf @m ->
astGatherKnobsS knobs shm shn shp
u0 (vars, w :.$ prest)
| otherwise ->
withSNat j $ \(SNat @j) ->
gcastWith (unsafeCoerceRefl :: (j <=? m) :~: True) $
astLetFun u0 $ \u ->
let varm2 = reboundsVarName (0, j - 1) varm
varm3 = reboundsVarName (0, valueOf @m - j - 1) varm
in astGatherKnobsS knobs (SNat @j :$$ shsTail shm) shn shp u
( varm2 ::$ mrest
, substituteAstIxS (astVar varm2)
varm (w :.$ prest) )
`astAppendS`
astGatherKnobsS knobs (SNat @(m - j) :$$ shsTail shm) shn shp u
( varm3 ::$ mrest
, substituteAstIxS (AstConcreteK j + astVar varm3)
varm (v :.$ prest) )
astGatherKnobsS knobs shm@(SNat @m :$$ _) shn shp u0
( vars@(varm ::$ mrest)
, Ast.AstLet varN uN
(Ast.AstCondK (AstLeqInt (AstConcreteK j) (AstIntVar varp)) v w)
:.$ prest )
| varm == varp
, let varIn1 = varN `varNameInAst` v
varIn2 = varN `varNameInAst` w
, j <= 0 || j >= valueOf @m
|| ixIsSmall prest && not (varIn1 && varIn2) =
if | j <= 0 ->
astGatherKnobsS knobs shm shn shp
u0 (vars, astLet varN uN v :.$ prest)
| j >= valueOf @m ->
astGatherKnobsS knobs shm shn shp
u0 (vars, astLet varN uN w :.$ prest)
| otherwise ->
withSNat j $ \(SNat @j) ->
gcastWith (unsafeCoerceRefl :: (j <=? m) :~: True) $
astLetFun u0 $ \u ->
let varm2 = reboundsVarName (0, j - 1) varm
varm3 = reboundsVarName (0, valueOf @m - j - 1) varm
in astGatherKnobsS knobs (SNat @j :$$ shsTail shm) shn shp u
( varm2 ::$ mrest
, substituteAstIxS (astVar varm2)
varm ((if varIn2
then astLet varN uN w
else w) :.$ prest) )
`astAppendS`
astGatherKnobsS knobs (SNat @(m - j) :$$ shsTail shm) shn shp u
( varm3 ::$ mrest
, substituteAstIxS (AstConcreteK j + astVar varm3)
varm ((if varIn1
then astLet varN uN v
else v) :.$ prest) )
astGatherKnobsS knobs
shm@(SNat @m :$$ (_ :: ShS shmTail))
shn
shp@(SNat @p :$$ (_ :: ShS shpTail))
v0
( varm ::$ mrest
, AstIntVar varp :.$ prest )
| knobPhase knobs `notElem` [PhaseVectorization, PhaseExpansion]
-- prevent a loop
, varm == varp
, not (varm `varNameInIxS` prest)
, FTKS _ x <- ftkAst v0 =
withSNat (min (valueOf @p) (valueOf @m)) $ \(SNat @m2) ->
gcastWith (unsafeCoerceRefl :: (m2 <=? p) :~: True) $
gcastWith (unsafeCoerceRefl :: (m2 <=? m) :~: True) $
Permutation.permFromListCont (permCycle
$ shsLength (shsTail shm) + 1)
$ \(permVars :: Permutation.Perm permVars) ->
Permutation.permFromListCont (backpermCycle
$ shsLength (shsTail shp) + 1)
$ \(permIx :: Permutation.Perm permIx) ->
gcastWith (unsafeCoerceRefl
:: m2 ': shmTail ++ shn
:~: Permutation.PermutePrefix
permVars (shmTail ++ (m2 ': shn))) $
gcastWith (unsafeCoerceRefl
:: shpTail ++ (m2 ': shn)
:~: Permutation.PermutePrefix
permIx (m2 ': shpTail ++ shn)) $
gcastWith (unsafeCoerceRefl
:: (Rank permVars <=? Rank (shmTail ++ (m2 ': shn)))
:~: True) $
gcastWith (unsafeCoerceRefl
:: (Rank permIx <=? Rank (m2 ': shpTail ++ shn)) :~: True) $
fromMaybe (error "astGatherKnobsS: impossible non-permutation")
$ Permutation.permCheckPermutation permVars
$ fromMaybe (error "astGatherKnobsS: impossible non-permutation")
$ Permutation.permCheckPermutation permIx
$ let v2 = astTransposeS permIx
$ astSliceS (SNat @0) (SNat @m2) (SNat @(p - m2)) v0
u = astGatherKnobsS
knobs (shsTail shm) (SNat @m2 :$$ shn) (shsTail shp)
v2 (mrest, prest)
ftk = FTKS (SNat @(m - m2) :$$ shsTail shm `shsAppend` shn) x
in astTransposeS permVars u
`astAppendS`
fromPlain (astConcrete ftk (tdefTarget ftk))
astGatherKnobsS knobs (m :$$ shm4) ZSS _shp
v4@(Ast.AstFromVectorK (_ :$$ ZSS) l)
( var4 ::$ _vrest4
, i4 :.$ rest4 )
| knobPhase knobs `notElem` [PhaseVectorization, PhaseExpansion]
, V.length l < 100
, let g i = case substituteAst (AstConcreteK i) var4 i4 of
AstConcreteK j -> Just j
_ -> Nothing
, Just lj <- mapM g [0 .. fromSNat' m - 1]
, Refl <- lemAppNil @shm = assert (null rest4) $
let FTKS _ x = ftkAst v4
ftk = FTKS shm4 x
zero = fromPlain $ astConcrete ftk (tdefTarget ftk)
f j =
if j >= V.length l
then zero
else astReplicateK shm4 (l V.! j)
in astFromVectorS (m :$$ ZSS)
$ V.fromListN (fromSNat' m) $ map f lj
astGatherKnobsS knobs shm@(m :$$ (shm4 :: ShS shm4)) shn shp
v4@(Ast.AstFromVectorS (_ :$$ ZSS) l)
( var4 ::$ vrest4
, i4 :.$ rest4 )
| knobPhase knobs `notElem` [PhaseVectorization, PhaseExpansion]
, V.length l < 100
, let g i = case substituteAst (AstConcreteK i) var4 i4 of
AstConcreteK j -> Just j
_ -> Nothing
, Just lj <- mapM g [0 .. fromSNat' m - 1]
, ixIsSmall rest4 =
let FTKS _ x = ftkAst v4
ftk = FTKS (shsTail shm `shsAppend` shn) x
zero = fromPlain $ astConcrete ftk (tdefTarget ftk)
f i j =
let subRest4 = substituteAstIxS (AstConcreteK i) var4 rest4
in if j >= V.length l
then zero
else astGatherKnobsS knobs shm4 shn (shsTail shp)
(l V.! j) (vrest4, subRest4)
in astFromVectorS (m :$$ ZSS)
$ V.fromListN (fromSNat' m) $ zipWith f [0 .. fromSNat' m - 1] lj
astGatherKnobsS knobs shm shn shp v0 (vars0, i1 :.$ rest1)
| knobPhase knobs `notElem` [PhaseVectorization, PhaseExpansion]
-- prevent a loop
, not (Foldable.any (`varNameInAst` i1) vars0) =
astGatherKnobsS
knobs shm shn (shsTail shp)
(astIndexKnobsS knobs (shsTail shp `shsAppend` shn) v0 (i1 :.$ ZIS))
(vars0, rest1)
astGatherKnobsS knobs shm@(m :$$ _) shn shp v0
(vars, ix@(i1 :.$ prest))
| knobPhase knobs `notElem` [PhaseVectorization, PhaseExpansion]
-- prevent a loop
, let intInteresting = \case
AstPlusK (AstConcreteK i) i2
| Just (lb, _) <- intBounds i2
, lb >= 0 || i <= 0-> True
Ast.AstLet _ _ (AstPlusK (AstConcreteK i) i2)
| Just (lb, _) <- intBounds i2
, lb >= 0 || i <= 0-> True
Ast.AstCondK (AstLeqInt (AstConcreteK j0) (AstIntVar var)) _ _
| j0 <= 0 || j0 >= fromSNat' m || ixIsSmall prest
, Foldable.any (== var) vars -> True
Ast.AstLet varN _
(Ast.AstCondK (AstLeqInt (AstConcreteK j0) (AstIntVar var)) v w)
| let varIn1 = varN `varNameInAst` v
varIn2 = varN `varNameInAst` w
, j0 <= 0 || j0 >= fromSNat' m
|| ixIsSmall prest && not (varIn1 && varIn2)
, Foldable.any (== var) vars -> True
Ast.AstCondK
(Ast.AstBoolAndK
(AstLeqInt (AstConcreteK j0) (AstIntVar var)) _) _ _
| j0 <= 0 || j0 >= fromSNat' m || ixIsSmall prest
, Foldable.any (== var) vars -> True
Ast.AstCondK
(Ast.AstBoolAndK
(AstLeqInt (AstConcreteK j0)
(Ast.AstN1K NegateOp (AstIntVar var))) _) _ _
| - j0 + 1 <= 0 || - j0 + 1 >= fromSNat' m || ixIsSmall prest
, Foldable.any (== var) vars -> True
Ast.AstLet varN _
(Ast.AstCondK
(Ast.AstBoolAndK
(AstLeqInt (AstConcreteK j0) (AstIntVar var)) _) v w)
| let varIn1 = varN `varNameInAst` v
varIn2 = varN `varNameInAst` w
, j0 <= 0 || j0 >= fromSNat' m
|| ixIsSmall prest && not (varIn1 && varIn2)
, Foldable.any (== var) vars -> True
Ast.AstLet varN _
(Ast.AstCondK
(Ast.AstBoolAndK
(AstLeqInt (AstConcreteK j0)
(Ast.AstN1K NegateOp (AstIntVar var))) _) v w)
| let varIn1 = varN `varNameInAst` v
varIn2 = varN `varNameInAst` w
, - j0 + 1 <= 0 || - j0 + 1 >= fromSNat' m
|| ixIsSmall prest && not (varIn1 && varIn2)
, Foldable.any (== var) vars -> True
AstIntVar var
| knobPhase knobs `elem` [PhaseSimplification, PhaseContraction]
, null $ drop 1 $ filter (var `varNameInAst`) (Foldable.toList ix)
, Foldable.any (== var) vars -> True
ik | knobPhase knobs `elem` [PhaseSimplification, PhaseContraction]
, not (Foldable.any (`varNameInAst` ik) vars) -> True
-- We can't reorder ix for the gather(fromVector) rule above,
-- because it becomes gather(transpose); we can only reorder vars.
_ -> False
, not (intInteresting i1) -- now vars may need to be reordered, too
, Just i <- findIndex intInteresting
(Foldable.toList ix) = assert (i > 0) $
Permutation.permFromListCont (backpermCycle $ i + 1)
$ \(perm :: Permutation.Perm perm) ->
gcastWith (unsafeCoerceRefl
:: Permutation.PermutePrefix perm (shp ++ shn)
:~: Permutation.PermutePrefix perm shp ++ shn) $
gcastWith (unsafeCoerceRefl :: (Rank perm <=? Rank (shp ++ shn)) :~: True) $
fromMaybe (error "astGatherKnobsS: impossible non-permutation")
$ Permutation.permCheckPermutation perm
$ let v2 = astTransposeS perm v0
in astGatherKnobsS knobs shm shn (shsPermutePrefix perm shp)
v2 (vars, ixsPermutePrefix perm ix)
-- this call is guaranteed to simplify as above, so the transpose
-- won't reduce it back to the original and cause a loop
astGatherKnobsS knobs shm@(m :$$ _) shn shp v0
(vars, ix@(i1 :.$ prest))
| let varInteresting = \case
Ast.AstCondK (AstLeqInt (AstConcreteK j0) (AstIntVar var)) _ _
| j0 <= 0 || j0 >= fromSNat' m || ixIsSmall prest ->
Just var
Ast.AstLet varN _
(Ast.AstCondK (AstLeqInt (AstConcreteK j0) (AstIntVar var)) v w)
| let varIn1 = varN `varNameInAst` v
varIn2 = varN `varNameInAst` w
, j0 <= 0 || j0 >= fromSNat' m
|| ixIsSmall prest && not (varIn1 && varIn2) ->
Just var
Ast.AstCondK
(Ast.AstBoolAndK
(AstLeqInt (AstConcreteK j0) (AstIntVar var)) _) _ _
| j0 <= 0 || j0 >= fromSNat' m || ixIsSmall prest ->
Just var
Ast.AstCondK
(Ast.AstBoolAndK
(AstLeqInt (AstConcreteK j0)
(Ast.AstN1K NegateOp (AstIntVar var))) _) _ _
| - j0 + 1 <= 0 || - j0 + 1 >= fromSNat' m || ixIsSmall prest ->
Just var
Ast.AstLet varN _
(Ast.AstCondK
(Ast.AstBoolAndK
(AstLeqInt (AstConcreteK j0) (AstIntVar var)) _) v w)
| let varIn1 = varN `varNameInAst` v
varIn2 = varN `varNameInAst` w
, j0 <= 0 || j0 >= fromSNat' m
|| ixIsSmall prest && not (varIn1 && varIn2) ->
Just var
Ast.AstLet varN _
(Ast.AstCondK
(Ast.AstBoolAndK
(AstLeqInt (AstConcreteK j0)
(Ast.AstN1K NegateOp (AstIntVar var))) _) v w)
| let varIn1 = varN `varNameInAst` v
varIn2 = varN `varNameInAst` w
, - j0 + 1 <= 0 || - j0 + 1 >= fromSNat' m
|| ixIsSmall prest && not (varIn1 && varIn2) ->
Just var
AstIntVar var
| knobPhase knobs `elem` [PhaseSimplification, PhaseContraction]
, not (var `varNameInIxS` prest) -> Just var
i4 -- has to be last, because ix can't be reordered
| knobPhase knobs `elem` [PhaseSimplification, PhaseContraction]
, Ast.AstFromVectorK (_ :$$ ZSS) l <- v0
, V.length l < 100
, let g var4 i = case substituteAst (AstConcreteK i) var4 i4 of
AstConcreteK{} -> True
_ -> False
mvar = Foldable.find (\var ->
all (g var) [0 .. fromSNat' m - 1]) vars
, Just{} <- mvar -> mvar
| knobPhase knobs `elem` [PhaseSimplification, PhaseContraction]
, Ast.AstFromVectorS (_ :$$ ZSS) l <- v0
, V.length l < 100
, ixIsSmall prest
, let g var4 i = case substituteAst (AstConcreteK i) var4 i4 of
AstConcreteK{} -> True
_ -> False
mvar = Foldable.find (\var ->
all (g var) [0 .. fromSNat' m - 1]) vars
, Just{} <- mvar -> mvar
_ -> Nothing
, Just varp <- varInteresting i1
, Just i <- elemIndex varp (Foldable.toList vars) = assert (i > 0) $
Permutation.permFromListCont (backpermCycle $ i + 1)
$ \(permWhole :: Permutation.Perm permWhole) ->
permInverse permWhole $ \(invperm :: Nested.Perm invperm) _ ->
gcastWith (unsafeCoerceRefl
:: shm ++ shn
:~: Permutation.PermutePrefix permWhole
(Permutation.PermutePrefix invperm shm ++ shn)) $
gcastWith (unsafeCoerceRefl
:: (Rank permWhole
<=? Rank (Permutation.PermutePrefix invperm shm ++ shn))
:~: True) $
fromMaybe (error "astGatherKnobsS: impossible non-permutation")
$ Permutation.permCheckPermutation permWhole
$ astTransposeS permWhole
$ astGatherKnobsS knobs (shsPermutePrefix invperm shm) shn shp
v0 (AstVarListS
$ ixsPermutePrefix invperm (unAstVarListS vars), ix)
-- this call is guaranteed to simplify as above, so the transpose
-- won't reduce it back to the original and cause a loop
astGatherKnobsS knobs shm shn shp@(SNat @in1 :$$ (shp1 :: ShS shp1))
v4 (vars4, ix4@(i4 :.$ rest4))
| FTKS _ x <- ftkAst v4
, let tryRecursing :: AstTensor AstMethodLet s (TKS2 (shp1 ++ shn) x)
-> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)
tryRecursing u =
let ftk = FTKS (shm `shsAppend` shn) x
defArr = fromPlain $ astConcrete ftk (tdefTarget ftk)
-- This boolean term may have free variables that act as
-- universally quantified.
in case 0 <=. i4 &&* i4 <=. valueOf @in1 - 1 of
AstConcreteK b ->
if b then astGather shm shn shp1 u (vars4, rest4) else defArr
_ -> Ast.AstGatherS shm shn shp v4 (vars4, ix4) = case v4 of
Ast.AstProject1{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstProject2{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
{- Ast.AstFromVector{} | gatherFromNF (shsTail shp) vars4 ix4 ->
-- normal form
Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4) -}
{- this rule seems counterproductive in many cases, so disabled until
we can detect cases where it helps:
Ast.AstFromVector snat STKS{} l ->
-- Term rest4 is duplicated without sharing and we can't help it,
-- because it needs to be in scope of vars4, so we can't use tlet.
funToVarsIxS @shm shm $ \varsFresh (IxS !ixFresh) ->
let f v = astGather @shm @shn @shp1 shn v (vars4, rest4)
-- This subst doesn't currently break sharing because it's a rename.
subst i =
Foldable.foldr (\(i2, var2) v2 -> substituteAst i2 var2 v2)
i
(zipSizedS ixFresh vars4)
i5 = subst i4
in astGather @shm @shn @(p1' ': shm)
shn (astFromVector snat (STKS (shsFromListS varsFresh
`shsAppend` shn)
(ftkToSTK x))
$ V.map f l)
(varsFresh, i5 :.$ IxS ixFresh) -}
Ast.AstApply{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstVar{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstBuild1 _ STKScalar (var2, v) | ZIS <- rest4 ->
tryRecursing $ sfromK $ astLet var2 i4 v
Ast.AstBuild1 _ STKS{} (var2, v) ->
tryRecursing $ astLet var2 i4 v
Ast.AstLet var u v ->
astLetRefresh var u v $ \v' -> astGather shm shn shp v' (vars4, ix4)
Ast.AstPrimalPart{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstDualPart{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstPlainPart{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstFromPrimal v -> fromPrimal $ astGather shm shn shp v (vars4, ix4)
Ast.AstFromDual v -> fromDual $ astGather shm shn shp v (vars4, ix4)
Ast.AstFromPlain v -> fromPlain $ astGather shm shn shp v (vars4, ix4)
-- Going inside a binary ops usually makes a term more expensive
-- to interpret and inverting that requires comparing two arguments,
-- so it's not practical.
AstPlusS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
AstTimesS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstN1S{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstR1S{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstR2S{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstI2S{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
AstConcreteS{} -> case unRepl1 v4 of
Just u -> tryRecursing u
_ -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
-- free variables possible in the index, so can't compute the array
Ast.AstFloorS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstFromIntegralS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstCastS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstArgMinS @n1 @shz v | FTKS nsh _ <- ftkAst v
, nl@(SNat @nl) <- shsLast nsh ->
let shnl = shn `shsAppend` (nl :$$ ZSS)
in gcastWith (unsafeCoerceRefl
:: Head (shm ++ (shn ++ '[nl]))
': Tail (shm ++ (shn ++ '[nl]))
:~: shm ++ (shn ++ '[nl])) $
gcastWith (unsafeCoerceRefl
:: Init (shm ++ (shn ++ '[nl])) :~: shm ++ shn) $
gcastWith (unsafeCoerceRefl
:: shp ++ (shn ++ '[nl]) :~: n1 ': shz) $
Ast.AstArgMinS @(Head (shm ++ (shn ++ '[nl])))
@(Tail (shm ++ (shn ++ '[nl])))
$ astGather shm shnl shp v (vars4, ix4)
Ast.AstArgMaxS @n1 @shz v | FTKS nsh _ <- ftkAst v
, nl@(SNat @nl) <- shsLast nsh ->
let shnl = shn `shsAppend` (nl :$$ ZSS)
in gcastWith (unsafeCoerceRefl
:: Head (shm ++ (shn ++ '[nl]))
': Tail (shm ++ (shn ++ '[nl]))
:~: shm ++ (shn ++ '[nl])) $
gcastWith (unsafeCoerceRefl
:: Init (shm ++ (shn ++ '[nl])) :~: shm ++ shn) $
gcastWith (unsafeCoerceRefl
:: shp ++ (shn ++ '[nl]) :~: n1 ': shz) $
Ast.AstArgMaxS @(Head (shm ++ (shn ++ '[nl])))
@(Tail (shm ++ (shn ++ '[nl])))
$ astGather shm shnl shp v (vars4, ix4)
{- is reverted in astGatherKnobsS immediatedly; only do in expansion phase?
Ast.AstIndexS @shm2 _shn2 v2 (i2 :.$ ZIS) ->
astGather @shm @shn @(shm2 ++ shp) shn v2 (vars4, i2 :.$ ix4) -}
Ast.AstIndexS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstCondS b v w | ixIsSmall ix4 ->
astCondS b (astGather shm shn shp v (vars4, ix4))
(astGather shm shn shp w (vars4, ix4))
Ast.AstCondS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstFromVectorK{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstFromVectorS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
-- This accomplishes fusion if v is a gather or anything
-- that gather can fuse with, but at the cost of an extra transpose
-- that doesn't fuse here unless astTransposeAsGatherS is used.
-- Since the transpose is O(1), let's leave this as is.
Ast.AstSumS @shm2 shm2 v ->
let perm3 = backpermCycleN (shsLength shm2)
(shsLength shm2 + shsLength shp)
perm4 = permCycleN (shsLength shm2) (shsLength shm2 + shsLength shm)
in Permutation.permFromListCont perm3
$ \(perm3S :: Permutation.Perm perm3P) ->
gcastWith (unsafeCoerceRefl
:: (Rank perm3P <=? Rank (shm2 ++ (shp ++ shn))) :~: True) $
gcastWith (unsafeCoerceRefl
:: Permutation.PermutePrefix perm3P (shm2 ++ (shp ++ shn))
:~: shp ++ (shm2 ++ shn)) $
fromMaybe (error "astGatherKnobsS: impossible non-permutation")
$ Permutation.permCheckPermutation perm3S
$ Permutation.permFromListCont perm4
$ \(perm4S :: Permutation.Perm perm4P) ->
gcastWith (unsafeCoerceRefl
:: (Rank perm4P <=? Rank (shm ++ (shm2 ++ shn))) :~: True) $
gcastWith (unsafeCoerceRefl
:: Permutation.PermutePrefix perm4P (shm ++ (shm2 ++ shn))
:~: shm2 ++ (shm ++ shn)) $
fromMaybe (error "astGatherKnobsS: impossible non-permutation")
$ Permutation.permCheckPermutation perm4S
$ let innerGather =
astGather shm (shm2 `shsAppend` shn) shp
(astTransposeS perm3S v) (vars4, ix4)
in astSumS shm2
$ astTransposeS perm4S innerGather
{- TODO: disabled until we can reliably fuse back to transpose
if not (knobExpand knobs)
then astTransposeS perm4S innerGather
else astTransposeAsGatherS knobs perm4S innerGather -}
Ast.AstScatterS shm7 shn7 shp7 v (vars, i5 :.$ ix2) | eqY i5 i4 ->
astGather shm shn (shsTail shp)
(astScatterKnobsS knobs shm7 shn7 (shsTail shp7) v (vars, ix2))
(vars4, rest4)
Ast.AstScatterS _ _ _ _ (_, AstConcreteK{} :.$ _)
| AstConcreteK{} <- i4 -> -- from above we know i5 /= i4
let ftk = FTKS (shm `shsAppend` shn) x
in fromPlain $ astConcrete ftk (tdefTarget ftk)
Ast.AstScatterS{} -> -- normal form
Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstReplicateK (_ :$$ shmRest) v ->
tryRecursing $ astReplicateK shmRest v
Ast.AstReplicateS ZSS v -> astGather shm shn shp v (vars4, ix4)
Ast.AstReplicateS (_ :$$ shmRest) v ->
tryRecursing $ astReplicateS shmRest v
Ast.AstGatherS @shm2 @shn2 @shp2 shm2 shn2 shp2 v2 (vars2, ix2)
| SNat @rank4 <- ixsRank ix4
, SNat @rank2 <- ixsRank (unAstVarListS vars2) ->
let subst :: AstIxS AstMethodLet shm7 -> AstVarListS shm7
-> AstInt AstMethodLet
-> AstInt AstMethodLet
subst ix vars t0 =
foldr (\ (v, i) -> substituteAst i v)
t0 (ixsZip (unAstVarListS vars) ix)
inBounds :: AstIxS AstMethodLet shm7 -> AstVarListS shm7 -> Bool
inBounds ix vars =
let inb (v, i) | Just (lbv, ubv) <- varNameToBounds v
, Just (lbi, ubi) <- intBounds i =
lbv <= lbi && ubi <= ubv
inb _ = True
in all inb (ixsZip (unAstVarListS vars) ix)
composedGather :: -- rank4 <= rank2
Maybe (AstTensor AstMethodLet s (TKS2 (shm ++ shn) x))
composedGather | SNat <- shsRank shm =
-- we have: shm2 ++ shn2 == shp ++ shn
-- so from ranks:
gcastWith (unsafeCoerceRefl :: Take (Rank shp) shm2 :~: shp) $
-- and from congruence:
-- gcastWith (unsafeCoerceRefl
-- :: Drop (Rank shp) shm2 ++ shn2 :~: shn) $
-- from congruence:
gcastWith (unsafeCoerceRefl
:: (shm ++ Drop (Rank shp) shm2) ++ shn2
:~: shm ++ shn) $
let vars2p = AstVarListS $ ixsTake @(Rank shp) $ unAstVarListS vars2
vars22 = AstVarListS $ ixsDrop @(Rank shp) $ unAstVarListS vars2
ix22 = fmap (subst ix4 vars2p) ix2
list422 = AstVarListS
$ unAstVarListS vars4 `ixsAppend` unAstVarListS vars22
in if ixIsSmall ix4 && inBounds ix4 vars2p
then Just $ astGather (shm `shsAppend` shsDrop @(Rank shp) shm2)
shn2
shp2
v2 (list422, ix22)
else Nothing
assimilatedGather :: -- rank2 <= rank4
Maybe (AstTensor AstMethodLet s (TKS2 (shm ++ shn) x))
assimilatedGather | SNat <- shsRank shm =
-- we have: shm2 ++ shn2 == shp ++ shn
-- so from ranks:
gcastWith (unsafeCoerceRefl :: Take (Rank shm2) shp :~: shm2) $
-- and from congruence:
-- gcastWith (unsafeCoerceRefl
-- :: Drop (Rank shm2) shp ++ shn :~: shn2) $
-- from congruence:
gcastWith (unsafeCoerceRefl
:: (shp2 ++ Drop (Rank shm2) shp) ++ shn
:~: shp2 ++ shn2) $
let ix42 = ixsTake @(Rank shm2) ix4
ix44 = ixsDrop @(Rank shm2) ix4
ix22 = fmap (subst ix42 vars2) ix2
ix2244 = ix22 `ixsAppend` ix44
in if ixIsSmall ix42 && inBounds ix42 vars2
then Just $ astGather shm
shn
(shp2 `shsAppend` shsDrop @(Rank shm2) shp)
v2 (vars4, ix2244)
else Nothing
in fromMaybe (Ast.AstGatherS shm shn shp v4 (vars4, ix4))
$ case cmpNat (Proxy @rank4) (Proxy @rank2) of
LTI -> composedGather
EQI -> assimilatedGather
GTI -> gcastWith (flipCompare (Proxy @rank4) (Proxy @rank2))
assimilatedGather
Ast.AstIotaS{} -> -- probably nothing can be simplified; a normal form
Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstAppendS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
-- fusing would result in gather([gather, gather]), so no gain
Ast.AstSliceS{}-> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
-- slicing is O(1) so no point fusing and complicating the expression;
-- if it did not simplify further with slice, it wouldn't with gather
Ast.AstReverseS{}-> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
-- reversing is O(1)
Ast.AstTransposeS @perm @sh perm v
| knobPhase knobs `elem` [PhaseVectorization, PhaseExpansion]
, fromSNat' (Permutation.permRank perm) <= length ix4
, FTKS sh _ <- ftkAst v
, Refl <- lemRankMapJust $ shsTakeLenPerm perm sh ->
gcastWith (unsafeCoerceRefl :: Rank (TakeLen perm sh) :~: Rank perm) $
permInverse perm
$ \(invperm :: Nested.Perm invperm) proof ->
case proof (ssxFromShX $ shxFromShS $ shsTakeLenPerm perm sh) of
Refl ->
-- from PermutePrefix and ranks:
gcastWith
(unsafeCoerceRefl
:: Permutation.PermutePrefix invperm shp ++ shn
:~: Permutation.PermutePrefix invperm (shp ++ shn)) $
-- from AstTransposeS:
-- gcastWith
-- (unsafeCoerceRefl
-- :: Permutation.PermutePrefix invperm (shp ++ shn)
-- :~: Permutation.PermutePrefix invperm (Permutation.PermutePrefix perm sh)) $
-- from PermutePrefix:
-- gcastWith
-- (unsafeCoerceRefl
-- :: Permutation.PermutePrefix invperm (Permutation.PermutePrefix perm sh)
-- :~: {-1-} Permutation.Permute invperm (TakeLen invperm (Permutation.PermutePrefix perm sh))
-- ++ {-2-} DropLen invperm (Permutation.PermutePrefix perm sh)) $
-- 1. from PermutePrefix:
-- gcastWith
-- (unsafeCoerceRefl
-- :: Permutation.Permute invperm (TakeLen invperm (Permutation.PermutePrefix perm sh))
-- :~: Permutation.Permute invperm (TakeLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh))) $
-- ranks
gcastWith
(unsafeCoerceRefl
:: Permutation.Permute invperm (TakeLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh))
:~: Permutation.Permute invperm (Permutation.Permute perm (TakeLen perm sh))) $
-- from permInverse but MapJust-unwrapped:
gcastWith
(unsafeCoerceRefl
:: Permutation.Permute invperm (Permutation.Permute perm (TakeLen perm sh))
:~: TakeLen perm sh) $
-- end of 1.
-- 2. from PermutePrefix
-- gcastWith
-- (unsafeCoerceRefl
-- :: DropLen invperm (Permutation.PermutePrefix perm sh)
-- :~: DropLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh)) $
-- ranks
gcastWith
(unsafeCoerceRefl
:: DropLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh)
:~: DropLen perm sh) $
-- end of 2.
-- from TakeLen:
gcastWith
(unsafeCoerceRefl
:: TakeLen perm sh ++ DropLen perm sh :~: sh) $
let invix4 = ixsPermutePrefix invperm ix4
in astGather shm shn (shsPermutePrefix invperm shp)
v (vars4, invix4)
Ast.AstTransposeS perm v
| knobPhase knobs `elem` [PhaseVectorization, PhaseExpansion] ->
astGather shm shn shp (astTransposeAsGatherS knobs perm v) (vars4, ix4)
Ast.AstTransposeS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstReshapeS sh v ->
if shsLength sh <= 1 -- this is flatten
then astGather shm shn shp (astReshapeAsGatherS knobs sh v) (vars4, ix4)
else Ast.AstGatherS shm shn shp v4 (vars4, ix4)
-- Conversions to shaped need to stay down, so this is NF.
Ast.AstConvert{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstDot1InS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4) -- TODO
Ast.AstMatmul2S{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4) -- TODO
Ast.AstBoolNotS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstBoolAndS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
Ast.AstLeqS{} -> Ast.AstGatherS shm shn shp v4 (vars4, ix4)
where
astGather
:: forall shm' shn' shp' s' x'. KnownSpan s'
=> ShS shm' -> ShS shn' -> ShS shp'
-> AstTensor AstMethodLet s' (TKS2 (shp' ++ shn') x')
-> (AstVarListS shm', AstIxS AstMethodLet shp')
-> AstTensor AstMethodLet s' (TKS2 (shm' ++ shn') x')
astGather = astGatherKnobsS knobs
-- Normal form of chains of appends has the append constructor on the right.
astAppendS :: KnownSpan s
=> AstTensor AstMethodLet s (TKS2 (m ': sh) x)
-> AstTensor AstMethodLet s (TKS2 (n ': sh) x)
-> AstTensor AstMethodLet s (TKS2 ((m + n) ': sh) x)
astAppendS u v | FTKS (SZ :$$ _) _ <- ftkAst u = v
astAppendS u v | FTKS (SZ :$$ _) _ <- ftkAst v = u
astAppendS (Ast.AstFromVectorS (SNat @k1 :$$ ZSS) l1)
(Ast.AstFromVectorS (SNat @k2 :$$ ZSS) l2) =
astFromVectorS (SNat @(k1 + k2) :$$ ZSS) $ l1 V.++ l2
astAppendS (Ast.AstReplicateS (SNat' @1 :$$ shmRest) a1)
(Ast.AstFromVectorS (SNat @k2 :$$ ZSS) l2) =
astFromVectorS (SNat @(1 + k2) :$$ ZSS) $ astReplicateS shmRest a1 `V.cons` l2
astAppendS (Ast.AstReplicateK (SNat' @1 :$$ shmRest) a1)
(Ast.AstFromVectorS (SNat @k2 :$$ ZSS) l2) =
astFromVectorS (SNat @(1 + k2) :$$ ZSS) $ astReplicateK shmRest a1 `V.cons` l2
astAppendS (Ast.AstFromVectorS (SNat @k1 :$$ ZSS) l1)
(Ast.AstReplicateS (SNat' @1 :$$ shmRest) a2) =
astFromVectorS (SNat @(k1 + 1) :$$ ZSS) $ l1 `V.snoc` astReplicateS shmRest a2
astAppendS (Ast.AstFromVectorS (SNat @k1 :$$ ZSS) l1)
(Ast.AstReplicateK (SNat' @1 :$$ shmRest) a2) =
astFromVectorS (SNat @(k1 + 1) :$$ ZSS) $ l1 `V.snoc` astReplicateK shmRest a2
astAppendS (Ast.AstFromVectorK (SNat @k1 :$$ ZSS) l1)
(Ast.AstFromVectorK (SNat @k2 :$$ ZSS) l2) =
astFromVectorK (SNat @(k1 + k2) :$$ ZSS) $ l1 V.++ l2
astAppendS (Ast.AstReplicateK (SNat' @1 :$$ ZSS) a1)
(Ast.AstFromVectorK (SNat @k2 :$$ ZSS) l2) =
astFromVectorK (SNat @(1 + k2) :$$ ZSS) $ a1 `V.cons` l2
astAppendS (Ast.AstFromVectorK (SNat @k1 :$$ ZSS) l1)
(Ast.AstReplicateK (SNat' @1 :$$ ZSS) a2) =
astFromVectorK (SNat @(k1 + 1) :$$ ZSS) $ l1 `V.snoc` a2
astAppendS (Ast.AstReplicateS (SNat' @1 :$$ shmRest1) a1)
(Ast.AstReplicateS (SNat' @1 :$$ shmRest2) a2) =
astFromVectorS (SNat @2 :$$ ZSS)
$ V.fromList [astReplicateS shmRest1 a1, astReplicateS shmRest2 a2]
astAppendS (Ast.AstReplicateK (SNat' @1 :$$ shmRest1) a1)
(Ast.AstReplicateK (SNat' @1 :$$ shmRest2) a2) =
astFromVectorS (SNat @2 :$$ ZSS)
$ V.fromList [astReplicateK shmRest1 a1, astReplicateK shmRest2 a2]
astAppendS (Ast.AstFromVectorS (SNat @k1 :$$ ZSS) l1)
(Ast.AstAppendS (Ast.AstFromVectorS (SNat @k2 :$$ ZSS) l2) w) =
astAppendS (astFromVectorS (SNat @(k1 + k2) :$$ ZSS) $ l1 V.++ l2) w
astAppendS (Ast.AstReplicateS (SNat' @1 :$$ shmRest) a1)
(Ast.AstAppendS (Ast.AstFromVectorS (SNat @k2 :$$ ZSS) l2) w) =
astAppendS (astFromVectorS (SNat @(1 + k2) :$$ ZSS)
$ astReplicateS shmRest a1 `V.cons` l2) w
astAppendS (Ast.AstReplicateK (SNat' @1 :$$ shmRest) a1)
(Ast.AstAppendS (Ast.AstFromVectorS (SNat @k2 :$$ ZSS) l2) w) =
astAppendS (astFromVectorS (SNat @(1 + k2) :$$ ZSS)
$ astReplicateK shmRest a1 `V.cons` l2) w
astAppendS (Ast.AstFromVectorS (SNat @k1 :$$ ZSS) l1)
(Ast.AstAppendS
(Ast.AstReplicateS (SNat' @1 :$$ shmRest) a2) w) =
astAppendS (astFromVectorS (SNat @(k1 + 1) :$$ ZSS)
$ l1 `V.snoc` astReplicateS shmRest a2) w
astAppendS (Ast.AstFromVectorS (SNat @k1 :$$ ZSS) l1)
(Ast.AstAppendS
(Ast.AstReplicateK (SNat' @1 :$$ shmRest) a2) w) =
astAppendS (astFromVectorS (SNat @(k1 + 1) :$$ ZSS)
$ l1 `V.snoc` astReplicateK shmRest a2) w
astAppendS (Ast.AstFromVectorK (SNat @k1 :$$ ZSS) l1)
(Ast.AstAppendS (Ast.AstFromVectorK (SNat @k2 :$$ ZSS) l2) w) =
astAppendS (astFromVectorK (SNat @(k1 + k2) :$$ ZSS) $ l1 V.++ l2) w
astAppendS (Ast.AstReplicateK (SNat' @1 :$$ ZSS) a1)
(Ast.AstAppendS (Ast.AstFromVectorK (SNat @k2 :$$ ZSS) l2) w) =
astAppendS (astFromVectorK (SNat @(1 + k2) :$$ ZSS) $ a1 `V.cons` l2) w
astAppendS (Ast.AstFromVectorK (SNat @k1 :$$ ZSS) l1)
(Ast.AstAppendS (Ast.AstReplicateK (SNat' @1 :$$ ZSS) a2) w) =
astAppendS (astFromVectorK (SNat @(k1 + 1) :$$ ZSS) $ l1 `V.snoc` a2) w
astAppendS (Ast.AstReplicateS (SNat' @1 :$$ shmRest1) a1)
(Ast.AstAppendS (Ast.AstReplicateS (SNat' @1 :$$ shmRest2) a2) w) =
astAppendS
(astFromVectorS (SNat @2 :$$ ZSS)
$ V.fromList [astReplicateS shmRest1 a1, astReplicateS shmRest2 a2]) w
astAppendS (Ast.AstReplicateK (SNat' @1 :$$ shmRest1) a1)
(Ast.AstAppendS (Ast.AstReplicateK (SNat' @1 :$$ shmRest2) a2) w) =
astAppendS
(astFromVectorS (SNat @2 :$$ ZSS)
$ V.fromList [astReplicateK shmRest1 a1, astReplicateK shmRest2 a2]) w
astAppendS (Ast.AstFromPrimal u) (Ast.AstFromPrimal v) =
fromPrimal $ astAppendS u v
astAppendS (Ast.AstFromDual u) (Ast.AstFromDual v) =
fromDual $ astAppendS u v
astAppendS (Ast.AstFromPlain u) (Ast.AstFromPlain v) =
fromPlain $ astAppendS u v
astAppendS u v | Just u0 <- unAstS u
, Just v0 <- unAstS v
, FTKS _ FTKScalar <- ftkAst u =
fromPlain $ astConcreteS (tsappend (Concrete u0) (Concrete v0))
astAppendS u (Ast.AstAppendS v w) | Just u0 <- unAstS u
, Just v0 <- unAstS v
, FTKS _ FTKScalar <- ftkAst u =
astAppendS (fromPlain $ astConcreteS (tsappend (Concrete u0) (Concrete v0))) w
astAppendS (Ast.AstAppendS v u) w = astAppendS v (astAppendS u w)
astAppendS u v = Ast.AstAppendS u v
astSliceS :: forall i n k sh s x. KnownSpan s
=> SNat i -> SNat n -> SNat k
-> AstTensor AstMethodLet s (TKS2 (i + n + k ': sh) x)
-> AstTensor AstMethodLet s (TKS2 (n ': sh) x)
astSliceS SZ SNat SZ v = v
astSliceS _ SZ _ v | FTKS (_ :$$ sh) x <- ftkAst v =
let ftk = FTKS (SZ :$$ sh) x
in fromPlain $ astConcrete ftk (tdefTarget ftk)
astSliceS SNat SNat SNat (Ast.AstFromVectorK (_ :$$ ZSS) l) =
astFromVectorK (SNat @n :$$ ZSS) $ V.take (valueOf @n) $ V.drop (valueOf @i) l
astSliceS SNat SNat SNat (Ast.AstFromVectorS (_ :$$ ZSS) l) =
astFromVectorS (SNat @n :$$ ZSS) $ V.take (valueOf @n) $ V.drop (valueOf @i) l
astSliceS SNat (SNat' @1) SNat v | FTKS (_ :$$ sh) _ <- ftkAst v =
astReplicateS (SNat @1 :$$ ZSS) (astIndexS sh v (valueOf @i :.$ ZIS))
astSliceS SNat SNat SNat (Ast.AstReplicateK (_ :$$ shmRest) v) =
astReplicateK (SNat @n :$$ shmRest) v
astSliceS SNat SNat SNat (Ast.AstReplicateS (_ :$$ shmRest) v) =
astReplicateS (SNat @n :$$ shmRest) v
astSliceS SNat SNat SNat (Ast.AstGatherS shm shn shp v (var ::$ vars, ix)) =
let varn = reboundsVarName (0, valueOf @n - 1) var
ivar = valueOf @i + astVar varn
ix2 = substituteAstIxS ivar var ix -- cheap subst, because ivar is tiny
in astGatherS (SNat @n :$$ shsTail shm) shn shp v (varn ::$ vars, ix2)
astSliceS i n@(SNat @n0) _k (Ast.AstAppendS v1 v2)
| FTKS (m1@(SNat @m1) :$$ _) _ <- ftkAst v1
, FTKS (m2@(SNat @m2) :$$ _) _ <- ftkAst v2 =
let i1 = fromSNat' i `min` fromSNat' m1
n1 = fromSNat' n `min` (fromSNat' m1 - i1)
k1 = fromSNat' m1 - i1 - n1
i2' = fromSNat' i `max` fromSNat' m1
i2 = i2' - fromSNat' m1
n2 = fromSNat' n - n1
k2 = fromSNat' m2 - i2 - n2
in withSNat i1 $ \si1@(SNat @i1) ->
withSNat n1 $ \sn1@(SNat @n1) ->
withSNat k1 $ \sk1@(SNat @k1) ->
withSNat i2 $ \si2@(SNat @i2) ->
withSNat n2 $ \sn2@(SNat @n2) ->
withSNat k2 $ \sk2@(SNat @k2) ->
gcastWith (unsafeCoerceRefl :: n1 + n2 :~: n0) $
gcastWith (unsafeCoerceRefl :: i1 + n1 + k1 :~: m1) $
gcastWith (unsafeCoerceRefl :: i2 + n2 + k2 :~: m2) $
astAppendS (astSliceS si1 sn1 sk1 v1) (astSliceS si2 sn2 sk2 v2)
astSliceS i n k (Ast.AstSliceS i2 _n2 k2 v) =
astSliceS (snatPlus i i2) n (snatPlus k k2) v
astSliceS i n k (Ast.AstReverseS v) = astReverseS (astSliceS k n i v)
-- This enlarges the term and increases computation, but sometimes
-- it permits eliminating the AstFromVectorS node altogether, so we risk it
-- for cases that commonly emerge from conditionals.
astSliceS i n@SNat k (Ast.AstTransposeS
perm@(SNat' @1 `PCons` SNat' @0 `PCons` PNil)
t@(Ast.AstFromVectorS (_ :$$ ZSS) l))
| FTKS (SNat' @2 :$$ _ :$$ (_ :: ShS sh2)) _ <- ftkAst t =
Ast.AstTransposeS perm
$ astFromVectorS @_ @(n : sh2) (SNat @2 :$$ ZSS) (V.map (astSliceS i n k) l)
-- TODO: generalize (maybe the above, too) using unReplN, but it's hard.
-- TODO: does it really work only for replicate-like things in-between?
astSliceS i n@SNat k
(Ast.AstTransposeS perm@(SNat' @1 `PCons` SNat' @0 `PCons` PNil) t)
| STKS (snat :$$ _ :$$ _) _ <- ftkToSTK $ ftkAst t
, Just u <- unRepl1 t =
astTransposeS perm $ astReplicateS (snat :$$ ZSS) $ astSliceS i n k u
astSliceS i n k v1 = case v1 of
Ast.AstLet var u v -> astLet var u (astSliceS i n k v)
Ast.AstFromPrimal v -> fromPrimal $ astSliceS i n k v
Ast.AstFromDual v -> fromDual $ astSliceS i n k v
Ast.AstFromPlain v -> fromPlain $ astSliceS i n k v
AstPlusS u v -> astSliceS i n k u `astPlusS` astSliceS i n k v
AstTimesS u v -> astSliceS i n k u `astTimesS` astSliceS i n k v
Ast.AstN1S opCode u -> astN1S opCode (astSliceS i n k u)
Ast.AstR1S opCode u -> astR1S opCode (astSliceS i n k u)
Ast.AstR2S opCode u v -> astR2S opCode (astSliceS i n k u)
(astSliceS i n k v)
Ast.AstI2S opCode u v -> astI2S opCode (astSliceS i n k u)
(astSliceS i n k v)
Ast.AstFloorS a -> astFloorS $ astSliceS i n k a
Ast.AstFromIntegralS v -> astFromIntegralS $ astSliceS i n k v
Ast.AstCastS v -> astCastS $ astSliceS i n k v
AstConcreteS v -> astConcreteS (tsslice i n k $ Concrete v)
Ast.AstCondS b a2 a3 -> astCondS b (astSliceS i n k a2) (astSliceS i n k a3)
_ -> Ast.AstSliceS i n k v1
astReverseS :: forall n sh s r. KnownSpan s
=> AstTensor AstMethodLet s (TKS2 (n ': sh) r)
-> AstTensor AstMethodLet s (TKS2 (n ': sh) r)
astReverseS (Ast.AstLet var u v) = astLet var u (astReverseS v)
astReverseS (Ast.AstFromPrimal v) = fromPrimal $ astReverseS v
astReverseS (Ast.AstFromDual v) = fromDual $ astReverseS v
astReverseS (Ast.AstFromPlain v) = fromPlain $ astReverseS v
astReverseS (Ast.AstCondS b a2 a3) =
astCondS b (astReverseS a2) (astReverseS a3)
astReverseS (Ast.AstFromVectorK shm@(_ :$$ ZSS) l) =
astFromVectorK shm $ V.reverse l
astReverseS (Ast.AstFromVectorS shm@(_ :$$ ZSS) l) =
astFromVectorS shm $ V.reverse l
astReverseS (Ast.AstReplicateK shm v) = astReplicateK shm v
astReverseS (Ast.AstReplicateS shm v) = astReplicateS shm v
astReverseS (Ast.AstGatherS shm@(SNat @k :$$ _) shn shp v (var ::$ vars, ix)) =
let ivar = valueOf @k - 1 - astVar var
ix2 = substituteAstIxS ivar var ix -- cheap subst, because ivar is tiny
in astGatherS shm shn shp v (var ::$ vars, ix2)
astReverseS (Ast.AstReverseS v) = v
astReverseS (Ast.AstTransposeS perm@(SNat' @1 `PCons` SNat' @0 `PCons` PNil) t)
| STKS (snat :$$ (_ :$$ _)) _ <- ftkToSTK $ ftkAst t
, Just u <- unRepl1 t =
astTransposeS perm $ astReplicateS (snat :$$ ZSS) $ astReverseS u
astReverseS (AstConcreteS v) = astConcreteS (tsreverse $ Concrete v)
astReverseS v = Ast.AstReverseS v
-- TODO: try to completely cover the AstGatherS case here, which would permit
-- not expanding to astTransposeAsGatherS in astGatherCase
-- | Beware, this does not do full simplification, which often requires
-- the gather form, so astTransposeAsGather needs to be called in addition
-- if full simplification is required.
astTransposeS
:: forall perm sh s r.
(Permutation.IsPermutation perm, Rank perm <= Rank sh, KnownSpan s)
=> Permutation.Perm perm -> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (Permutation.PermutePrefix perm sh) r)
astTransposeS perm t =
gcastWith (unsafeCoerceRefl
:: TakeLen perm sh ++ DropLen perm sh :~: sh) $ case perm of
PNil -> t
PCons (SNat' @0) PNil ->
gcastWith (unsafeCoerceRefl :: Permutation.PermutePrefix '[0] sh :~: sh)
t
_ | FTKS sh _ <- ftkAst t
, Just u2 <- unReplN @_ @(DropLen perm sh) (shsTakeLenPerm perm sh) t ->
astReplicateS (shsPermute perm (shsTakeLenPerm perm sh)) u2
_ -> case t of
_ | SNat' @0 `PCons` _ <- perm
, STKS ((:$$) @n @sh2 snat _) _ <- ftkToSTK $ ftkAst t
, Just u <- unRepl1 t -> case permUnShift1 perm of
(perm2 :: Permutation.Perm perm2) ->
fromMaybe (error "astTransposeS: impossible non-permutation")
$ Permutation.permCheckPermutation perm2
$ gcastWith (unsafeCoerceRefl :: Rank perm2 + 1 :~: Rank perm)
-- for GHC 9.10 only:
$ gcastWith (unsafeCoerceRefl :: (Rank perm2 <=? Rank sh2) :~: True)
$ gcastWith (unsafeCoerceRefl
:: Permutation.PermutePrefix perm (n : sh2)
:~: n : Permutation.PermutePrefix perm2 sh2)
$ astReplicateS (snat :$$ ZSS) (astTransposeS perm2 u)
Ast.AstLet var u v -> astLet var u (astTransposeS perm v)
Ast.AstFromPrimal v -> fromPrimal $ astTransposeS perm v
Ast.AstFromDual v -> fromDual $ astTransposeS perm v
Ast.AstFromPlain v -> fromPlain $ astTransposeS perm v
AstPlusS u v -> astTransposeS perm u `astPlusS` astTransposeS perm v
AstTimesS u v -> astTransposeS perm u `astTimesS` astTransposeS perm v
Ast.AstN1S opCode u -> astN1S opCode (astTransposeS perm u)
Ast.AstR1S opCode u -> astR1S opCode (astTransposeS perm u)
Ast.AstR2S opCode u v ->
astR2S opCode (astTransposeS perm u) (astTransposeS perm v)
Ast.AstI2S opCode u v ->
astI2S opCode (astTransposeS perm u) (astTransposeS perm v)
AstConcreteS v -> astConcreteS (tstranspose perm $ Concrete v)
-- Ast.AstFloorS v -> astFloorS $ astTransposeS perm v
-- Ast.AstFromIntegralS v -> astFromIntegralS $ astTransposeS perm v
-- Ast.AstCastS v -> astCastS $ astTransposeS perm v
Ast.AstIndexS @shm shn v ix | n <- ixsRank ix ->
Permutation.permFromListCont
(Permutation.permToList'
$ iterate (unsafeCoerce Permutation.permShift1) perm
!! fromSNat' n) -- this has a fake type, but that's fine
$ \ (permn :: Perm permn) ->
fromMaybe (error "astTransposeS: impossible non-permutation")
$ Permutation.permCheckPermutation permn
$ gcastWith (unsafeCoerceRefl
:: Permutation.PermutePrefix permn (shm ++ sh)
:~: shm ++ Permutation.PermutePrefix perm sh)
-- should suffice, but it doesn't
-- $ gcastWith (unsafeCoerceRefl :: Rank permn :~: n + Rank perm)
$ gcastWith (unsafeCoerceRefl
:: (Rank permn <=? Rank (shm ++ sh)) :~: True)
$ astIndexS (shsPermutePrefix perm shn) (astTransposeS permn v) ix
Ast.AstCondS b u v -> astCondS b (astTransposeS perm u) (astTransposeS perm v)
Ast.AstFromVectorS @_ @sh2 shm@(SNat @n :$$ ZSS) l
| SNat' @0 `PCons` _ <- perm -> case permUnShift1 perm of
(perm2 :: Permutation.Perm perm2) ->
fromMaybe (error "astTransposeS: impossible non-permutation")
$ Permutation.permCheckPermutation perm2
$ gcastWith (unsafeCoerceRefl :: Rank perm2 + 1 :~: Rank perm)
-- for GHC 9.10 only:
$ gcastWith (unsafeCoerceRefl :: (Rank perm2 <=? Rank sh2) :~: True)
$ gcastWith (unsafeCoerceRefl
:: Permutation.PermutePrefix perm (n : sh2)
:~: n : Permutation.PermutePrefix perm2 sh2)
$ astFromVectorS shm (V.map (astTransposeS perm2) l)
Ast.AstSumS @shm2 shm2 v ->
let zsuccP :: Permutation.Perm (Permutation.PermId (Rank shm2)
++ Permutation.MapPlusN (Rank shm2) perm)
zsuccP = Permutation.permShiftN (shsRank shm2) perm
in
gcastWith (unsafeCoerceRefl
:: (Rank (Permutation.PermId (Rank shm2)
++ Permutation.MapPlusN (Rank shm2) perm)
<=? Rank (shm2 ++ sh)) :~: True) $
gcastWith (unsafeCoerceRefl
:: Permutation.PermutePrefix
(Permutation.PermId (Rank shm2)
++ Permutation.MapPlusN (Rank shm2) perm) (shm2 ++ sh)
:~: shm2 ++ Permutation.PermutePrefix perm sh) $
fromMaybe (error $ "astTransposeS: impossible non-permutation: "
++ show (Permutation.permToList' zsuccP))
$ Permutation.permCheckPermutation zsuccP
$ astSumS shm2 $ astTransposeS zsuccP v
Ast.AstScatterS @_ @shn @shp shm shn shp v (vars, ix)
-- TODO: should the below be backpermute or permute?
| fromSNat' (Permutation.permRank perm) <= length ix ->
let ix2 :: AstIxS AstMethodLet (Permutation.PermutePrefix perm shp)
ix2 = ixsPermutePrefix perm ix
in gcastWith (unsafeCoerceRefl
:: Permutation.PermutePrefix perm shp ++ shn
:~: Permutation.PermutePrefix perm (shp ++ shn)) $
astScatterS shm shn (shsPermutePrefix perm shp)
v (vars, ix2)
-- This increases term size and work, so limited to size 2.
-- TODO: generalize
Ast.AstReplicateS shm@(_ :$$ ZSS)
(Ast.AstFromVectorK shm2@(SNat' @2 :$$ ZSS) l)
| SNat' @1 `PCons` SNat' @0 `PCons` PNil <- perm ->
astFromVectorS shm2 (V.map (astReplicateK shm) l)
Ast.AstReplicateS shm@(_ :$$ ZSS)
(Ast.AstFromVectorS shm2@(SNat' @2 :$$ ZSS) l)
| SNat' @1 `PCons` SNat' @0 `PCons` PNil <- perm ->
astFromVectorS shm2 (V.map (astReplicateS shm) l)
Ast.AstGatherS @shm @shn shm shn shp v (vars, ix)
-- TODO: should the below be backpermute or permute?
| fromSNat' (Permutation.permRank perm) <= length vars ->
let vars2 :: AstVarListS (Permutation.PermutePrefix perm shm)
vars2 = AstVarListS $ ixsPermutePrefix perm $ unAstVarListS vars
in gcastWith (unsafeCoerceRefl
:: Permutation.PermutePrefix perm shm ++ shn
:~: Permutation.PermutePrefix perm (shm ++ shn)) $
astGatherS (shsPermutePrefix perm shm) shn shp
v (vars2, ix)
Ast.AstAppendS u v | SNat' @0 `PCons` _ <- perm
, FTKS ((:$$) @m @sh2 _ _) _ <- ftkAst u
, FTKS ((:$$) @n _ _) _ <- ftkAst v ->
gcastWith (unsafeCoerceRefl
:: Permutation.PermutePrefix perm ((m + n) : sh2)
:~: m + n : Tail (Permutation.PermutePrefix
perm (m : sh2))) $
gcastWith (unsafeCoerceRefl
:: Permutation.PermutePrefix perm ((m + n) : sh2)
:~: m + n : Tail (Permutation.PermutePrefix
perm (n : sh2))) $
astAppendS (astTransposeS perm u) (astTransposeS perm v)
Ast.AstSliceS i n@(SNat @n) k u | SNat' @0 `PCons` _ <- perm
, FTKS ((:$$) @ink @sh2 _ _) _ <- ftkAst u ->
gcastWith (unsafeCoerceRefl
:: Permutation.PermutePrefix perm (n : sh2)
:~: n : Tail (Permutation.PermutePrefix perm (ink : sh2))) $
astSliceS i n k (astTransposeS perm u)
Ast.AstReverseS u | SNat' @0 `PCons` _ <- perm ->
astReverseS (astTransposeS perm u)
Ast.AstTransposeS @_ @sh2 perm2 u | FTKS sh2 _ <- ftkAst u ->
-- TODO: try to perform at type level
let permV = Permutation.permToList' perm
perm2V = Permutation.permToList' perm2
perm2Matched =
perm2V
++ take (length permV - length perm2V) (drop (length perm2V) [0 ..])
perm3V = normalizePermutationHack
$ backpermutePrefixList permV perm2Matched
in Permutation.permFromListCont perm3V $ \(perm3
:: Permutation.Perm perm3) ->
fromMaybe (error "astTransposeS: impossible non-permutation")
$ Permutation.permCheckPermutation perm3
$ gcastWith (unsafeCoerceRefl
:: Permutation.PermutePrefix perm3 sh2
:~: Permutation.PermutePrefix perm sh) $
case compare (length perm3V) (shsLength sh2) of
LT -> gcastWith (unsafeCoerceRefl
:: Compare (Rank perm3) (Rank sh2) :~: LT) $
astTransposeS perm3 u
EQ -> gcastWith (unsafeCoerceRefl
:: Compare (Rank perm3) (Rank sh2) :~: EQ) $
astTransposeS perm3 u
GT -> error "astTransposeS: GT"
u -> Ast.AstTransposeS perm u
-- TODO: try to cover the AstGatherS case here, which would permit
-- not expanding to astReshapeAsGatherS in astGatherCase
-- | Beware, this does not do full simplification, which often requires
-- the gather form, so astReshapeAsGather needs to be called in addition
-- if full simplification is required.
astReshapeS :: forall sh sh2 x s. (Product sh ~ Product sh2, KnownSpan s)
=> ShS sh2 -> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
astReshapeS sh2 t = case t of
_ | Just u <- unRepl t
, Refl <- lemAppNil @sh2 -> astReplicateS sh2 u
_ | STKS ((:$$) @_ @sh1 k _) _ <- ftkToSTK $ ftkAst t
, Just u <- unRepl1 t
, (:$$) @_ @rest2 k2 rest2 <- sh2
, Just Refl <- testEquality k k2 ->
gcastWith (unsafeCoerceRefl :: Product rest2 :~: Product sh1) $
astReplicateS (k :$$ ZSS) $ astReshapeS rest2 u
Ast.AstLet var u v -> astLet var u (astReshapeS @_ @sh2 sh2 v)
Ast.AstFromPrimal v -> fromPrimal $ astReshapeS sh2 v
Ast.AstFromDual v -> fromDual $ astReshapeS sh2 v
Ast.AstFromPlain v -> fromPlain $ astReshapeS sh2 v
-- Reshaping can be costly, so we don't touch AstTimesS, etc.
Ast.AstN1S opCode u -> astN1S opCode (astReshapeS @_ @sh2 sh2 u)
Ast.AstR1S opCode u -> astR1S opCode (astReshapeS @_ @sh2 sh2 u)
-- TODO: enable once we rewrite fromVectors with non-singleton shm:
-- Ast.AstFromVectorK _shm l -> astFromVectorK sh2 l
Ast.AstReshapeS _ v -> astReshapeS @_ @sh2 sh2 v
_ | FTKS sh _ <- ftkAst t -> case testEquality sh sh2 of
Just Refl -> t
_ -> Ast.AstReshapeS sh2 t
astConvert
:: KnownSpan s
=> TKConversion y z -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvert c a | yftk <- ftkAst a = case (yftk, convertFTK c yftk) of
(_, zftk) | Just Refl <- matchingFTK yftk zftk -> a
-- this covers the ConvId case and more, so not simplifying c at worst
-- causes c to take more memory but doesn't inhibit rewriting
-- Below we heavily depend on c being semantically determined
-- by the domain and codomain. We pick the simplest such c.
(FTKS ZSS (FTKScalar @ry), FTKScalar @rz)
| Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->
kfromS a
(FTKR ZSR (FTKScalar @ry), FTKScalar @rz)
| Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->
kfromR a
(FTKX ZSX (FTKScalar @ry), FTKScalar @rz)
| Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->
kfromX a
(FTKScalar @ry, FTKS ZSS (FTKScalar @rz))
| Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->
sfromK a
(FTKR shr xy, FTKS sh@(_ :$$ _) xz)
| Just Refl <- matchingFTK xy xz
, Just Refl <- testEquality (shrRank shr) (shsRank sh) ->
astConvDownSFromR sh xz a
(FTKX shx xy, FTKS sh@(_ :$$ _) xz)
| Just Refl <- matchingFTK xy xz
, Just Refl <- testEquality (shxRank shx) (shsRank sh) ->
astConvDownSFromX sh xz a
(FTKS sh@(_ :$$ _) xy, FTKR shr xz)
| Just Refl <- matchingFTK xy xz
, Just Refl <- testEquality (shrRank shr) (shsRank sh) ->
astConvUpRFromS sh xz a
(FTKS sh@(_ :$$ _) xy, FTKX shx xz)
| Just Refl <- matchingFTK xy xz
, Just Refl <- testEquality (shxRank shx) (shsRank sh) ->
astConvUpXFromS shx xz a
(_, zftk) | Just c2 <- convDownMaybe yftk (ftkToSTK zftk) ->
astConvertDown c2 zftk a
(_, zftk) | Just c2 <- convUpMaybe yftk zftk ->
astConvertUp c2 zftk a
{- This has barely any effect:
(FTKR ZSR (FTKScalar @ry), FTKS ZSS (FTKScalar @rz))
| Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->
case kfromR a of
Ast.AstConvert{} -> Ast.AstConvert c a -- would likely loop
b -> sfromK b
(FTKX ZSX (FTKScalar @ry), FTKS ZSS (FTKScalar @rz))
| Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->
case kfromX a of
Ast.AstConvert{} -> Ast.AstConvert c a -- would likely loop
b -> sfromK b -}
_ -> case a of -- normalize somewhat even for, e.g., product to product
Ast.AstFromPrimal v -> fromPrimal $ astConvert c v
Ast.AstFromDual v -> fromDual $ astConvert c v
Ast.AstFromPlain v -> fromPlain $ astConvert c v
Ast.AstConvert c2 t2 -> astConvert (c `convCmp` c2) t2
_ -> Ast.AstConvert c a
-- We are pushing conversions to shaped tensors down, into concrete values
-- and towards variables, so that the conversions often cancel out.
astConvertDown :: forall y z s. KnownSpan s
=> TKConversion y z -> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertDown c zftk t = case (ftkAst t, zftk) of
(yftk, _) | Just Refl <- matchingFTK yftk zftk -> t
(FTKS ZSS (FTKScalar @ry), FTKScalar @rz) ->
case testEquality (typeRep @ry) (typeRep @rz) of
Just Refl -> astConvertDownKFromS c t
Nothing -> error "astConvertDown: tensor kinds don't match"
(FTKR ZSR (FTKScalar @ry), FTKScalar @rz) ->
case testEquality (typeRep @ry) (typeRep @rz) of
Just Refl -> astConvertDownKFromR c t
Nothing -> error "astConvertDown: tensor kinds don't match"
(FTKX ZSX (FTKScalar @ry), FTKScalar @rz) ->
case testEquality (typeRep @ry) (typeRep @rz) of
Just Refl -> astConvertDownKFromX c t
Nothing -> error "astConvertDown: tensor kinds don't match"
(FTKR shr yx, FTKS shz@(_ :$$ _) zx)->
case (matchingFTK yx zx, testEquality (shsRank shz) (shrRank shr)) of
(Just Refl, Just Refl) -> astConvertDownSFromR c shz zx t
_ -> error "astConvertDown: tensor kinds don't match"
(FTKX shy yx, FTKS shz@(_ :$$ _) zx) ->
case (matchingFTK yx zx, testEquality (shsRank shz) (shxRank shy)) of
(Just Refl, Just Refl) -> astConvertDownSFromX c shz zx t
_ -> error "astConvertDown: tensor kinds don't match"
(FTKProduct{}, _) -> case t of
Ast.AstPair a1 a2 | FTKProduct zftk1 zftk2 <- zftk ->
-- Here we can't always use the c the user presumably wrote,
-- so we always create a canonical one.
astPair (astConvDown zftk1 a1) (astConvDown zftk2 a2)
Ast.AstLet var u v -> astLet var u (astConvertDown c zftk v)
Ast.AstPrimalPart v -> astPrimalPart $ astConvertDown c zftk v
Ast.AstDualPart v -> astDualPart $ astConvertDown c zftk v
Ast.AstPlainPart v -> astPlainPart $ astConvertDown c zftk v
Ast.AstFromPrimal v -> fromPrimal $ astConvertDown c zftk v
Ast.AstFromDual v -> fromDual $ astConvertDown c zftk v
Ast.AstFromPlain v -> fromPlain $ astConvertDown c zftk v
Ast.AstConvert c2 t2 -> astConvert (c `convCmp` c2) t2
_ -> Ast.AstConvert c t -- don't introduce let just to push a conversion
(yftk, _) ->
error $ "astConvertDown: wrong tensor kinds: " ++ show (yftk, zftk, t)
astConvertDownKFromS
:: forall r s. (GoodScalar r, KnownSpan s)
=> TKConversion (TKS '[] r) (TKScalar r)
-> AstTensor AstMethodLet s (TKS '[] r)
-> AstTensor AstMethodLet s (TKScalar r)
astConvertDownKFromS c a = case a of
Ast.AstProject1{} -> Ast.AstConvert c a
Ast.AstProject2{} -> Ast.AstConvert c a
Ast.AstApply (AstLambda !var !v) ll ->
astApply (AstLambda var (astConvertDownKFromS c v)) ll
Ast.AstVar{} -> Ast.AstConvert c a
Ast.AstLet var u v -> astLet var u (astConvertDownKFromS c v)
Ast.AstPrimalPart v -> astPrimalPart $ astConvertDownKFromS c v
Ast.AstDualPart v -> astDualPart $ astConvertDownKFromS c v
Ast.AstPlainPart v -> astPlainPart $ astConvertDownKFromS c v
Ast.AstFromPrimal v -> fromPrimal $ astConvertDownKFromS c v
Ast.AstFromDual v -> fromDual $ astConvertDownKFromS c v
Ast.AstFromPlain v -> fromPlain $ astConvertDownKFromS c v
AstPlusS u v -> astConvertDownKFromS c u `astPlusK` astConvertDownKFromS c v
AstTimesS u v -> astConvertDownKFromS c u `astTimesK` astConvertDownKFromS c v
Ast.AstN1S opCode u -> astN1K opCode (astConvertDownKFromS c u)
Ast.AstR1S opCode u -> astR1K opCode (astConvertDownKFromS c u)
Ast.AstR2S opCode u v ->
astR2K opCode (astConvertDownKFromS c u) (astConvertDownKFromS c v)
Ast.AstI2S QuotOp u v ->
astI2K QuotOp (astConvertDownKFromS c u) (astConvertDownKFromS c v)
Ast.AstI2S RemOp u v ->
astI2K RemOp (astConvertDownKFromS c u) (astConvertDownKFromS c v)
AstConcreteS v -> AstConcreteK (Nested.sunScalar v)
Ast.AstFloorS v -> astFloorK (kfromS v)
Ast.AstFromIntegralS v -> astFromIntegralK (kfromS v)
Ast.AstCastS v -> astCastK (kfromS v)
Ast.AstArgMinS v | FTKS (_ :$$ ZSS) FTKScalar <- ftkAst v -> astArgMinK v
Ast.AstArgMaxS v | FTKS (_ :$$ ZSS) FTKScalar <- ftkAst v -> astArgMaxK v
Ast.AstIndexS @shm ZSS v ix | Refl <- lemAppNil @shm -> astIndexK v ix
Ast.AstCondS b v1 v2 -> astCondK b (astConvertDownKFromS c v1)
(astConvertDownKFromS c v2)
Ast.AstFromVectorK ZSS l -> l V.! 0
Ast.AstFromVectorS ZSS l -> astConvertDownKFromS c $ l V.! 0
Ast.AstSumS _ v | Dict0 <- numFromTKAllNum (Proxy @r) -> astSumK v
Ast.AstScatterS{} -> Ast.AstConvert c a -- not a NF
Ast.AstReplicateK ZSS v -> v
Ast.AstReplicateS ZSS v -> astConvertDownKFromS c v
Ast.AstGatherS{} -> Ast.AstConvert c a -- not a NF
Ast.AstTransposeS{} -> Ast.AstConvert c a -- not a NF
Ast.AstReshapeS{} -> Ast.AstConvert c a -- not a NF
Ast.AstDot1InS{} -> Ast.AstConvert c a -- not a NF
Ast.AstBoolNotS b -> astBoolNotK $ astConvertDownKFromS c b
Ast.AstBoolAndS b1 b2 ->
astBoolAndK (astConvertDownKFromS c b1) (astConvertDownKFromS c b2)
Ast.AstLeqS _shb _sh arg1 arg2 -> astLeq arg1 arg2
Ast.AstConvert c2 a2 -> astConvert (c `convCmp` c2) a2
astConvertDownKFromR
:: KnownSpan s
=> TKConversion (TKR 0 r) (TKScalar r)
-> AstTensor AstMethodLet s (TKR 0 r)
-> AstTensor AstMethodLet s (TKScalar r)
astConvertDownKFromR c a = case a of
Ast.AstProject1{} -> Ast.AstConvert c a
Ast.AstProject2{} -> Ast.AstConvert c a
Ast.AstApply (AstLambda !var !v) ll ->
astApply (AstLambda var (astConvertDownKFromR c v)) ll
Ast.AstVar{} -> Ast.AstConvert c a
Ast.AstLet var u v -> astLet var u (astConvertDownKFromR c v)
Ast.AstPrimalPart v -> astPrimalPart $ astConvertDownKFromR c v
Ast.AstDualPart v -> astDualPart $ astConvertDownKFromR c v
Ast.AstPlainPart v -> astPlainPart $ astConvertDownKFromR c v
Ast.AstFromPrimal v -> fromPrimal $ astConvertDownKFromR c v
Ast.AstFromDual v -> fromDual $ astConvertDownKFromR c v
Ast.AstFromPlain v -> fromPlain $ astConvertDownKFromR c v
Ast.AstConvert c2 a2 -> astConvert (c `convCmp` c2) a2
astConvertDownKFromX
:: KnownSpan s
=> TKConversion (TKX '[] r) (TKScalar r)
-> AstTensor AstMethodLet s (TKX '[] r)
-> AstTensor AstMethodLet s (TKScalar r)
astConvertDownKFromX c a = case a of
Ast.AstProject1{} -> Ast.AstConvert c a
Ast.AstProject2{} -> Ast.AstConvert c a
Ast.AstApply (AstLambda !var !v) ll ->
astApply (AstLambda var (astConvertDownKFromX c v)) ll
Ast.AstVar{} -> Ast.AstConvert c a
Ast.AstLet var u v -> astLet var u (astConvertDownKFromX c v)
Ast.AstPrimalPart v -> astPrimalPart $ astConvertDownKFromX c v
Ast.AstDualPart v -> astDualPart $ astConvertDownKFromX c v
Ast.AstPlainPart v -> astPlainPart $ astConvertDownKFromX c v
Ast.AstFromPrimal v -> fromPrimal $ astConvertDownKFromX c v
Ast.AstFromDual v -> fromDual $ astConvertDownKFromX c v
Ast.AstFromPlain v -> fromPlain $ astConvertDownKFromX c v
Ast.AstConvert c2 a2 -> astConvert (c `convCmp` c2) a2
-- We are pushing conversions to shaped tensors down, into concrete values
-- and towards variables, so that the conversions often cancel out.
astConvertDownSFromR :: forall sh x s. KnownSpan s
=> TKConversion (TKR2 (Rank sh) x) (TKS2 sh x)
-> ShS sh -> FullShapeTK x
-> AstTensor AstMethodLet s (TKR2 (Rank sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertDownSFromR c sh x a0 = case a0 of
Ast.AstConvert c2 a2 -> astConvert (c `convCmp` c2) a2
Ast.AstProject1{} -> Ast.AstConvert c a0
Ast.AstProject2{} -> Ast.AstConvert c a0
Ast.AstApply{} -> Ast.AstConvert c a0
Ast.AstVar{} -> Ast.AstConvert c a0
Ast.AstBuild1 snat@SNat (STKR @n _ xstk) (var, a) -> case sh of
snat2@SNat :$$ rest | Just Refl <- sameNat snat snat2
, Refl <- lemRankReplicate (Proxy @n) ->
let c2 = convCmp (ConvXS' (FTKS rest x)) ConvRX
!a2 = astConvert c2 a
in Ast.AstBuild1 snat (STKS rest xstk) (var, a2)
_ -> error "astConvertDownSFromR: impossible shape"
Ast.AstLet var u v -> astLet var u (astConvertDownSFromR c sh x v)
Ast.AstPrimalPart a -> astPrimalPart $ astConvertDownSFromR c sh x a
Ast.AstDualPart a -> astDualPart $ astConvertDownSFromR c sh x a
Ast.AstPlainPart a -> astPlainPart $ astConvertDownSFromR c sh x a
Ast.AstFromPrimal a -> fromPrimal $ astConvertDownSFromR c sh x a
Ast.AstFromDual a -> fromDual $ astConvertDownSFromR c sh x a
Ast.AstFromPlain a -> fromPlain $ astConvertDownSFromR c sh x a
astConvertDownSFromX :: forall sh shx x s. (KnownSpan s, Rank shx ~ Rank sh)
=> TKConversion (TKX2 shx x) (TKS2 sh x)
-> ShS sh -> FullShapeTK x
-> AstTensor AstMethodLet s (TKX2 shx x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertDownSFromX c sh x a0 = case a0 of
Ast.AstConvert c2 a2 -> astConvert (c `convCmp` c2) a2
Ast.AstProject1{} -> Ast.AstConvert c a0
Ast.AstProject2{} -> Ast.AstConvert c a0
Ast.AstApply{} -> Ast.AstConvert c a0
Ast.AstVar{} -> Ast.AstConvert c a0
Ast.AstBuild1 snat@SNat (STKX @shx2 _ xstk) (var, a) -> case sh of
(:$$) @_ @rest snat2@SNat rest | Just Refl <- sameNat snat snat2 ->
-- This is needed only for GHC 9.10.
gcastWith (unsafeCoerceRefl :: Rank rest :~: Rank shx2) $
let c2 = ConvXS' (FTKS rest x)
!a2 = astConvert c2 a
in Ast.AstBuild1 snat (STKS rest xstk) (var, a2)
_ -> error "astConvertDownSFromX: impossible shape"
Ast.AstLet var u v -> astLet var u (astConvertDownSFromX c sh x v)
Ast.AstPrimalPart a -> astPrimalPart $ astConvertDownSFromX c sh x a
Ast.AstDualPart a -> astDualPart $ astConvertDownSFromX c sh x a
Ast.AstPlainPart a -> astPlainPart $ astConvertDownSFromX c sh x a
Ast.AstFromPrimal a -> fromPrimal $ astConvertDownSFromX c sh x a
Ast.AstFromDual a -> fromDual $ astConvertDownSFromX c sh x a
Ast.AstFromPlain a -> fromPlain $ astConvertDownSFromX c sh x a
-- We are pulling conversions from shaped tensors up.
-- z is shaped or scalar or a product (presumably with some shaped components,
-- but it's not checked; the other components are supposed to be
-- converted identically, which is not checked in c, either).
astConvertUp
:: KnownSpan s
=> TKConversion y z -> FullShapeTK z -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertUp c zftk t = case (ftkAst t, zftk, t) of
(yftk, _, _) | Just Refl <- matchingFTK yftk zftk -> t
-- Rare cases where we don't pull up but push down so that conversions
-- don't end up interspersed with AstFromPrimal and similar.
(_, _, Ast.AstFromPrimal v) ->
fromPrimal $ astConvertUp c zftk v
(_, _, Ast.AstFromDual v) ->
fromDual $ astConvertUp c zftk v
(_, _, Ast.AstFromPlain v )->
fromPlain $ astConvertUp c zftk v
(_, _, Ast.AstConvert c2 a2) ->
astConvert (c `convCmp` c2) a2
_ -> Ast.AstConvert c t
astConvDown :: forall y z s. KnownSpan s
=> FullShapeTK z -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
{-# INLINE astConvDown #-}
astConvDown zftk t = case convDownMaybe (ftkAst t) (ftkToSTK zftk) of
Just c -> astConvertDown c zftk t
Nothing -> error $ "astConvDown: unexpected types "
++ "(" ++ show (ftkAst t) ++ ", " ++ show zftk ++ ")"
astConvDownSFromR :: forall sh s x. KnownSpan s
=> ShS sh -> FullShapeTK x
-> AstTensor AstMethodLet s (TKR2 (Rank sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvDownSFromR sh x | Refl <- lemRankReplicate (Proxy @(Rank sh)) =
astConvertDownSFromR (ConvCmp (ConvXS' (FTKS sh x)) ConvRX) sh x
astConvDownSFromX :: forall sh sh' s x. (KnownSpan s, Rank sh ~ Rank sh')
=> ShS sh -> FullShapeTK x
-> AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvDownSFromX sh x = astConvertDownSFromX (ConvXS' (FTKS sh x)) sh x
astConvUp :: forall y z s. KnownSpan s
=> FullShapeTK z -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
{-# INLINE astConvUp #-}
astConvUp zftk t = case convUpMaybe (ftkAst t) zftk of
Just c -> astConvertUp c zftk t
Nothing -> error $ "astConvUp: unexpected types "
++ "(" ++ show (ftkAst t) ++ ", " ++ show zftk ++ ")"
astConvUpRFromS :: forall sh x s. KnownSpan s
=> ShS sh -> FullShapeTK x
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank sh) x)
astConvUpRFromS sh x t | Refl <- lemRankMapJust sh =
astConvertUp
(ConvCmp (ConvXR (ftkToSTK x)) ConvSX) (FTKR (shrFromShS sh) x) t
astConvUpXFromS :: forall sh sh' x s. (Rank sh ~ Rank sh', KnownSpan s)
=> IShX sh' -> FullShapeTK x
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh' x)
astConvUpXFromS sh' x =
gcastWith (unsafeCoerceRefl :: Rank (MapJust sh) :~: Rank sh) $
astConvertUp (ConvCmp (ConvXX' (FTKX sh' x)) ConvSX) (FTKX sh' x)
astSumKContract :: forall shm r s. (NumScalar r, KnownSpan s)
=> AstTensor AstMethodLet s (TKS shm r)
-> AstTensor AstMethodLet s (TKScalar r)
astSumKContract t0 = case t0 of
Ast.AstLet var u v -> astLet var u (astSumKContract v)
AstTimesS t1 t2 -> astDot0 t1 t2
Ast.AstDot1InS _ _ t1 t2 -> astDot0 t1 t2
Ast.AstMatmul2S m@SNat SNat p@SNat m1 m2 ->
astDot0 (astTransposeS (Permutation.makePerm @'[1, 0])
(astReplicateS (p :$$ ZSS) m1))
(astTransposeS (Permutation.makePerm @'[0, 2, 1])
(astReplicateS (m :$$ ZSS) m2))
_ -> astSumK t0
astDot0 :: (NumScalar r, KnownSpan s)
=> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKScalar r)
astDot0 t1 t2 = case (t1, t2) of
_ | Just u1 <- unRepl t1 ->
kfromS u1 `astTimesK` astSumKContract t2
_ | Just u2 <- unRepl t2 ->
kfromS u2 `astTimesK` astSumKContract t1
_ | FTKS (snat :$$ _) _ <- ftkAst t1
, Just u1 <- unRepl1 t1
, Just u2 <- unRepl1 t2 ->
astTimesK (astDot0 u1 u2)
(fromPlain $ AstConcreteK $ fromIntegral $ fromSNat' snat)
(Ast.AstFromPrimal u1, Ast.AstFromPrimal u2) ->
fromPrimal $ astDot0 u1 u2
(Ast.AstFromDual u1, Ast.AstFromDual u2) ->
fromDual $ astDot0 u1 u2
(Ast.AstFromPlain u1, Ast.AstFromPlain u2) ->
fromPlain $ astDot0 u1 u2
(Ast.AstN1S NegateOp u1, Ast.AstN1S NegateOp u2) -> astDot0 u1 u2
(u, v) | Just u0 <- unAstS u
, Just v0 <- unAstS v ->
withKnownShS (Nested.sshape u0) $
fromPlain $ astConcreteK $ tsdot0 (Concrete u0) (Concrete v0)
(Ast.AstFromVectorK _ l1, Ast.AstFromVectorK _ l2) | V.length l1 < 100 ->
V.foldr (\(a1, a2) acc -> acc + a1 `astTimesK` a2) 0 $ V.zip l1 l2
-- TODO: instead take the longest common prefix (sufix?) of shm1 and shm2:
(Ast.AstSumS shm1 u1, Ast.AstSumS shm2 u2)
| Just Refl <- testEquality shm1 shm2 ->
astDot0 u1 u2
{- KnownNat would be needed (or SNat):
(Ast.AstAppendS @m1 u1 v1, Ast.AstAppendS @m2 u2 v2)
| Just Refl <- sameNat (SNat @m1) (SNat @m2) ->
astDot0 u1 u2 + astDot0 v1 v2 -}
(Ast.AstReverseS u1, Ast.AstReverseS u2) -> astDot0 u1 u2
(Ast.AstTransposeS @_ @sh1 perm1 u1, Ast.AstTransposeS @_ @sh2 perm2 u2)
| Just Refl <- testEquality perm1 perm2 ->
gcastWith (unsafeCoerceRefl :: sh1 :~: sh2) $
astDot0 u1 u2
_ -> case ftkAst t1 of
FTKS ZSS FTKScalar -> kfromS t1 `astTimesK` kfromS t2
FTKS (SNat' @1 :$$ ZSS) FTKScalar ->
astTimesK (astIndexK t1 (0 :.$ ZIS)) (astIndexK t2 (0 :.$ ZIS))
_ -> Ast.AstDot0 t1 t2
astDot1InS :: forall sh n r s. (NumScalar r, KnownSpan s)
=> ShS sh -> SNat n
-> AstTensor AstMethodLet s (TKS (sh ++ '[n]) r)
-> AstTensor AstMethodLet s (TKS (sh ++ '[n]) r)
-> AstTensor AstMethodLet s (TKS sh r)
astDot1InS sh n@SNat t1 t2 = case (t1, t2) of
(Ast.AstFromPrimal u1, Ast.AstFromPrimal u2) ->
fromPrimal $ astDot1InS sh n u1 u2
(Ast.AstFromDual u1, Ast.AstFromDual u2) ->
fromDual $ astDot1InS sh n u1 u2
(Ast.AstFromPlain u1, Ast.AstFromPlain u2) ->
fromPlain $ astDot1InS sh n u1 u2
(Ast.AstN1S NegateOp u1, Ast.AstN1S NegateOp u2) -> astDot1InS sh n u1 u2
{- TODO: instead take the longest common prefix (sufix?) of shm1 and shm2.
-- TODO: possibly this increases work, so wait until we have a generalization
-- of tsdot1In.
(Ast.AstSumS @shm1 @shn1 shm1 u1, Ast.AstSumS shm2 u2)
| Just Refl <- testEquality shm1 shm2 ->
gcastWith (unsafeCoerceRefl :: shm1 ++ sh ++ '[n] :~: shm1 ++ shn1) $
astSumS shm1 $ astDot1InS (shm1 `shsAppend` sh) n u1 u2 -}
(u, v) | Just u0 <- unAstS u
, Just v0 <- unAstS v ->
withKnownShS sh $
fromPlain
$ astConcreteS $ tsdot1In @_ @sh (SNat @n) (Concrete u0) (Concrete v0)
_ -> case sh of
ZSS -> sfromK $ astDot0 t1 t2
snat :$$ shRest -> case (t1, t2) of
_ | Just u1 <- unRepl1 t1
, Just u2 <- unRepl1 t2 ->
astReplicateS (snat :$$ ZSS) $ astDot1InS shRest n u1 u2
(Ast.AstReverseS u1, Ast.AstReverseS u2) ->
astReverseS $ astDot1InS sh n u1 u2
_ -> Ast.AstDot1InS sh n t1 t2
astMatmul2S :: (NumScalar r, KnownSpan s)
=> SNat m -> SNat n -> SNat p
-> AstTensor AstMethodLet s (TKS '[m, n] r)
-> AstTensor AstMethodLet s (TKS '[n, p] r)
-> AstTensor AstMethodLet s (TKS '[m, p] r)
astMatmul2S m@SNat n@SNat p@SNat t1 t2 = case (t1, t2) of
(Ast.AstFromPrimal u1, Ast.AstFromPrimal u2) ->
fromPrimal $ astMatmul2S m n p u1 u2
(Ast.AstFromDual u1, Ast.AstFromDual u2) ->
fromDual $ astMatmul2S m n p u1 u2
(Ast.AstFromPlain u1, Ast.AstFromPlain u2) ->
fromPlain $ astMatmul2S m n p u1 u2
(u, v) | Just u0 <- unAstS u
, Just v0 <- unAstS v ->
fromPlain $ astConcreteS $ tsmatmul2 (Concrete u0) (Concrete v0)
_ -> Ast.AstMatmul2S m n p t1 t2
astBoolNotK :: AstBool AstMethodLet -> AstBool AstMethodLet
astBoolNotK = \case
Ast.AstLet var n b -> astLet var n (astBoolNotK b)
AstConcreteK b -> AstConcreteK $ not b
Ast.AstCondK b v w -> astCondK b (astBoolNotK v) (astBoolNotK w)
AstLeqInt n k -> astLeqK (1 + k) n
Ast.AstBoolNotK b -> b
b -> Ast.AstBoolNotK b
astBoolNotS :: AstTensor AstMethodLet PlainSpan (TKS sh Bool)
-> AstTensor AstMethodLet PlainSpan (TKS sh Bool)
astBoolNotS = Ast.AstBoolNotS
astBoolAndK :: AstBool AstMethodLet -> AstBool AstMethodLet
-> AstBool AstMethodLet
astBoolAndK = \cases
u v | eqY u v -> u
(Ast.AstLet var n b) c -> astLetRefresh var n b $ \b' -> astBoolAndK b' c
b (Ast.AstLet var n c) -> astLetRefresh var n c $ \c' -> astBoolAndK b c'
(AstConcreteK True) b -> b
(AstConcreteK False) _b -> AstConcreteK False
b (AstConcreteK True) -> b
_b (AstConcreteK False) -> AstConcreteK False
(Ast.AstBoolAndK b c) d -> b `astBoolAndK` (c `astBoolAndK` d)
-- Normal forms:
b@(Ast.AstLeqK AstConcreteK{} Ast.AstVar{}) c ->
Ast.AstBoolAndK b c
b@(Ast.AstLeqK AstConcreteK{} (Ast.AstN1K NegateOp Ast.AstVar{})) c ->
Ast.AstBoolAndK b c
b@(Ast.AstBoolNotK
(Ast.AstBoolAndK (Ast.AstLeqK AstConcreteK{} Ast.AstVar{}) _)) c ->
Ast.AstBoolAndK b c
b@(Ast.AstBoolNotK
(Ast.AstBoolAndK
(Ast.AstLeqK AstConcreteK{}
(Ast.AstN1K NegateOp Ast.AstVar{})) _)) c ->
Ast.AstBoolAndK b c
-- Rewriting to the normal forms by commutativity:
b c@(Ast.AstLeqK AstConcreteK{} Ast.AstVar{}) ->
Ast.AstBoolAndK c b
b c@(Ast.AstLeqK AstConcreteK{} (Ast.AstN1K NegateOp Ast.AstVar{})) ->
Ast.AstBoolAndK c b
b c@(Ast.AstBoolNotK
(Ast.AstBoolAndK (Ast.AstLeqK AstConcreteK{} Ast.AstVar{}) _)) ->
Ast.AstBoolAndK c b
b c@(Ast.AstBoolNotK
(Ast.AstBoolAndK
(Ast.AstLeqK AstConcreteK{}
(Ast.AstN1K NegateOp Ast.AstVar{})) _)) ->
Ast.AstBoolAndK c b
-- Rewriting to the normal forms by commutativity and associativity:
b (Ast.AstBoolAndK
c@(Ast.AstLeqK AstConcreteK{} Ast.AstVar{}) d) ->
astBoolAndK c (b `astBoolAndK` d)
b (Ast.AstBoolAndK
c@(Ast.AstLeqK AstConcreteK{}
(Ast.AstN1K NegateOp Ast.AstVar{})) d) ->
astBoolAndK c (b `astBoolAndK` d)
b (Ast.AstBoolAndK
c@(Ast.AstBoolNotK
(Ast.AstBoolAndK (Ast.AstLeqK AstConcreteK{} Ast.AstVar{}) _)) d) ->
astBoolAndK c (b `astBoolAndK` d)
b (Ast.AstBoolAndK
c@(Ast.AstBoolNotK
(Ast.AstBoolAndK
(Ast.AstLeqK AstConcreteK{}
(Ast.AstN1K NegateOp Ast.AstVar{})) _)) d) ->
astBoolAndK c (b `astBoolAndK` d)
b c -> Ast.AstBoolAndK b c
astBoolAndS :: AstTensor AstMethodLet PlainSpan (TKS sh Bool)
-> AstTensor AstMethodLet PlainSpan (TKS sh Bool)
-> AstTensor AstMethodLet PlainSpan (TKS sh Bool)
astBoolAndS = Ast.AstBoolAndS
astLeqK :: forall r. NumScalar r
=> AstTensor AstMethodLet PlainSpan (TKScalar r)
-> AstTensor AstMethodLet PlainSpan (TKScalar r)
-> AstBool AstMethodLet
astLeqK = \cases
u v | eqY u v -> true
u v | Just (u1, u2) <- bounds u
, Just (v1, v2) <- bounds v
, u2 <= v1 || u1 > v2 -> AstConcreteK (u2 <= v1)
-- This is wrong, because LHS is a particular valuation and RHS is all.
-- Ast.AstLet _ _ u <=. Ast.AstLet _ _ v -> u <=. v
(Ast.AstLet var n u) v -> astLetRefresh var n u $ \u' -> astLeqK u' v
u (Ast.AstLet var n v) -> astLetRefresh var n v $ \v' -> astLeqK u v'
u (AstPlusK (AstConcreteK v) w) ->
astLeqK (astPlusK u (AstConcreteK $ negate v)) w
(AstPlusK (AstConcreteK u) w) v ->
astLeqK (AstConcreteK u) (astPlusK (negate w) v)
u (AstConcreteK v) ->
astLeqK (AstConcreteK (negate v)) (negate u)
(AstConcreteK u) (AstTimesK (AstConcreteK v) w)
| v > 0 && u >= 0
, Just Refl <- testEquality (typeRep @r) (typeRep @Int) ->
astLeqK (AstConcreteK ((u + v - 1) `quotH` v)) w
(AstConcreteK u) (AstTimesK (AstConcreteK v) w)
| v > 0 && u < 0
, Just Refl <- testEquality (typeRep @r) (typeRep @Int) ->
astLeqK (AstConcreteK (u `quotH` v)) w
(AstConcreteK u) (AstTimesK (AstConcreteK v) w)
| v < 0
, Just Refl <- testEquality (typeRep @r) (typeRep @Int) ->
astLeqK (AstConcreteK u)
(AstTimesK (AstConcreteK $ negate v) -- not astTimesK!
(negate w))
u v -> Ast.AstLeqK u v
astLeq :: forall sh r. NumScalar r
=> AstTensor AstMethodLet PlainSpan (TKS sh r)
-> AstTensor AstMethodLet PlainSpan (TKS sh r)
-> AstBool AstMethodLet
astLeq = \cases
u v | eqY u v -> true
(AstConcreteS u) (AstConcreteS v) ->
AstConcreteK $ Shaped.stoPrimitive u <= Shaped.stoPrimitive v
u (AstPlusS (AstConcreteS v) w) ->
astLeq (astPlusS u (AstConcreteS $ negate v)) w
(AstPlusS (AstConcreteS u) w) v ->
astLeq (AstConcreteS u) (astPlusS (negate w) v)
u (AstConcreteS v) ->
astLeq (AstConcreteS (negate v)) (negate u)
(AstConvUpSFromK w) (AstConvUpSFromK v) -> astLeqK w v
w (AstConvUpSFromK v) -> astLeqK (kfromS w) v
(AstConvUpSFromK v) w -> astLeqK v (kfromS w)
u v -> Ast.AstLeq u v
astLeqS :: forall shb sh r. NumScalar r
=> ShS shb -> ShS sh
-> AstTensor AstMethodLet PlainSpan (TKS (shb ++ sh) r)
-> AstTensor AstMethodLet PlainSpan (TKS (shb ++ sh) r)
-> AstTensor AstMethodLet PlainSpan (TKS shb Bool)
astLeqS = Ast.AstLeqS
-- * Scalar and shaped AstMethodLet AST instances
-- This instance is needed for unwinding in astConcrete.
instance KnownSpan s => ConvertTensor (AstTensor AstMethodLet s) where
tconvert c _astk = astConvert c
-- These are somewhat faster than their default implementations.
kfromR = astConvertDownKFromR (ConvCmp ConvX0 ConvRX)
kfromS = astConvertDownKFromS (ConvCmp ConvX0 ConvSX)
kfromX = astConvertDownKFromX ConvX0
sfromK = astConvertUp (ConvCmp ConvXS (Conv0X STKScalar)) (FTKS ZSS FTKScalar)
rfromK = astConvertUp (ConvCmp (ConvXR STKScalar) (Conv0X STKScalar))
(FTKR ZSR FTKScalar)
rfromS t = case ftkAst t of
FTKS sh x -> astConvUpRFromS sh x t
rfromX t = case ftkAst t of
FTKX sh' x ->
withShsFromShX sh' $ \(sh :: ShS sh) ->
astConvUpRFromS sh x $ astConvDownSFromX sh x t
xfromR t = case ftkAst t of
FTKR shr x ->
withShsFromShR shr $ \(sh :: ShS sh) ->
astConvUpXFromS (shCastSX knownShX sh) x
$ astConvDownSFromR sh x t
sfromR t = case ftkAst t of
FTKR _ x -> astConvDownSFromR knownShS x t
sfromX t = case ftkAst t of
FTKX _ x -> astConvDownSFromX knownShS x t
xfromK = astConvertUp (Conv0X STKScalar) (FTKX ZSX FTKScalar)
xfromS t = case ftkAst t of
FTKS sh x -> astConvUpXFromS (shCastSX knownShX sh) x t
rzip @_ @_ @n a
| Refl <- lemRankReplicate (Proxy @n) = case ftkAst a of
FTKProduct (FTKR _sh y) (FTKR _ z) ->
let c = convCmp
(ConvXR (ftkToSTK (FTKProduct y z)))
(convCmp
(ConvZip (ftkToSTK y) (ftkToSTK z))
(ConvT2 ConvRX ConvRX))
in astConvert c a
runzip @_ @_ @n a
| Refl <- lemRankReplicate (Proxy @n) = case ftkAst a of
FTKR _sh (FTKProduct y z) ->
let c = convCmp
(ConvT2 (ConvXR (ftkToSTK y)) (ConvXR (ftkToSTK z)))
(convCmp
(ConvUnzip (ftkToSTK y) (ftkToSTK z))
ConvRX)
in astConvert c a
szip a = case ftkAst a of
FTKProduct (FTKS _sh y) (FTKS _ z) ->
let c = convCmp
ConvXS
(convCmp
(ConvZip (ftkToSTK y) (ftkToSTK z))
(ConvT2 ConvSX ConvSX))
in astConvert c a
sunzip a = case ftkAst a of
FTKS _sh (FTKProduct y z) ->
let c = convCmp
(ConvT2 ConvXS ConvXS)
(convCmp
(ConvUnzip (ftkToSTK y) (ftkToSTK z))
ConvSX)
in astConvert c a
xzip a = case ftkAst a of
FTKProduct (FTKX _sh y) (FTKX _ z) ->
let c = ConvZip (ftkToSTK y) (ftkToSTK z)
in astConvert c a
xunzip a = case ftkAst a of
FTKX _sh (FTKProduct y z) ->
let c = ConvUnzip (ftkToSTK y) (ftkToSTK z)
in astConvert c a
xnestR @sh1 @m @x sh1 a
| Refl <- lemRankReplicate (Proxy @m) =
let c :: TKConversion (TKX2 (sh1 ++ Replicate m Nothing) x)
(TKX2 sh1 (TKR2 m x))
c = convCmp
(ConvXX (ConvXR (knownSTK @x)))
(ConvNest @_ @_ @(Replicate m Nothing)
(STKX sh1 (knownSTK @x)))
in astConvert c a
xnestS @_ @_ @x sh1 a =
let c = convCmp (ConvXX ConvXS)
(ConvNest (STKX sh1 (knownSTK @x)))
in astConvert c a
xnest @_ @_ @x sh1 a =
let c = ConvNest (STKX sh1 (knownSTK @x))
in astConvert c a
xunNestR a =
let c = convCmp ConvUnnest
(ConvXX ConvRX)
in astConvert c a
xunNestS a =
let c = convCmp ConvUnnest
(ConvXX ConvSX)
in astConvert c a
xunNest a =
let c = ConvUnnest
in astConvert c a
tpairConv = astPair
tunpairConv t = (astProject1 t, astProject2 t)
-- These instances improve readability of rules.
instance (NumScalar r, KnownSpan s)
=> Num (AstTensor AstMethodLet s (TKScalar r)) where
(+) = astPlusK
(*) = astTimesK
negate = astN1K NegateOp
abs = astN1K AbsOp
signum = astN1K SignumOp
{-# INLINE fromInteger #-}
fromInteger i = fromPlain $ AstConcreteK (fromInteger i)
instance (NumScalar r, IntegralH r, Nested.IntElt r, KnownSpan s)
=> IntegralH (AstTensor AstMethodLet s (TKScalar r)) where
quotH = astI2K QuotOp
remH = astI2K RemOp
instance (NumScalar r, Differentiable r, KnownSpan s)
=> Fractional (AstTensor AstMethodLet s (TKScalar r)) where
(/) = astR2K DivideOp
recip = astR1K RecipOp
{-# INLINE fromRational #-}
fromRational r = fromPlain $ AstConcreteK (fromRational r)
instance (NumScalar r, Differentiable r, KnownSpan s)
=> Floating (AstTensor AstMethodLet s (TKScalar r)) where
pi = error "pi is not defined for tensors"
exp = astR1K ExpOp
log = astR1K LogOp
sqrt = astR1K SqrtOp
(**) = astR2K PowerOp
logBase = astR2K LogBaseOp
sin = astR1K SinOp
cos = astR1K CosOp
tan = astR1K TanOp
asin = astR1K AsinOp
acos = astR1K AcosOp
atan = astR1K AtanOp
sinh = astR1K SinhOp
cosh = astR1K CoshOp
tanh = astR1K TanhOp
asinh = astR1K AsinhOp
acosh = astR1K AcoshOp
atanh = astR1K AtanhOp
instance (NumScalar r, Differentiable r, KnownSpan s)
=> RealFloatH (AstTensor AstMethodLet s (TKScalar r)) where
atan2H = astR2K Atan2Op
instance (NumScalar r, KnownSpan s)
=> Num (AstTensor AstMethodLet s (TKS sh r)) where
(+) = astPlusS
(*) = astTimesS
negate = astN1S NegateOp
abs = astN1S AbsOp
signum = astN1S SignumOp
fromInteger i = error $ "fromInteger is not defined for shaped tensors: "
++ show i
instance (NumScalar r, IntegralH r, Nested.IntElt r, KnownSpan s)
=> IntegralH (AstTensor AstMethodLet s (TKS sh r)) where
quotH = astI2S QuotOp
remH = astI2S RemOp
instance (NumScalar r, Differentiable r, KnownSpan s)
=> Fractional (AstTensor AstMethodLet s (TKS sh r)) where
(/) = astR2S DivideOp
recip = astR1S RecipOp
fromRational r = error $ "fromRational is not defined for shaped tensors: "
++ show r
instance (NumScalar r, Differentiable r, KnownSpan s)
=> Floating (AstTensor AstMethodLet s (TKS sh r)) where
pi = error "pi is not defined for tensors"
exp = astR1S ExpOp
log = astR1S LogOp
sqrt = astR1S SqrtOp
(**) = astR2S PowerOp
logBase = astR2S LogBaseOp
sin = astR1S SinOp
cos = astR1S CosOp
tan = astR1S TanOp
asin = astR1S AsinOp
acos = astR1S AcosOp
atan = astR1S AtanOp
sinh = astR1S SinhOp
cosh = astR1S CoshOp
tanh = astR1S TanhOp
asinh = astR1S AsinhOp
acosh = astR1S AcoshOp
atanh = astR1S AtanhOp
instance (NumScalar r, Differentiable r, KnownSpan s)
=> RealFloatH (AstTensor AstMethodLet s (TKS sh r)) where
atan2H = astR2S Atan2Op
instance Boolean (AstBool AstMethodLet) where
true = AstConcreteK True
false = AstConcreteK False
notB = astBoolNotK
(&&*) = astBoolAndK
b ||* c = notB (notB b &&* notB c)
-- Since u and v are duplicated here, they need to be shared.
-- We share their difference, which would most likely appear in the
-- inequalities once they are rewritten, to ensure it's shared and whatever
-- vectorization substitutes into it is shared as well.
-- Otherwise, if u and v are variables, the sharing would vanish
-- before vectoriation complicates the expression a bit, making it
-- worth sharing.
instance (KnownSpan s, NumScalar r)
=> EqH (AstTensor AstMethodLet s) (TKScalar r) where
u ==. v | eqY u v = true
uUnshared ==. vUnshared = astLetFun (uUnshared - vUnshared) $ \uv ->
0 <=. uv &&* uv <=. 0
instance (KnownSpan s, NumScalar r)
=> EqH (AstTensor AstMethodLet s) (TKS sh r) where
u ==. v | eqY u v = true
uUnshared ==. vUnshared = astLetFun (uUnshared - vUnshared) $ \uv ->
let zero = fromPlain $ AstConcreteS $ defTargetRep $ ftkAst uUnshared
in zero <=. uv &&* uv <=. zero
instance (KnownSpan s, NumScalar r)
=> OrdH (AstTensor AstMethodLet s) (TKScalar r) where
u <=. v = fromPlain $ astLeqK (astPlainPart u) (astPlainPart v)
instance (KnownSpan s, NumScalar r)
=> OrdH (AstTensor AstMethodLet s) (TKS sh r) where
u <=. v = fromPlain $ astLeq (astPlainPart u) (astPlainPart v)
-- * Helper combinators
-- All but the last case are shortcuts for common forms.
astConcrete :: FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PlainSpan y
{-# INLINE astConcrete #-}
astConcrete ftk v = case ftk of
FTKScalar -> astConcreteK v
FTKR ZSR FTKScalar ->
rfromK $ AstConcreteK $ Nested.runScalar $ unConcrete v
FTKR sh' FTKScalar ->
withShsFromShR sh' $ \sh ->
withKnownShS sh $
astConvUpRFromS sh FTKScalar $ astConcreteS $ sfromR v
FTKS ZSS FTKScalar ->
sfromK $ AstConcreteK $ Nested.sunScalar $ unConcrete v
FTKS _ FTKScalar -> astConcreteS v
FTKX ZSX FTKScalar ->
xfromK $ AstConcreteK $ Nested.munScalar $ unConcrete v
FTKX sh' FTKScalar ->
withShsFromShX sh' $ \(sh :: ShS sh) ->
withKnownShS sh $
astConvUpXFromS sh' FTKScalar
$ astConcreteS $ sfromX @_ @sh v
FTKProduct ftk1 ftk2 ->
astPair (astConcrete ftk1 (tproject1 v))
(astConcrete ftk2 (tproject2 v))
_ -> concreteTarget astConcreteK astConcreteS
(\sh -> astConvUpRFromS sh FTKScalar)
(\sh' -> astConvUpXFromS sh' FTKScalar)
(ftkToSTK ftk) v
-- INLINE here would bloat the binary a lot, probably negating any
-- gains from directly calling the function. Also, this is not a bottleneck.
-- It would also trigger specialization of many other functions from this
-- module, which would fail, because their unfolding is not exposed.
-- To some extend, this is even happening with no pragma (wrappers and
-- similar functions derived from this one fail to specialize in other modules).
astLetFun :: forall y z s s2. (KnownSpan s, KnownSpan s2)
=> AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
{-# NOINLINE astLetFun #-}
astLetFun a f | astIsSmall True a = f a
astLetFun a f = case a of
AstConvUp c ftkz v ->
unsafePerformIO $ do
var <- funToAstNoBoundsIO (ftkAst v)
pure $! astLet var v (f $ astConvertUp c ftkz $ astVar var)
Ast.AstFromPrimal v -> astLetFun v (f . fromPrimal)
Ast.AstFromDual v -> astLetFun v (f . fromDual)
Ast.AstFromPlain v -> astLetFun v (f . fromPlain)
_ -> unsafePerformIO $ case ftkAst a of
ftk@FTKScalar -> do
var <- funToAstAutoBoundsIO ftk a
pure $! astLet var a (f $ astVar var)
FTKR sh' x ->
withShsFromShR sh' $ \(sh :: ShS sh) -> do
let v = astConvDownSFromR sh x a
var <- funToAstNoBoundsIO (FTKS sh x)
pure $! astLet var v (f $ astConvUpRFromS sh x $ astVar var)
-- safe, because subsitution ruled out above
FTKX sh' x ->
withShsFromShX sh' $ \(sh :: ShS sh) -> do
let v = astConvDownSFromX sh x a
var <- funToAstNoBoundsIO (FTKS sh x)
pure $! astLet var v (f $ astConvUpXFromS sh' x $ astVar var)
FTKS ZSS x@FTKScalar -> do
let v = kfromS a
var <- funToAstAutoBoundsIO x v
pure $! astLet var v (f $ sfromK $ astVar var)
-- calling recursively for product may be not worth it
ftk -> do
var <- funToAstNoBoundsIO ftk
pure $! astLet var a (f $ astVar var)
astLetRefresh :: forall y x z s s2 s3. (KnownSpan s2, KnownSpan s3)
=> AstVarName '(s, y) -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s3 x
-> (AstTensor AstMethodLet s3 x -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetRefresh var u v f =
withKnownSpan (varNameToSpan var) $
astLetFun u $ \ !ast1 -> f (substituteAst ast1 var v)
unRepl1 :: forall n sh x s. KnownSpan s
=> AstTensor AstMethodLet s (TKS2 (n ': sh) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
unRepl1 (Ast.AstLet var u t) =
Ast.AstLet var u <$> unRepl1 t -- we may be before inlining
unRepl1 (Ast.AstPrimalPart t) = astPrimalPart <$> unRepl1 t
unRepl1 (Ast.AstDualPart t) = astDualPart <$> unRepl1 t
unRepl1 (Ast.AstPlainPart t) = astPlainPart <$> unRepl1 t
unRepl1 (Ast.AstFromPrimal t) = fromPrimal <$> unRepl1 t
unRepl1 (Ast.AstFromDual t) = fromDual <$> unRepl1 t
unRepl1 (Ast.AstFromPlain t) = fromPlain <$> unRepl1 t
unRepl1 (AstConcreteS a) = AstConcreteS <$> sunReplicate1 a
unRepl1 (Ast.AstReplicateK (_ :$$ shmRest) u) = Just $ astReplicateK shmRest u
unRepl1 (Ast.AstReplicateS (_ :$$ shmRest) u) = Just $ astReplicateS shmRest u
unRepl1 _ = Nothing
-- The result must not be equal to the argument.
unRepl :: forall sh x s. KnownSpan s
=> AstTensor AstMethodLet s (TKS2 sh x)
-> Maybe (AstTensor AstMethodLet s (TKS2 '[] x))
-- This is too costly and not needed in all the places where unRepl is used,
-- hence the restriction to different result and argument:
-- unRepl t | FTKS ZSS _ <- ftkAst t = Just t
unRepl (Ast.AstLet var u t) = Ast.AstLet var u <$> unRepl t
unRepl (Ast.AstPrimalPart t) = astPrimalPart <$> unRepl t
unRepl (Ast.AstDualPart t) = astDualPart <$> unRepl t
unRepl (Ast.AstPlainPart t) = astPlainPart <$> unRepl t
unRepl (Ast.AstFromPrimal t) = fromPrimal <$> unRepl t
unRepl (Ast.AstFromDual t) = fromDual <$> unRepl t
unRepl (Ast.AstFromPlain t) = fromPlain <$> unRepl t
unRepl (AstConcreteS a) | _ :$$ _ <- Nested.sshape a =
AstConcreteS . Nested.sscalar <$> sunReplicate a
unRepl (Ast.AstReplicateK _ u) = Just $ sfromK u
unRepl (Ast.AstReplicateS _ u) = unRepl u
unRepl _ = Nothing
-- The result must not be equal to the argument.
unReplN :: forall shm shn x s. KnownSpan s
=> ShS shm -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 shn x))
unReplN ZSS a = Just a
unReplN (_ :$$ sh) (Ast.AstReplicateK (_ :$$ shmRest) u) =
unReplN sh (astReplicateK shmRest u)
unReplN (_ :$$ sh) (Ast.AstReplicateS (_ :$$ shmRest) u) =
unReplN sh (astReplicateS shmRest u)
unReplN shm (Ast.AstLet var u t) = Ast.AstLet var u <$> unReplN shm t
unReplN shm (Ast.AstPrimalPart t) = astPrimalPart <$> unReplN shm t
unReplN shm (Ast.AstDualPart t) = astDualPart <$> unReplN shm t
unReplN shm (Ast.AstPlainPart t) = astPlainPart <$> unReplN shm t
unReplN shm (Ast.AstFromPrimal t) = fromPrimal <$> unReplN shm t
unReplN shm (Ast.AstFromDual t) = fromDual <$> unReplN shm t
unReplN shm (Ast.AstFromPlain t) = fromPlain <$> unReplN shm t
unReplN shm (AstConcreteS a) = AstConcreteS <$> sunReplicateN shm a
unReplN _ _ = Nothing
-- * Substitution wrappers
-- | We assume no variable is shared between a binding and its nested binding
-- and nobody substitutes into variables that are bound.
-- This keeps the substitution code simple, because we never need to compare
-- variables to any variable in the bindings.
substituteAst :: forall s s2 y z. KnownSpan s
=> AstTensor AstMethodLet s2 z -> AstVarName '(s2, z)
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
substituteAst i var v1 =
fromMaybe v1 $ substitute1Ast i var v1
substituteAstIxS
:: AstTensor AstMethodLet s y -> AstVarName '(s, y) -> AstIxS AstMethodLet sh
-> AstIxS AstMethodLet sh
substituteAstIxS i var ix =
fromMaybe ix $ substitute1AstIxS i var ix
-- * Substitution workers
-- Invariant: if the variable has bounds, the expression can only have
-- values within the bounds (regardless of what the `bounds` call would say).
substitute1Ast :: forall s s2 y z. KnownSpan s
=> AstTensor AstMethodLet s2 z -> AstVarName '(s2, z)
-> AstTensor AstMethodLet s y
-> Maybe (AstTensor AstMethodLet s y)
substitute1Ast i var = subst where
substIxS :: AstIxS AstMethodLet sh3 -> Maybe (AstIxS AstMethodLet sh3)
substIxS ix =
let mix = fmap subst ix
in if any isJust mix
then Just $ ixsZipWith fromMaybe ix mix
else Nothing
subst :: forall s3 y2. KnownSpan s3
=> AstTensor AstMethodLet s3 y2 -> Maybe (AstTensor AstMethodLet s3 y2)
subst = \case
Ast.AstPair u v ->
case (subst u, subst v) of
(Nothing, Nothing) -> Nothing
(mu, mv) -> Just $ astPair (fromMaybe u mu) (fromMaybe v mv)
Ast.AstProject1 a -> astProject1 <$> subst a
Ast.AstProject2 a -> astProject2 <$> subst a
Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->
case ( substitute1AstHFun i var f, substitute1AstHFun i var df
, substitute1AstHFun i var rf, subst acc0
, subst es ) of
(Nothing, Nothing, Nothing, Nothing, Nothing) -> Nothing
(mf, mdf, mrf, macc0, mes) ->
Just $ astMapAccumLDer k bftk eftk
(fromMaybe f mf)
(fromMaybe df mdf)
(fromMaybe rf mrf)
(fromMaybe acc0 macc0)
(fromMaybe es mes)
Ast.AstApply t ll ->
case ( substitute1AstHFun i var t
, subst ll ) of
(Nothing, Nothing) -> Nothing
(mt, mll) -> Just $ astApply (fromMaybe t mt) (fromMaybe ll mll)
Ast.AstVar var2 ->
-- We can't assert anything here about bounds, because only runtime concrete
-- concrete values need to be in bounds and bounds approximations
-- don't have to agree.
if varNameToAstVarId var == varNameToAstVarId var2
then case testEquality var var2 of
Just Refl -> case (i, var2) of
(Ast.AstVar var3, AstVarName _ (FtkAndBoundsBounds lb2 ub2)) ->
let bs = case var3 of
AstVarName _ (FtkAndBoundsBounds lb3 ub3) ->
(max lb2 lb3, min ub2 ub3)
-- We know all bounds approximations have to be correct
-- so we can intersect them. We ignore approximations
-- from @var, because they should be equal to those of var2,
-- because we never shadow variables nor update their types.
_ -> (lb2, ub2)
in Just $ astVar $ reboundsVarName bs var3
_ -> Just i
_ -> error $ "substitute1Ast: kind of the variable "
++ show var2 ++ ": " ++ show (varNameToFTK var)
++ ", payload kind: " ++ show (varNameToFTK var2)
++ ", payload: " ++ show i
else Nothing
Ast.AstBuild1 k stk (var2, v) ->
assert (varNameToAstVarId var2 /= varNameToAstVarId var) $
case subst v of
Just !v2 -> Just $ Ast.AstBuild1 k stk (var2, v2)
Nothing -> Nothing
Ast.AstLet var2 u v ->
assert (varNameToAstVarId var2 /= varNameToAstVarId var) $
case (withKnownSpan (varNameToSpan var2) $ subst u, subst v) of
(Nothing, Nothing) -> Nothing
(mu, mv) -> Just $ astLet var2 (fromMaybe u mu) (fromMaybe v mv)
Ast.AstPrimalPart a -> astPrimalPart <$> subst a
Ast.AstDualPart a -> astDualPart <$> subst a
Ast.AstPlainPart a -> astPlainPart <$> subst a
Ast.AstFromPrimal a -> fromPrimal <$> subst a
Ast.AstFromDual a -> fromDual <$> subst a
Ast.AstFromPlain a -> fromPlain <$> subst a
AstPlusK u v ->
let mu = subst u
mv = subst v
in if isJust mu || isJust mv
then Just $ astPlusK (fromMaybe u mu) (fromMaybe v mv)
else Nothing
AstTimesK u v ->
let mu = subst u
mv = subst v
in if isJust mu || isJust mv
then Just $ astTimesK (fromMaybe u mu) (fromMaybe v mv)
else Nothing
Ast.AstN1K opCode u -> astN1K opCode <$> subst u
Ast.AstR1K opCode u -> astR1K opCode <$> subst u
Ast.AstR2K opCode u v ->
let mu = subst u
mv = subst v
in if isJust mu || isJust mv
then Just $ astR2K opCode (fromMaybe u mu) (fromMaybe v mv)
else Nothing
Ast.AstI2K QuotOp u v ->
let mu = subst u
mv = subst v
in if isJust mu || isJust mv
then Just $ astI2K QuotOp (fromMaybe u mu) (fromMaybe v mv)
else Nothing
Ast.AstI2K RemOp u v ->
let mu = subst u
mv = subst v
in if isJust mu || isJust mv
then Just $ astI2K RemOp (fromMaybe u mu) (fromMaybe v mv)
else Nothing
AstConcreteK{} -> Nothing
Ast.AstFloorK a -> astFloorK <$> subst a
Ast.AstFromIntegralK v -> astFromIntegralK <$> subst v
Ast.AstCastK v -> astCastK <$> subst v
Ast.AstArgMinK v -> astArgMinK <$> subst v
Ast.AstArgMaxK v -> astArgMaxK <$> subst v
Ast.AstIndexK v ix ->
case (subst v, substIxS ix) of
(Nothing, Nothing) -> Nothing
(mv, mix) -> Just $ astIndexK (fromMaybe v mv) (fromMaybe ix mix)
AstPlusS u v ->
let mu = subst u
mv = subst v
in if isJust mu || isJust mv
then Just $ astPlusS (fromMaybe u mu) (fromMaybe v mv)
else Nothing
AstTimesS u v ->
let mu = subst u
mv = subst v
in if isJust mu || isJust mv
then Just $ astTimesS (fromMaybe u mu) (fromMaybe v mv)
else Nothing
Ast.AstN1S opCode u -> astN1S opCode <$> subst u
Ast.AstR1S opCode u -> astR1S opCode <$> subst u
Ast.AstR2S opCode u v ->
let mu = subst u
mv = subst v
in if isJust mu || isJust mv
then Just $ astR2S opCode (fromMaybe u mu) (fromMaybe v mv)
else Nothing
Ast.AstI2S QuotOp u v ->
let mu = subst u
mv = subst v
in if isJust mu || isJust mv
then Just $ astI2S QuotOp (fromMaybe u mu) (fromMaybe v mv)
else Nothing
Ast.AstI2S RemOp u v ->
let mu = subst u
mv = subst v
in if isJust mu || isJust mv
then Just $ astI2S RemOp (fromMaybe u mu) (fromMaybe v mv)
else Nothing
AstConcreteS{} -> Nothing
Ast.AstFloorS a -> astFloorS <$> subst a
Ast.AstFromIntegralS a -> astFromIntegralS <$> subst a
Ast.AstCastS v -> astCastS <$> subst v
Ast.AstArgMinS a -> Ast.AstArgMinS <$> subst a
Ast.AstArgMaxS a -> Ast.AstArgMaxS <$> subst a
Ast.AstIndexS shn t ix ->
case (subst t, substIxS ix) of
(Nothing, Nothing) -> Nothing
(mt, mix) -> Just $ astIndexS shn (fromMaybe t mt) (fromMaybe ix mix)
Ast.AstCondK b v w ->
case (subst b, subst v, subst w) of
(Nothing, Nothing, Nothing) -> Nothing
(mb, mv, mw) ->
Just
$ astCondKInitial (fromMaybe b mb) (fromMaybe v mv) (fromMaybe w mw)
Ast.AstCondS b v w ->
case (subst b, subst v, subst w) of
(Nothing, Nothing, Nothing) -> Nothing
(mb, mv, mw) ->
Just
$ astCondSInitial (fromMaybe b mb) (fromMaybe v mv) (fromMaybe w mw)
Ast.AstFromVectorK shm args ->
let margs = V.map subst args
in if V.any isJust margs
then Just $ astFromVectorK shm $ V.zipWith fromMaybe args margs
else Nothing
Ast.AstFromVectorS shm args ->
let margs = V.map subst args
in if V.any isJust margs
then Just $ astFromVectorS shm $ V.zipWith fromMaybe args margs
else Nothing
Ast.AstSumK v -> astSumK <$> subst v
Ast.AstSumS shm v -> astSumS shm <$> subst v
Ast.AstScatterS shm shn shp t (vars, ix) ->
assert (all (\v -> varNameToAstVarId var /= varNameToAstVarId v) vars) $
case (subst t, substIxS ix) of
(Nothing, Nothing) -> Nothing
(mt, mix) -> Just $ astScatterS shm shn shp
(fromMaybe t mt)
(vars, fromMaybe ix mix)
Ast.AstReplicateK shm v -> astReplicateK shm <$> subst v
Ast.AstReplicateS shm v -> astReplicateS shm <$> subst v
Ast.AstGatherS shm shn shp t (vars, ix) ->
assert (all (\v -> varNameToAstVarId var /= varNameToAstVarId v) vars) $
case (subst t, substIxS ix) of
(Nothing, Nothing) -> Nothing
(mt, mix) -> Just $ astGatherS shm shn shp
(fromMaybe t mt)
(vars, fromMaybe ix mix)
Ast.AstIotaS{} -> Nothing
Ast.AstAppendS x y ->
case (subst x, subst y) of
(Nothing, Nothing) -> Nothing
(mx, my) -> Just $ astAppendS (fromMaybe x mx) (fromMaybe y my)
Ast.AstSliceS i2 n k v -> astSliceS i2 n k <$> subst v
Ast.AstReverseS v -> astReverseS <$> subst v
Ast.AstTransposeS perm v -> astTransposeS perm <$> subst v
Ast.AstReshapeS sh v -> astReshapeS sh <$> subst v
Ast.AstConvert c v -> astConvert c <$> subst v
Ast.AstDot0 u v ->
let mu = subst u
mv = subst v
in if isJust mu || isJust mv
then Just $ astDot0 (fromMaybe u mu) (fromMaybe v mv)
else Nothing
Ast.AstDot1InS sh n u v ->
let mu = subst u
mv = subst v
in if isJust mu || isJust mv
then Just $ astDot1InS sh n (fromMaybe u mu) (fromMaybe v mv)
else Nothing
Ast.AstMatmul2S m n p u v ->
let mu = subst u
mv = subst v
in if isJust mu || isJust mv
then Just $ astMatmul2S m n p (fromMaybe u mu) (fromMaybe v mv)
else Nothing
Ast.AstBoolNotK arg -> astBoolNotK <$> subst arg
Ast.AstBoolNotS arg -> astBoolNotS <$> subst arg
Ast.AstBoolAndK arg1 arg2 ->
let mb1 = subst arg1
mb2 = subst arg2
in if isJust mb1 || isJust mb2
then Just $ astBoolAndK (fromMaybe arg1 mb1) (fromMaybe arg2 mb2)
else Nothing
Ast.AstBoolAndS arg1 arg2 ->
let mb1 = subst arg1
mb2 = subst arg2
in if isJust mb1 || isJust mb2
then Just $ astBoolAndS (fromMaybe arg1 mb1) (fromMaybe arg2 mb2)
else Nothing
Ast.AstLeqK arg1 arg2 ->
let mr1 = subst arg1
mr2 = subst arg2
in if isJust mr1 || isJust mr2
then Just $ astLeqK (fromMaybe arg1 mr1) (fromMaybe arg2 mr2)
else Nothing
Ast.AstLeq arg1 arg2 ->
let mr1 = subst arg1
mr2 = subst arg2
in if isJust mr1 || isJust mr2
then Just $ astLeq (fromMaybe arg1 mr1) (fromMaybe arg2 mr2)
else Nothing
Ast.AstLeqS shb sh arg1 arg2 ->
let mr1 = subst arg1
mr2 = subst arg2
in if isJust mr1 || isJust mr2
then Just $ astLeqS shb sh (fromMaybe arg1 mr1) (fromMaybe arg2 mr2)
else Nothing
substitute1AstIxS
:: AstTensor AstMethodLet s2 y -> AstVarName '(s2, y)
-> AstIxS AstMethodLet sh
-> Maybe (AstIxS AstMethodLet sh)
substitute1AstIxS i var ix =
let mix = fmap (substitute1Ast i var) ix
in if any isJust mix
then Just $ ixsZipWith fromMaybe ix mix
else Nothing
substitute1AstHFun
:: forall s s3 x y z.
AstTensor AstMethodLet s3 z -> AstVarName '(s3, z) -> AstHFun s x y
-> Maybe (AstHFun s x y)
substitute1AstHFun _i var (AstLambda var2 _) =
assert (varNameToAstVarId var2 /= varNameToAstVarId var)
Nothing -- no outside free variables