Как правильно использовать прагму GHC SPECIALIZE? (Пример: специализация чистой функции из монадических с использованием Identity.)

В качестве примера предположим, что я хочу написать монадическую и не монадическую карту по спискам. Я начну с монадического:

import Control.Monad
import Control.Monad.Identity

mapM' :: (Monad m) => (a -> m b) -> ([a] -> m [b])
mapM' _ [] = return []
mapM' f (x:xs) = liftM2 (:) (f x) (mapM f xs)

Теперь я хочу повторно использовать код для записи чистого map (вместо повторения кода):

map' :: (a -> b) -> ([a] -> [b])
map' f = runIdentity . mapM' (Identity . f)

Что необходимо сделать, чтобы map' был оптимизирован так, как если бы он был написанный явно как map? В частности:

Нужно ли писать
```
{-# SPECIALIZE mapM' :: (a -> Identity b) -> ([a] -> Identity [b]) #-}
```
или GHC оптимизирует сам map' (полностью исключая Identity)?
Необходимо добавить что-нибудь еще (более прагмы)?
Как я могу проверить, насколько хорошо скомпилированный map' оптимизирован по явно написанному коду для map?

Ответ 1

Ну, давайте спросим самого компилятора.

Компиляция модуля

module PMap where

import Control.Monad
import Control.Monad.Identity

mapM' :: (Monad m) => (a -> m b) -> ([a] -> m [b])
mapM' _ [] = return []
mapM' f (x:xs) = liftM2 (:) (f x) (mapM f xs)

map' :: (a -> b) -> ([a] -> [b])
map' f = runIdentity . mapM' (Identity . f)

с ghc -O2 -ddump-simpl -ddump-to-file PMap.hs (ghc-7.6.1, 7.4.2 производит то же самое, за исключением уникальных имен) создает следующее ядро для map'

PMap.map'
  :: forall a_afB b_afC. (a_afB -> b_afC) -> [a_afB] -> [b_afC]
[GblId,
 Arity=2,
 Caf=NoCafRefs,
 Str=DmdType LS,
 Unf=Unf{Src=<vanilla>, TopLvl=True, Arity=2, Value=True,
         ConLike=True, WorkFree=True, Expandable=True,
         Guidance=IF_ARGS [60 30] 160 40}]
PMap.map' =
  \ (@ a_c) (@ b_d) (f_afK :: a_c -> b_d) (eta_B1 :: [a_c]) ->
    case eta_B1 of _ {
      [] -> GHC.Types.[] @ b_d;
      : x_afH xs_afI ->
        GHC.Types.:
          @ b_d
          (f_afK x_afH)
          (letrec {
             go_ahZ [Occ=LoopBreaker]
               :: [a_c] -> Data.Functor.Identity.Identity [b_d]
             [LclId, Arity=1, Str=DmdType S]
             go_ahZ =
               \ (ds_ai0 :: [a_c]) ->
                 case ds_ai0 of _ {
                   [] ->
                     (GHC.Types.[] @ b_d)
                     `cast` (Sym <(Data.Functor.Identity.NTCo:Identity <[b_d]>)>
                             :: [b_d] ~# Data.Functor.Identity.Identity [b_d]);
                   : y_ai5 ys_ai6 ->
                     (GHC.Types.:
                        @ b_d
                        (f_afK y_ai5)
                        ((go_ahZ ys_ai6)
                         `cast` (<Data.Functor.Identity.NTCo:Identity <[b_d]>>
                                 :: Data.Functor.Identity.Identity [b_d] ~# [b_d])))
                     `cast` (Sym <(Data.Functor.Identity.NTCo:Identity <[b_d]>)>
                             :: [b_d] ~# Data.Functor.Identity.Identity [b_d])
                 }; } in
           (go_ahZ xs_afI)
           `cast` (<Data.Functor.Identity.NTCo:Identity <[b_d]>>
                   :: Data.Functor.Identity.Identity [b_d] ~# [b_d]))
    }

Yup, только cast s, никаких реальных накладных расходов. Вы получаете локального работника go, который действует точно так же, как map.

Подведение итогов: вам нужно только -O2, и вы можете проверить, насколько хорошо оптимизирован код, посмотрев на ядро (-ddump-simpl) или, если вы можете его прочитать, на сборке (-ddump-asm) соответственно, бит LLVM -ddump-llvm).

Наверное, неплохо немного разобраться. Что касается

Нужно ли писать
{-# SPECIALIZE mapM' :: (a -> Identity b) -> ([a] -> Identity [b]) #-}
или GHC оптимизирует сам map' (полностью исключая Identity)?

ответ заключается в том, что если вы используете специализацию в том же модуле, что и общая функция, тогда вам вообще не нужна прагма {-# SPECIALISE #-}, GHC самостоятельно создает специализацию, если она видит какую-либо выгоду в что. В приведенном выше модуле GHC создал правило специализации

"SPEC PMap.mapM' [Data.Functor.Identity.Identity]" [ALWAYS]
    forall (@ a_abG)
           (@ b_abH)
           ($dMonad_sdL :: GHC.Base.Monad Data.Functor.Identity.Identity).
      PMap.mapM' @ Data.Functor.Identity.Identity
                 @ a_abG
                 @ b_abH
                 $dMonad_sdL
      = PMap.mapM'_$smapM' @ a_abG @ b_abH

что также приносит пользу любому использованию mapM' в монаде Identity вне определяющего модуля (если скомпилировано с оптимизацией, а монада распознается как Identity во время запуска правила).

Однако, если GHC не понимает тип специализации достаточно хорошо, он может не видеть никакой пользы и не специализироваться (я не знаю, это достаточно хорошо, чтобы сказать, будет ли он в любом случае пытаться - до сих пор я нашел специализация каждый раз, когда я смотрел).

Если вы хотите быть уверенным, посмотрите на ядро.

Если вам нужна специализация в другом модуле, GHC не имеет никакой причины специализировать функцию при компиляции определяющего модуля, поэтому в этом случае необходима прагма. Вместо прагмы {-# SPECIALISE #-}, требующей специализации для нескольких типов, выбранных вручную, вероятно, лучше - с ghc-7 - использовать прагму {-# INLINABLE #-}, так что (слегка измененный) исходный код становится доступным в импортирующие модули, которые позволяют специализации для любых требуемых типов.

Нужно добавить что-нибудь еще (более прагмы)?

Для разных целей, конечно, могут потребоваться разные прагмы, но, как правило, {#- INLINABLE #-} - это тот, который вам больше всего нужен. И, конечно, {-# RULES #-} может делать магию, которую компилятор не может сделать сам по себе.

Как я могу проверить, насколько хорошо скомпилированный map' оптимизирован по явно написанному коду для map?

Посмотрите на созданный ядро, asm или llvm bitcode, в зависимости от того, что вы лучше понимаете (ядро относительно легко).
Сравнивайте полученный код с рукописной специализацией, если вы не уверены в своей сути и должны знать. В конечном счете, если вы не получите одинаковые промежуточные результаты на каком-то этапе (core/cmm/asm/llvm), бенчмаркинг - единственный способ узнать наверняка.