Transitivity of Auto-Specialization in GHC

21,861

Short answers:

The question's key points, as I understand them, are the following:

  • "is the auto-specialization transitive?"
  • Should I only expect (+) to be specialized transitively with an explicit pragma?
  • (apparently intended) Is this a bug of GHC? Is it inconsistent with the documentation?

AFAIK, the answers are No, mostly yes but there are other means, and No.

Code inlining and type application specialization is a trade-off between speed (execution time) and code size. The default level gets some speedup without bloating the code. Choosing a more exhaustive level is left to the programmer's discretion via SPECIALISE pragma.

Explanation:

The optimiser also considers each imported INLINABLE overloaded function, and specialises it for the different types at which it is called in M.

Suppose f is a function whose type includes a type variable a constrained by a type class C a. GHC by default specializes f with respect to a type application (substituting a for t) if f is called with that type application in the source code of (a) any function in the same module, or (b) if f is marked INLINABLE, then any other module that imports f from B. Thus, auto-specialization is not transitive, it only touches INLINABLE functions imported and called for in the source code of A.

In your example, if you rewrite the instance of Num as follows:

instance (Num r, Unbox r) => Num (Qux r) where
    (+) = quxAdd

quxAdd (Qux x) (Qux y) = Qux $ U.zipWith (+) x y
  • quxAdd is not specifically imported by Main. Main imports the instance dictionary of Num (Qux Int), and this dictionary contains quxAdd in the record for (+). However, although the dictionary is imported, the contents used in the dictionary are not.
  • plus does not call quxAdd, it uses the function stored for the (+) record in the instance dictionary of Num t. This dictionary is set at the call site (in Main) by the compiler.
Share:
21,861
crockeea
Author by

crockeea

Cryptography researcher with an interest in homomorphic encryption and lattice cryptography.

Updated on July 08, 2022

Comments

  • crockeea
    crockeea almost 2 years

    From the docs for GHC 7.6:

    [Y]ou often don't even need the SPECIALIZE pragma in the first place. When compiling a module M, GHC's optimiser (with -O) automatically considers each top-level overloaded function declared in M, and specialises it for the different types at which it is called in M. The optimiser also considers each imported INLINABLE overloaded function, and specialises it for the different types at which it is called in M.

    and

    Moreover, given a SPECIALIZE pragma for a function f, GHC will automatically create specialisations for any type-class-overloaded functions called by f, if they are in the same module as the SPECIALIZE pragma, or if they are INLINABLE; and so on, transitively.

    So GHC should automatically specialize some/most/all(?) functions marked INLINABLE without a pragma, and if I use an explicit pragma, the specialization is transitive. My question is: is the auto-specialization transitive?

    Specifically, here's a small example:

    Main.hs:

    import Data.Vector.Unboxed as U
    import Foo
    
    main =
        let y = Bar $ Qux $ U.replicate 11221184 0 :: Foo (Qux Int)
            (Bar (Qux ans)) = iterate (plus y) y !! 100
        in putStr $ show $ foldl1' (*) ans
    

    Foo.hs:

    module Foo (Qux(..), Foo(..), plus) where
    
    import Data.Vector.Unboxed as U
    
    newtype Qux r = Qux (Vector r)
    -- GHC inlines `plus` if I remove the bangs or the Baz constructor
    data Foo t = Bar !t
               | Baz !t
    
    instance (Num r, Unbox r) => Num (Qux r) where
        {-# INLINABLE (+) #-}
        (Qux x) + (Qux y) = Qux $ U.zipWith (+) x y
    
    {-# INLINABLE plus #-}
    plus :: (Num t) => (Foo t) -> (Foo t) -> (Foo t)
    plus (Bar v1) (Bar v2) = Bar $ v1 + v2
    

    GHC specializes the call to plus, but does not specialize (+) in the Qux Num instance which kills performance.

    However, an explicit pragma

    {-# SPECIALIZE plus :: Foo (Qux Int) -> Foo (Qux Int) -> Foo (Qux Int) #-}
    

    results in transitive specialization as the docs indicate, so (+) is specialized and the code is 30x faster (both compiled with -O2). Is this expected behavior? Should I only expect (+) to be specialized transitively with an explicit pragma?


    UPDATE

    The docs for 7.8.2 haven't changed, and the behavior is the same, so this question is still relevant.