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é ver­ze.

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:

Pozn:

  1. 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).

Flattr this!

This entry was posted in JVM, Scala, VM. Bookmark the permalink.

Leave a Reply

Your email address will not be published. Required fields are marked *