`collapse_NP` doesn't participate in foldr/build fusion

[jtnuttall]

I was wondering whether there is a reason the collapse implementations should not use GHC.Exts.build directly?

collapse_NP is defined with plain cons/nil:

collapse_NP :: NP (K a) xs -> [a]
collapse_NP Nil         = []
collapse_NP (K x :* xs) = x : collapse_NP xs

Because it doesn't use build, it always produces an intermediate list allocation.

Reproducer

{-# LANGUAGE GHC2024 #-}
{-# OPTIONS_GHC -O2 -ddump-simpl -ddump-rule-firings
    -dsuppress-all -dsuppress-uniques -dno-suppress-type-signatures #-}

module Main (main) where

import Generics.SOP
import GHC.Exts qualified

{-# NOINLINE consumeHcollapse #-}
consumeHcollapse :: SListI xs => NP (K Int) xs -> Int
consumeHcollapse = sum . filter even . hcollapse

{-# NOINLINE consumeWithBuild #-}
consumeWithBuild :: SListI xs => NP (K Int) xs -> Int
consumeWithBuild = sum . filter even . collapse_build

collapse_build :: forall a xs. NP (K a) xs -> [a]
collapse_build np = GHC.Exts.build (\(c :: a -> r -> r) n ->
  let go :: forall ys. NP (K a) ys -> r
      go Nil         = n
      go (K x :* xs) = x `c` go xs
  in go np)

sample :: NP (K Int) '[(), (), ()]
sample = K 1 :* K 2 :* K 3 :* Nil

main :: IO ()
main = do
  print (consumeHcollapse sample)
  print (consumeWithBuild sample)

Rule firings

consumeHcollapse: filter and sum fuse, but collapse_NP does not participate

consumeWithBuild: collapse, filter, and sum fuse

Resulting Core

consumeHcollapse - list allocation, two loops

Rec {
$wgo1 :: [Int] -> Int# -> Int#
$wgo1
  = \ (ds :: [Int]) (ww :: Int#) ->
      case ds of {
        [] -> ww;
        : y ys ->
          case y of { I# ipv ->
          case remInt# ipv 2# of {
            __DEFAULT -> $wgo1 ys ww;
            0# -> $wgo1 ys (+# ww ipv)
          }
          }
      }
end Rec }

$wconsumeHcollapse :: forall k (xs :: [k]). NP (K Int) xs -> Int#
$wconsumeHcollapse
  = \ (@k) (@(xs :: [k])) (eta :: NP (K Int) xs) ->
      $wgo1 (collapse_NP eta) 0#

consumeWithBuild - no list allocation, fused loop

Rec {
$wgo :: forall k (ys :: [k]). NP (K Int) ys -> Int# -> Int#
$wgo
  = \ (@k) (@(ys :: [k])) (ds :: NP (K Int) ys) (ww :: Int#) ->
      case ds of {
        Nil co -> ww;
        :* @x @xs co ds1 xs1 ->
          case ds1 `cast` <Co:4> :: ... of { I# ipv ->
          case remInt# ipv 2# of {
            __DEFAULT -> $wgo xs1 ww;
            0# -> $wgo xs1 (+# ww ipv)
          }
          }
      }
end Rec }

$wconsumeWithBuild :: forall k (xs :: [k]). NP (K Int) xs -> Int#
$wconsumeWithBuild
  = \ (@k) (@(xs :: [k])) (x :: NP (K Int) xs) -> $wgo x 0#

Relationship to #74

This is narrower than the general SOP overhead elimination discussed in #74. It would require a change to collapse_NP and an INLINE pragma on collapse_NP. There's a code bloat concern but I think it might be justifiable in enough cases to make this worth it.

I'd be happy to put together a PR/fork but I wanted to ask whether there was a reason not to do this first.

Tested on GHC 9.14.1, sop-core 0.5.0.2, generics-sop 0.5.1.4.

[kosmikus]

Thanks for the detailed report. I think you're right that the specific case of collapse_NP looks much simpler than trying to look at the whole library.

I'm personally not sure whether it would then be better to just directly implement collapse_NP in terms of build, or use rewrite rules to rewrite to a build-based version. Probably it doesn't matter much.

[jtnuttall]

Thanks for the response.

I went with explicit build usage because I felt like collapse_NP is structurally more similar to unfoldr than map.

I can take a peek at using a rule and see if that makes a difference when fusion fails.

I'll put together a PR when I get the chance. Is there anything additional you'd like to see with the PR?

[jtnuttall]

Hi @kosmikus, checking in on this issue. Is there anything I can do to make #188 more palatable or easier to review? Happy to update it however would be helpful, e.g. trimming it to just the collapse_NP change, benchmarking codegen on wider records and/or modules with many call sites, new inspection tests, etc.