BenchProdTools.hs

{-# LANGUAGE TemplateHaskell #-}
{-# OPTIONS_GHC -Wno-missing-export-lists #-}
-- {-# OPTIONS_GHC -ddump-stranal #-}
-- | A contrived benchmark: a product of a list of scalars.
module BenchProdTools where

import Prelude

import Control.DeepSeq (NFData (..))
import Criterion.Main
import Data.Default qualified as Default
import Data.Foldable qualified as Foldable
import Data.List (foldl1')
import Data.Proxy (Proxy (Proxy))
import GHC.Exts (IsList (..), WithDict)
import GHC.TypeLits (KnownNat)
import Test.Inspection
import Type.Reflection (Typeable)

import Data.Array.Nested qualified as Nested
import Data.Array.Nested.Ranked.Shape

import HordeAd
import HordeAd.Core.Adaptor
import HordeAd.Core.Ops

bgroup100, bgroup1000, bgroup1e4, bgroup1e5, bgroup1e6, bgroup1e7, bgroup5e7 :: [Double] -> Benchmark
bgroup100 = envProd 100 $ bgroup "100" . benchProd

bgroup1000 = envProd 1000 $ bgroup "1000" . benchProdShort

bgroup1e4 = envProd 1e4 $ bgroup "1e4" . benchProdShort

bgroup1e5 = envProd 1e5 $ bgroup "1e5" . benchProdShortest

bgroup1e6 = envProd 1e6 $ bgroup "1e6" . benchProdShortest

bgroup1e7 = envProd 1e7 $ bgroup "1e7" . benchProdShortest

bgroup5e7 = envProd 5e7 $ bgroup "5e7" . benchProdShortest
  -- 5e7 == 5 * 10^7 == 0.5 * 10^8 == 0.5e8

envProd :: r ~ Double
        => Rational
        -> (forall n.
            ( SNat n
            , [Concrete (TKScalar r)]
            , IxR n (Concrete (TKScalar r))
            , IxR n (Concrete (TKS '[] r))
            , Concrete (TKS '[n] r) )
            -> Benchmark)
        -> [r]
        -> Benchmark
envProd rat f allxs =
  let k = round rat
  in withSNat k $ \(SNat @k) ->
    env (return $!
      let l = take k allxs
          lt = map sscalar l
      in ( SNat @k
         , map Concrete l
         , fromList (map Concrete l)
         , fromList lt
         , sfromList . fromList $ lt) )
        (f @k)

benchProd
  :: r ~ Double
  => ( SNat n
     , [Concrete (TKScalar r)]
     , IxR n (Concrete (TKScalar r))
     , IxR n (Concrete (TKS '[] r))
     , Concrete (TKS '[n] r) )
  -> [Benchmark]
benchProd ~(snat, list, l, lt, t) = case snat of
  SNat ->
    [ bench "cgrad s MapAccum" $ nf (crevSMapAccum snat) t
    , bench "grad s MapAccum" $ nf (revSMapAccum snat) t
    , bench "cgrad k MapAccum" $ nf (crevScalarMapAccum snat) t
    , bench "grad k MapAccum" $ nf (revScalarMapAccum snat) t
    , bench "cgrad k list" $ nf crevScalarList list
    , bench "grad k list" $ nf revScalarList list
    , bench "cgrad k L" $ nf (crevScalarL snat) l
    , bench "grad k L" $ nf (revScalarL snat) l
    , bench "cgrad k R" $ nf (crevScalarR snat) l
    , bench "grad k R" $ nf (revScalarR snat) l
    , bench "cgrad k NotShared" $ nf (crevScalarNotShared snat) l
    , bench "cgrad s L" $ nf (crevSL snat) lt
    , bench "grad s L" $ nf (revSL snat) lt
    , bench "cgrad s R" $ nf (crevSR snat) lt
    , bench "grad s R" $ nf (revSR snat) lt
    , bench "cgrad s NotShared" $ nf (crevSNotShared snat) lt
    ]

benchProdShort
  :: r ~ Double
  => ( SNat n
     , [Concrete (TKScalar r)]
     , IxR n (Concrete (TKScalar r))
     , IxR n (Concrete (TKS '[] r))
     , Concrete (TKS '[n] r) )
  -> [Benchmark]
benchProdShort ~(snat, list, l, lt, t) = case snat of
  SNat ->
    [ bench "cgrad s MapAccum" $ nf (crevSMapAccum snat) t
    , bench "grad s MapAccum" $ nf (revSMapAccum snat) t
    , bench "cgrad k MapAccum" $ nf (crevScalarMapAccum snat) t
    , bench "grad k MapAccum" $ nf (revScalarMapAccum snat) t
--    , bench "cgrad k list" $ nf crevScalarList list
--    , bench "grad k list" $ nf revScalarList list
    , bench "cgrad k L" $ nf (crevScalarL snat) l
    , bench "grad k L" $ nf (revScalarL snat) l
--    , bench "cgrad k R" $ nf (crevScalarR snat) l
--    , bench "grad k R" $ nf (revScalarR snat) l
    , bench "cgrad k NotShared" $ nf (crevScalarNotShared snat) l
--    , bench "cgrad s L" $ nf (crevSL snat) lt
--    , bench "grad s L" $ nf (revSL snat) lt
--    , bench "cgrad s R" $ nf (crevSR snat) lt
--    , bench "grad s R" $ nf (revSR snat) lt
--    , bench "cgrad s NotShared" $ nf (crevSNotShared snat) lt
    ]

benchProdShortest
  :: r ~ Double
  => ( SNat n
     , [Concrete (TKScalar r)]
     , IxR n (Concrete (TKScalar r))
     , IxR n (Concrete (TKS '[] r))
     , Concrete (TKS '[n] r) )
  -> [Benchmark]
benchProdShortest ~(snat, list, l, lt, t) = case snat of
  SNat ->
    [ bench "cgrad s MapAccum" $ nf (crevSMapAccum snat) t
    , bench "grad s MapAccum" $ nf (revSMapAccum snat) t
    , bench "cgrad k MapAccum" $ nf (crevScalarMapAccum snat) t
    , bench "grad k MapAccum" $ nf (revScalarMapAccum snat) t
--    , bench "cgrad k list" $ nf crevScalarList list
--    , bench "grad k list" $ nf revScalarList list
--    , bench "cgrad k L" $ nf (crevScalarL snat) l
--    , bench "grad k L" $ nf (revScalarL snat) l
--    , bench "cgrad k R" $ nf (crevScalarR snat) l
--    , bench "grad k R" $ nf (revScalarR snat) l
--    , bench "cgrad k NotShared" $ nf (crevScalarNotShared snat) l
--    , bench "cgrad s L" $ nf (crevSL snat) lt
--    , bench "grad s L" $ nf (revSL snat) lt
--    , bench "cgrad s R" $ nf (crevSR snat) lt
--    , bench "grad s R" $ nf (revSR snat) lt
--    , bench "cgrad s NotShared" $ nf (crevSNotShared snat) lt
    ]

-- Another variant, with foldl1' and indexing, would be a disaster.
-- We can define sproduct if this benchmark ends up used anywhere,
-- because the current codomain of gradientFromDelta rules out
-- low-level hacky pipeline tricks that could avoid indexing.
multSMapAccum :: (BaseTensor target, LetTensor target, NumScalar r)
              => SNat n -> target (TKS '[n] r) -> target (TKS '[] r)
multSMapAccum SNat = sfold (*) (sscalar 1)
{-# SPECIALIZE multSMapAccum :: SNat n -> ADVal Concrete (TKS '[n] Double) -> ADVal Concrete (TKS '[] Double) #-}
{-# SPECIALIZE multSMapAccum :: SNat n -> AstTensor AstMethodLet FullSpan (TKS '[n] Double) -> AstTensor AstMethodLet FullSpan (TKS '[] Double) #-}

crevSMapAccum
  :: SNat n -> Concrete (TKS '[n] Double) -> Concrete (TKS '[n] Double)
crevSMapAccum snat@SNat =
  cgrad  @_ @_ @_ @Concrete(kfromS . multSMapAccum snat)

revSMapAccum
  :: SNat n -> Concrete (TKS '[n] Double) -> Concrete (TKS '[n] Double)
revSMapAccum snat@SNat = grad (kfromS . multSMapAccum snat)

multScalarMapAccum :: forall target n r. (ADReady target, NumScalar r)
                   => SNat n -> target (TKS '[n] r) -> target (TKScalar r)
multScalarMapAccum snat = tfold snat STKScalar STKScalar (*) 1
{-# SPECIALIZE multScalarMapAccum :: SNat n -> ADVal Concrete (TKS '[n] Double) -> ADVal Concrete (TKScalar Double) #-}
{-# SPECIALIZE multScalarMapAccum :: SNat n -> AstTensor AstMethodLet FullSpan (TKS '[n] Double) -> AstTensor AstMethodLet FullSpan (TKScalar Double) #-}

crevScalarMapAccum
  :: SNat n -> Concrete (TKS '[n] Double) -> Concrete (TKS '[n] Double)
crevScalarMapAccum snat@SNat =
  cgrad @_ @_ @_ @Concrete (multScalarMapAccum snat)

revScalarMapAccum
  :: SNat n -> Concrete (TKS '[n] Double) -> Concrete (TKS '[n] Double)
revScalarMapAccum snat@SNat = grad (multScalarMapAccum snat)

multScalarList :: (BaseTensor target, NumScalar r)
               => [target (TKScalar r)] -> target (TKScalar r)
multScalarList = foldl1' (*)

crevScalarList
  :: [Concrete (TKScalar Double)] -> [Concrete (TKScalar Double)]
crevScalarList =
  cgrad @_ @_ @_ @Concrete multScalarList

revScalarList
  :: [Concrete (TKScalar Double)] -> [Concrete (TKScalar Double)]
revScalarList =
  grad multScalarList

multScalarL :: (BaseTensor target, NumScalar r)
            => IxR n (target (TKScalar r)) -> target (TKScalar r)
multScalarL = foldl1' (*) . Foldable.toList

crevScalarL
  :: SNat n -> IxR n (Concrete (TKScalar Double))
  -> IxR n (Concrete (TKScalar Double))
crevScalarL snat@SNat =
  withKnownSTK (stkOfIxR (knownSTK @(TKScalar Double)) snat) $
  cgrad @_ @_ @_ @Concrete multScalarL

revScalarL
  :: SNat n -> IxR n (Concrete (TKScalar Double))
  -> IxR n (Concrete (TKScalar Double))
revScalarL snat@SNat =
  withKnownSTK (stkOfIxR (knownSTK @(TKScalar Double)) snat) $
  grad multScalarL

multScalarR :: (BaseTensor target, NumScalar r)
            => IxR n (target (TKScalar r)) -> target (TKScalar r)
multScalarR = foldr1 (*)

crevScalarR
  :: SNat n -> IxR n (Concrete (TKScalar Double))
  -> IxR n (Concrete (TKScalar Double))
crevScalarR snat@SNat =
  withKnownSTK (stkOfIxR (knownSTK @(TKScalar Double)) snat) $
  cgrad @_ @_ @_ @Concrete multScalarR

revScalarR
  :: SNat n -> IxR n (Concrete (TKScalar Double))
  -> IxR n (Concrete (TKScalar Double))
revScalarR snat@SNat =
  withKnownSTK (stkOfIxR (knownSTK @(TKScalar Double)) snat) $
  grad multScalarR

multScalarNotShared :: (BaseTensor target, NumScalar r)
                    => IxR n (ADVal target (TKScalar r))
                    -> ADVal target (TKScalar r)
multScalarNotShared = foldr1 multNotShared

crevScalarNotShared
  :: SNat n -> IxR n (Concrete (TKScalar Double))
  -> IxR n (Concrete (TKScalar Double))
crevScalarNotShared snat@SNat =
  withKnownSTK (stkOfIxR (knownSTK @(TKScalar Double)) snat) $
  cgrad @_ @_ @_ @Concrete multScalarNotShared

multSL :: (BaseTensor target, NumScalar r)
       => IxR n (target (TKS '[] r)) -> target (TKS '[] r)
multSL = foldl1' (*) . Foldable.toList

crevSL
  :: SNat n -> IxR n (Concrete (TKS '[] Double))
  -> IxR n (Concrete (TKS '[] Double))
crevSL snat@SNat =
  withKnownSTK (stkOfIxR (knownSTK @(TKS '[] Double)) snat) $
  cgrad @_ @_ @_ @Concrete (kfromS . multSL)

revSL
  :: SNat n -> IxR n (Concrete (TKS '[] Double))
  -> IxR n (Concrete (TKS '[] Double))
revSL snat@SNat =
  withKnownSTK (stkOfIxR (knownSTK @(TKS '[] Double)) snat) $
  grad (kfromS . multSL)

multSR :: (BaseTensor target, NumScalar r)
       => IxR n (target (TKS '[] r)) -> target (TKS '[] r)
multSR = foldr1 (*)

crevSR
  :: SNat n -> IxR n (Concrete (TKS '[] Double))
  -> IxR n (Concrete (TKS '[] Double))
crevSR snat@SNat =
  withKnownSTK (stkOfIxR (knownSTK @(TKS '[] Double)) snat) $
  cgrad @_ @_ @_ @Concrete (kfromS . multSR)

revSR
  :: SNat n -> IxR n (Concrete (TKS '[] Double))
  -> IxR n (Concrete (TKS '[] Double))
revSR snat@SNat =
  withKnownSTK (stkOfIxR (knownSTK @(TKS '[] Double)) snat) $
  grad (kfromS . multSR)

multSNotShared :: (BaseTensor target, NumScalar r)
               => IxR n (ADVal target (TKS '[] r))
               -> ADVal target (TKS '[] r)
multSNotShared = foldr1 multNotShared

crevSNotShared
  :: SNat n -> IxR n (Concrete (TKS '[] Double))
  -> IxR n (Concrete (TKS '[] Double))
crevSNotShared snat@SNat =
  withKnownSTK (stkOfIxR (knownSTK @(TKS '[] Double)) snat) $
  cgrad @_ @_ @_ @Concrete (kfromS . multSNotShared)

{- TODO: re-enable once -fpolymorphic-specialisation works

-- KnownNat and KnownSpan are tag types, so it's fine not to specialize
-- for them. Some of the other classes come from existential types,
-- some of which it's not advantageous to specialize.
--
-- This is expected to fail with -O0 and to pass with -O1
-- and -fpolymorphic-specialisation.
-- This prevents running benchmarks without optimization, which is a good thing.
inspect $ hasNoTypeClassesExcept 'crevScalarL [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''NumScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default, ''Nested.Elt]
inspect $ hasNoTypeClassesExcept 'revScalarL [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''NumScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default, ''Nested.Elt,      ''KnownSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Storable, ''Integral]
inspect $ hasNoTypeClassesExcept 'crevScalarNotShared [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''NumScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default, ''Nested.Elt]
inspect $ hasNoTypeClassesExcept 'crevSL [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''NumScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default, ''Nested.Elt, ''Nested.Storable,    ''ShareTensor]
inspect $ hasNoTypeClassesExcept 'revSL [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''NumScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default, ''Nested.Elt,      ''KnownSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Storable, ''Integral]
inspect $ hasNoTypeClassesExcept 'crevSMapAccum [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''NumScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default, ''Nested.Elt, ''LetTensor, ''BaseTensor, ''ConvertTensor, ''Boolean, ''CommonTargetEqOrd, ''AllTargetShow, ''ShareTensor]
inspect $ hasNoTypeClassesExcept 'revSMapAccum [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''NumScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default, ''Nested.Elt, ''LetTensor, ''BaseTensor, ''ConvertTensor, ''Boolean, ''CommonTargetEqOrd, ''AllTargetShow, ''ShareTensor,      ''KnownSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Storable, ''Integral]
inspect $ hasNoTypeClassesExcept 'crevScalarMapAccum [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''NumScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default, ''Nested.Elt, ''LetTensor, ''BaseTensor, ''ConvertTensor, ''Boolean, ''CommonTargetEqOrd, ''AllTargetShow, ''ShareTensor]
inspect $ hasNoTypeClassesExcept 'revScalarMapAccum [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''NumScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''NFData, ''Default.Default, ''Nested.Elt, ''LetTensor, ''BaseTensor, ''ConvertTensor, ''Boolean, ''CommonTargetEqOrd, ''AllTargetShow, ''ShareTensor,      ''KnownSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Storable, ''Integral]
-- inspect $ coreOf 'revScalarL
-}