I ve Scale se dá psát rychlý generický kód za použití typeclass
Nedávno jsem tu psal o tom, jak se dá rychle vypočítat Jaccardova podobnost. Rychle bylo klíčové slova a jediná věc na které skutečně záleželo. Právě kvůli rychlosti jsem do knihovny sketches přidal verzi této rutiny (a pár dalších funkcí) specializovanou pro jediný primitivní typ. Výsledek je na jednu stranu rychlý, ale na druhou stranu neflexibilní. Co když budu chtít udělat to samé, ale místo čtyřbajtových intů s osmibajtovými longy? V tom případě mám smůlu a musím duplikovat kód.
Naštěstí Scala nabízí dvě vlastnosti - typeclassy a specializaci, které přesně tento problém v kombinaci s chytrým a agresivním JITem dokážou vyřešit. S trochou snahy je možné psát kód, který je obecný a zároveň stejně rychlý jako verze ručně specializovaná pro každý primitivní typ.
Ručně specializovaná metoda vypadala takhle:
def intersectionSize(a: Array[Int], b: Array[Int]): Int = { var size, ai, bi = 0 while (ai != a.length && bi != b.length) { val av = a(ai) val bv = b(bi) size += (if (av == bv) 1 else 0) ai += (if (av <= bv) 1 else 0) bi += (if (av >= bv) 1 else 0) } size }
Když bych to přepsal na generickou verzi, se signaturou
def intersectionSize[T](a: Array[T], b: Array[T]): Int
pole Array[T]
stále budou poli příslušných primitivních typů. V tomto případě
generický typ nezpůsobí, že obsah polí bude boxován do objektů, ale k přístupu
do nich budou použité statické metody array_apply
a array_length
ze třídy
ScalaRunTime, které mají efektivitě velice daleko.
Ale na tom teď ale příliš nezáleží, protože to stejně nebude fungovat, protože bych v těle
metody požíval porovnání ==
, <=
, >=
na objektech generického typu T
. T
může být cokoli a tedy nemůže garantovat, že na něm budou tyhle operace definovány. To
se dá vyřešit implicitním parametrem, který obsahuje implementaci těchto
operací pro daný typ. Jde o takzvanou typeclass.
def intersectionSize[T](a: Array[T], b: Array[T])(implicit num: Num[T]): Int
Num
je něco jako Numeric ze standardní knihovny. Při zavolání metody
bez posledního bloku patametrů bude automaticky doplněn kompilátorem. V těle metody se to použije následovně:
size += (if (num.eq (av, bv)) 1 else 0) ai += (if (num.lte(av, bv)) 1 else 0) bi += (if (num.gte(av, bv)) 1 else 0)
Tohle už je správně a kompilátor si přestane stěžovat. I když jsem vyřešil
problém genericity, problém s rychlostí stále zůstává. Každý element pole je
autoboxován, jako objekt předán metodám třídy Num
, které je rozbalí a
provedou na něm jednu instrukci užitečné práce.
I když i v tomto případě není všechno ztraceno a JIT může trochu pomoct - když
profilování ukáže, že se metoda volá jen pro jeden typ a je tedy použitá jedna
a ta samá instance Num
, JIT může metody eq
, lte
a gte
inlinovat, po
inlinování je možné odstranit zbytečný autoboxing a dostat něco, co vypadá
celkem efektivně. Problém je v tom, že jde o velice křehkou rovnováhu sil.
Takhle to funguje, když je Num
argument vždycky instancí jedné třídy a callsite pro eq
,
lte
a gte
jsou monomorfní. Když intersectionSize
používám v mnoha
kontextech a s mnoha typy, JIT to s inlinováním přestane mít lehké, najednou
musí počítat se všemi variantami, musí do kódu vkládat rozskoky, použít
inline cache nebo použít klasické volání virtuálních metod. Výsledná binárka začne trpět a efektivita začne klesat.
Ale na druhou stranu, když se JIT rozhodne celou metodu intersectionSize
inlinovat a tedy zkopírovat její tělo na místo odkud je volána, všechno může
být zase dobré. Inlinováním se vytvoří nový kontext a nové callsite metod eq
,
lte
a gte
, které najednou nejsou uvězněné v těle metody a tedy sdílené pro všechna
možná použití intersectionSize
.
A tady do hry vstupuje specializace.
Je třeba specializovat jak typeclass Num
, tak metodu, která ji používá:
trait Num[@specialized(Int, Long, Float, Double) T] { def eq (a: T, b: T): Boolean def gt (a: T, b: T): Boolean def gte(a: T, b: T): Boolean = !gt(b, a) def lt (a: T, b: T): Boolean = gt(b, a) def lte(a: T, b: T): Boolean = !gt(a, b) } def intersectionSize[@specialized(Int, Long, Float, Double) T](a: Array[T], b: Array[T])(implicit num: Num[T]): Int = ???
Signatura metody specializované pro čtyři primitivní typy1 má na délku kilometr, ale každé písmeno za to stojí. Tady se totiž dostávám k jádru pudla.
Specializace vygenerovala nezávislou metodu pro každý typ zvlášť. Verze pro
Int
bude vypadat nějak takhle:
def intersectionSize$mIc$sp(a: Array[Int], b: Array[Int])(implicit num: Num[Int]): Int = { var size, ai, bi = 0 while (ai != a.length && bi != b.length) { val av = a(ai) val bv = b(bi) size += (if (num.eq$mIc$sp (av, bv)) 1 else 0) ai += (if (num.lte$mIc$sp(av, bv)) 1 else 0) bi += (if (num.gte$mIc$sp(av, bv)) 1 else 0) } size }
Jak je vidět, místo eq
se teď volá metoda eq$mIc$sp
specializovaná pro
Int
. Ta je nejen efektivní tím, že jako argument chce nahý primitivní int,
ale také tím, že na jejím místě se vždycky za každých okolností volá jedna
jediná implementace pocházející z jedné typeclassy Num[Int]
. Konkrétní
callsite je tedy zaručeně monomorfní, JIT s ní může zacházet jako se
statickou metodu a velice snadno inlinovat její tělo. A v tomto případě jde o
stabilní a garantovanou optimalizaci, kterou nerozbije zavolání metody
intersectionSize
v naprosto nesouvisející části programu.
JVM v tomto případě odvede tak dobrou práci, že výsledná binárka, kterou vyprodukoval JIT, je k nerozeznání od binárky ručně specializované verze.
Ale zatímco tohle skvěle funguje pro typeclassy, nemusí to vůbec platit pro funkce vyšších řádů. Typeclass má typicky pro každý typ jednu implementaci a v metodě specializované pro jeden typ se použije právě jedna typeclassa. Funkce vyšších řádů fungují naproti tomu většinou přijímají spoustu různých funkcí jako argument a specializace tedy negarantuje dobré inlinování, jen eliminuje zbytečnou režii autoboxingu.
Dále k tématu:
- Jak JVM volá virtuální metody, jaká temná božstva musí vzývat, aby to bylo aspoň trochu rychlé
- Optimistic (Re)specialization - autor v tomto článku ukáže jak dělat specializaci dynamicky a optimisticky, bez toho, aby kompilátor musel přepsat volání metod na jejich specializované verze. V následujícím článku pak použil miniboxing, který dělá to samé a mnohem lépe.
- State of the Specialization - Specializace je spolu s value types jedna z věcí plánovaných do některé budoucí verze JVM.
- Bridging Islands of Specialized Code using Macros and Reified Types
Pozn:
- Používám jen čtyři nejčastější typy, abych omezil množství tříd a velikost bytekódu, kterou specializace vyprodukuje. Tohle je jedna ze slabých stránek specializace, protože je třeba vygenerovat všechny kombinace všech typů a parametrů pro které je specializováno. Z tohoto důvodu bylo navržena alternativa - tzv. miniboxing (viz & viz).