Data.IntMap.unionWithKey is slower than mergeWithKey

[Bodigrim]

Here is an exagerated example

#!/usr/bin/env cabal
{- cabal:
build-depends: base, containers >= 0.8, tasty-bench
default-language: GHC2024
ghc-options: -O2
ghc-options: "-with-rtsopts=-T"
ghc-options: -ddump-simpl -dsuppress-all -dno-suppress-type-signatures -ddump-to-file
-}

import Data.IntMap (IntMap)
import Data.IntMap qualified as IM
import Test.Tasty.Bench
import GHC.Exts (inline)

mkBench :: String -> (IntMap () -> IntMap () -> IntMap ()) -> Benchmark
mkBench name func = bench name $
  nf (\acc -> foldl' func acc (replicate 1000 mempty)) (mempty :: IntMap ())

main :: IO ()
main = defaultMain
  [ mkBench "union" IM.union -- good
  , mkBench "unionWith" (IM.unionWith (<>)) -- bad
  , mkBench "unionWith saturated" (\xs ys -> IM.unionWith (<>) xs ys) -- refuses to inline
  , mkBench "unionWith inlined" (inline IM.unionWith (<>)) -- still refuses to inline
  , mkBench "unionWithKey" (IM.unionWithKey (\_k x y -> x <> y)) -- bad
  , mkBench "unionWithKey saturated" (\xs ys -> IM.unionWithKey (\_k x y -> x <> y) xs ys) -- refuses to inline
  , mkBench "unionWithKey inlined" (inline IM.unionWithKey (\_k x y -> x <> y)) -- still refuses to inline
  , mkBench "mergeWithKey" (IM.mergeWithKey (\_k x y -> Just $ x <> y) id id) -- good
  ]

On my laptop it gives the following measurements:

$ cabal run UnionWith.hs
All
  union:                  OK
    3.59 μs ± 206 ns,   0 B  allocated
  unionWith:              OK
    8.09 μs ± 437 ns, 117 KB allocated
  unionWith saturated:    OK
    8.08 μs ± 434 ns, 117 KB allocated
  unionWith inlined:      OK
    8.12 μs ± 492 ns, 117 KB allocated
  unionWithKey:           OK
    8.17 μs ± 618 ns, 117 KB allocated
  unionWithKey saturated: OK
    8.09 μs ± 412 ns, 117 KB allocated
  unionWithKey inlined:   OK
    8.13 μs ± 545 ns, 117 KB allocated
  mergeWithKey:           OK
    4.10 μs ± 212 ns,   0 B  allocated

Looking at dumped Core, it seems that unionWithKey fails to inline, even despite saturation and inline. Presumably that's simply because for some reason the interface file is missing its unfolding.

Could containers enable -fexpose-all-unfolding in the Cabal file? It's a blunt weapon, but at least it would give a user an option to force inlining despite GHC heuristics, if their use case benefits from it.

[meooow25]

Looking at dumped Core, it seems that unionWithKey fails to inline, even despite saturation and inline. Presumably that's simply because for some reason the interface file is missing its unfolding.

I haven't checked the Core yet but that's probably the reason, it is not marked as INLINE or INLINABLE.

containers/containers/src/Data/IntMap/Internal.hs

Lines 1252 to 1257 in 85a1ab5

	unionWithKey :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
	unionWithKey f m1 m2
	= mergeWithKey' Bin f' id id m1 m2
	where
	f' (Tip k1 x1) (Tip _k2 x2) = Tip k1 (f k1 x1 x2)
	f' _ _ = error "not Tip"

Could containers enable -fexpose-all-unfolding in the Cabal file? It's a blunt weapon, but at least it would give a user an option to force inlining despite GHC heuristics, if their use case benefits from it.

That might be too blunt a weapon, but I think it would be fine to mark functions like this as INLINABLE. That also allows a user to force inline it.

I'm curious about how much this really affects performance, because the cost of having an unknown function to call isn't too bad. Of course, the benchmark here is very artificial, though I have to check where the 117KB allocations are coming from.

[Bodigrim]

INLINABLE should be sufficient, yes. I suspect the rest of the bunch (intersectWithKey, differenceWith, maybe unionsWith also) might benefit from it as well.

It's mostly redundant allocations I'm concerned about, they have a measurable effect when foldMaping a list of almost-empty IntMaps.