Haskell中如何保持高层函数多态性同时优化低层函数性能？

我从生产项目中提取了以下最小示例，该机器学习项目由线性代数库、深度学习库和应用程序组成。

线性代数库（Matrix模块） #

module Matrix where

import Data.Vector.Storable hiding (sum)

data Matrix a = Matrix { rows :: Int, cols :: Int, items :: Vector a } deriving (Eq, Show, Read)

item :: Storable a => Int -> Int -> Matrix a -> a
item i j m = unsafeIndex (items m) $ i * cols m + j

multiply :: Storable a => Num a => Matrix a -> Matrix a -> Matrix a
multiply a b = Matrix (rows a) (cols b) $ generate (rows a * cols b) (f . flip divMod (cols b)) where
    f (i, j) = sum $ (\ k -> item i k a * item k j b) <$> [0 .. cols a - 1]

深度学习库（Deep模块） #

module Deep where

import Foreign.Storable
import Matrix

transform :: Storable a => Num a => [Matrix a] -> Matrix a -> Matrix a
transform layers batch = foldr multiply batch layers

应用程序代码 #

import qualified Data.Vector.Storable as VS
import Test.Tasty.Bench
import Matrix
import Deep

main :: IO ()
main = defaultMain [bmultiply] where
    bmultiply = bench "bmultiply" $ nf (items . transform layers) batch where
        m k l c = Matrix k l $ VS.replicate (k * l) c :: Matrix Double
        layers = m 256 256 <$> [0.1, 0.2, 0.3]
        batch = m 256 100 0.4

性能痛点 #

我喜欢深度学习库（除FFI调用BLAS的情况外）和线性代数库无需关注Float或Double等具体类型的多态特性，但这也导致未特化时代码使用装箱值，性能比最优情况差约60倍（959 ms vs 16.7 ms）。

目前获得高性能的唯一方式是通过编译器编译指示（compiler pragmas）在整个调用层级强制内联或特化，这非常麻烦：原本仅和multiply函数相关的性能问题会“感染”整个代码库，哪怕是经过5层间接调用、跨多个中间库使用multiply的高层函数，都必须了解底层特化的技术细节。

实际生产代码中受影响的函数远多于这个示例，只要忘记为某个函数添加正确的编译指示，性能立刻下降。此外，开发库时无法预知用户会使用哪些类型，因此提前添加特化编译指示并非可行方案。

更遗憾的是，所有性能关键的紧循环都完全包含在multiply函数内——该函数本身仅被调用少数几次，只要调用multiply时动态拆箱值就不会影响性能，高层机器学习函数完全不需要值被特化和拆箱。我认为应该存在一种方式，在保持高层和中间层函数多态性的同时，将特化请求传递给低层函数。

核心疑问 #

• 在Haskell中通常如何解决这类问题？
• 如果开发一个使用vector包在紧循环中生成高速代码的库，如何在不丢失多态性或强制全量内联的前提下，将高性能传递给库的用户？
• 是否存在一种方式，仅在高层函数中承担多态的代价（即装箱），仅在与需要特化的函数的边界处进行特化和拆箱，而非让特化“感染”整个调用层级？

解决方案 #

1. 针对底层核心函数添加INLINE和SPECIALIZE编译指示 #

对于multiply这类性能关键的底层函数，直接在其定义处添加编译指示，针对常见数值类型预先生成特化版本，同时确保编译器能在调用点内联特化代码：

{-# INLINE multiply #-}
{-# SPECIALIZE multiply :: Matrix Double -> Matrix Double -> Matrix Double #-}
{-# SPECIALIZE multiply :: Matrix Float -> Matrix Float -> Matrix Float #-}
multiply :: Storable a => Num a => Matrix a -> Matrix a -> Matrix a
multiply a b = Matrix (rows a) (cols b) $ generate (rows a * cols b) (f . flip divMod (cols b)) where
    f (i, j) = sum $ (\ k -> item i k a * item k j b) <$> [0 .. cols a - 1]

这样编译器会预先生成指定类型的特化代码，高层函数调用multiply时会自动匹配对应版本，无需在整个调用链添加额外编译指示。对于用户自定义数值类型，开启-O2优化后，编译器会在调用点自动生成特化代码。

2. 启用字典传递优化选项 #

Haskell的多态依赖字典传递实现，开启-O2时GHC会自动尝试特化字典传递的代码。如果优化效果不佳，可以添加-fexpose-all-unfoldings编译选项，让编译器获取更多函数展开信息，提升特化成功率。

3. 采用“多态外壳+特化内核”的设计模式 #

将性能关键逻辑封装为内部特化函数，外部保留多态接口，仅在接口层处理装箱/拆箱和特化调度：

-- 内部特化核心函数，不暴露给用户
{-# INLINE multiplyCoreDouble #-}
multiplyCoreDouble :: Matrix Double -> Matrix Double -> Matrix Double
multiplyCoreDouble a b = Matrix (rows a) (cols b) $ generate (rows a * cols b) (f . flip divMod (cols b)) where
    f (i, j) = sum $ (\ k -> item i k a * item k j b) <$> [0 .. cols a - 1]

{-# INLINE multiplyCoreFloat #-}
multiplyCoreFloat :: Matrix Float -> Matrix Float -> Matrix Float
multiplyCoreFloat a b = Matrix (rows a) (cols b) $ generate (rows a * cols b) (f . flip divMod (cols b)) where
    f (i, j) = sum $ (\ k -> item i k a * item k j b) <$> [0 .. cols a - 1]

-- 外部多态接口，根据类型调度到对应特化版本或通用版本
multiply :: (Storable a, Num a) => Matrix a -> Matrix a -> Matrix a
multiply = case (undefined :: a) of
    (_ :: Double) -> multiplyCoreDouble
    (_ :: Float) -> multiplyCoreFloat
    _ -> multiplyGeneric -- 通用版本，保证正确性的同时兼容所有类型

这种方式让高层函数保持多态性，仅在底层核心逻辑做特化，避免特化需求扩散到整个代码库。

4. 用newtype辅助自定义类型特化 #

如果用户使用自定义数值类型，可建议用newtype包装，并为newtype实例添加特化编译指示。GHC对newtype的特化处理无额外开销，能高效生成优化代码。

内容的提问来源于stack exchange，提问作者Jules