Arreglos de enum eficientes en memoria en Zig

 (alic.dev)
3 puntos por GN⁺ 2023-09-19 | 1 comentarios | Compartir por WhatsApp
  • Al almacenar en masa un enum/tagged union con variantes de distinto tamaño, el espacio reservado según la variante más grande incrementa los costos de padding y fragmentación en Vec y HashMap
  • Zig puede inspeccionar tamaño de campos, alineación y discriminante mediante comptime y reflexión de tipos, y transformar de forma genérica contenedores de enum según el layout de memoria
  • Un Vec<Enum> simple hace que cada elemento use espacio equivalente al de la variante más grande, y SoA reduce el padding del tag pero deja la fragmentación entre variantes en la zona de valores
  • El enfoque de AoVA denso que agrupa variantes del mismo tamaño reduce en el enum de ejemplo 15 vectores a 3 clústeres de 2, 4 y 8 bytes, pero si varias variantes se mezclan en la misma allocation se vuelve difícil iterar con seguridad de tipos
  • En Rust, los proc macro tienen acceso limitado a información de tamaño y alineación de tipos, y también hay restricciones para calcular longitudes de arreglos genéricos, por lo que el staging con conciencia de tipos de Zig muestra mejor la composabilidad de la eficiencia de memoria en código de sistemas

El espacio que desperdician los arreglos de enum en Rust

  • Un enum/tagged union con variantes de distinto tamaño debe reservar suficiente memoria para contener la variante más grande
  • El enum de ejemplo Foo tiene variantes u8, u16, u32 y u64, y por el tag y la alineación el tipo termina midiendo 16 bytes
  • Si se meten muchos de estos enum en Vec o HashMap, cada elemento ocupa espacio según la variante más grande, aumentando el padding y la fragmentación
  • Transformarlo a struct of arrays (SoA), dejando el tag en una allocation separada, reduce parte del padding, pero no elimina la fragmentación en la zona de valores causada por la diferencia de tamaño entre variantes
  • En Rust también es posible construir a mano una estructura de datos para un enum específico, pero crear una estructura genérica lo más eficiente posible en memoria para enums arbitrarios es difícil o casi imposible en la práctica
    • A los proc macro les cuesta aplicar #[derive] sobre tipos de terceros o type alias, y su composabilidad es baja
    • No tienen conciencia de tipos, y los atajos basados en generic_const_expr esparcen cláusulas where verbosas por el grafo de llamadas y no encajan bien con parámetros de tipo genéricos

Por qué el problema se nota más en el AST de un compilador

  • Una de las grandes motivaciones para usar arreglos de enum eficientes es el uso de memoria del AST de un compilador
  • Un AST grande provoca latencia de memoria y eviction de caché durante la compilación, con un costo fuerte para el rendimiento del frontend
  • En un video sobre el compilador de Carbon, Chandler Carruth comenta que el AST parseado de clang suele consumir 50 veces más memoria que el código fuente original
  • Un ejemplo de cómo representar nodos de expresión en Rust usa un enum Expr
    • Unit
    • Number
    • Binary(Operation, ExprId, ExprId)
    • Ident(Symbol)
    • Eval(ExprId, ExprSlice)
    • BlockExpression(ExprId, StatementSlice)
  • En OCaml, el sistema de runtime y el GC se encargan de la memoria, así que se pueden expresar tipos recursivos sin indirection explícita
  • En Rust, Vec<Expr> hace que todos los elementos ocupen sizeof(Enum), incluyendo el tamaño de la variante más grande, el tag y el padding

Reducir la fragmentación con SoA y AoVA

  • Si un enum simple de 3 variantes tiene miembros de 8, 16 y 32 bits, un Vec convencional reserva mucho espacio para todos los elementos para acomodar la variante de 32 bits y sus requisitos de alineación
  • Una mejora común es usar tagged index y técnicas similares para mantener pequeña la propia variante del enum
    • El crate tagged_index del compilador de Rust
    • Casos de small-string optimization
    • Es una optimización frecuente en código de alto rendimiento como runtimes de lenguaje, GC, compiladores, motores de juego y kernels de SO
  • También se puede cambiar el contenedor a un enfoque SoA que guarde el discriminante y los valores en allocations separadas
    • El compilador self-hosted de Zig usa este enfoque
    • Reduce el padding causado por los tags, pero en la colección de valores sigue existiendo fragmentación entre variantes
  • La compilación staged de Zig permite crear de forma genérica contenedores que hagan la transformación SoA para tipos arbitrarios
  • Rust depende de proc macro como soa_derive, con la limitación de que no se puede agregar #[derive] sin modificar el código fuente de tipos de terceros

Arreglos por variante y agrupación por tamaño

  • Para reducir aún más la fragmentación en la zona de valores, se puede usar un vector por variante
  • Al insertar, se devuelve un tagged index que contiene tanto el tag del enum como el índice dentro del arreglo de esa variante
  • A este patrón se le llama array of variant arrays (AoVA)
  • AoVA puede implementarse con proc macro en Rust y con comptime en Zig
  • Si hay muchas variantes y varias comparten el mismo tamaño, tener un vector por variante hace que el número de vectores crezca demasiado
    • El enum Foo de ejemplo tiene 15 variantes
    • El enfoque de un vector por variante agrega 15 vectores
    • Pueden aumentar las reasignaciones y llamadas al sistema, y puede requerir más memoria para amortización que un Vec naive
    • Los vectores pueden quedar dispersos aleatoriamente en memoria, aumentando la posibilidad de conflictos de caché
    • El propio contenedor AoVA también puede usar bastante memoria y agrandar la estructura que lo contiene
  • Si se agrupan por tamaño, el enum de ejemplo se divide en tres clústeres: 2 bytes, 4 bytes y 8 bytes
    • c_2: Vec<[u8; 2]> almacena de A a D
    • c_4: Vec<[u8; 4]> almacena de E a I
    • c_8: Vec<[u8; 8]> almacena de J a O
  • El enfoque de AoVA denso puede reducir en 80% la cantidad total de vectores
  • Si se colocan variantes distintas dentro de la misma allocation, se vuelve difícil iterar el vector con seguridad de tipos
    • El acceso solo es posible mediante el tagged pointer generado al insertar
    • En estructuras de árbol basadas en índices flatten donde no hace falta iteración ciega, puede ser un trade-off aceptable
  • Si se necesita iteración con seguridad de tipos, se puede volver a incluir el tag y aceptar el costo de padding
  • Si el padding sigue siendo demasiado grande, se puede aplicar transformación SoA a cada arreglo de variante, aunque eso duplica la cantidad de vectores

La composabilidad del layout de memoria que habilita comptime en Zig

  • El prototipo en Zig está implementado en osmium
  • La clave es la reflexión en tiempo de compilación mediante built-ins del compilador para inspeccionar tipo de campo, tamaño en bytes, tamaño en bits y discriminante
  • El código de ejemplo verifica el tipo con @typeInfo(inner) y solo procesa el caso union
    • Recorre los campos del union
    • Calcula el espacio necesario con @max(field.alignment, @sizeOf(field.type))
    • Guarda la información de tamaños en un vector asignado en stack
    • Construye un mapping desde campos del union hacia índice de clúster
    • Si no es un union, lanza un error de compilación
  • El fragmento exacto de código está en este source
  • En Rust, crear el mismo ejemplo con un proc macro es básicamente imposible
    • Los proc macro no pueden acceder a información de tamaño o alineación del tipo
    • Se podría generar una const fn que calcule clústeres para un enum específico, pero no se puede usar para definir la longitud de arreglos de un tipo genérico
  • En Rust, es difícil que la implementación de un contenedor genérico cambie condicionalmente según si el tipo dado es enum o struct
  • En Zig, conceptualmente es posible elegir entre EfficientEnumArray<T> y EfficientStructArray<T> según T.isEnum()
  • La implementación de AoVA también puede seleccionarse según características del enum
    • Por ejemplo, se puede especializar para usarla solo cuando mezclar variantes distintas reduzca la cantidad de vectores en más de 90%
  • Si la capacidad máxima se conoce en tiempo de compilación, la función generadora de tipos puede decidir el ancho de bits necesario para el tagged index
  • Si ese tagged index se incluye en otra estructura de datos, por ejemplo dentro de otro enum, los bits sobrantes pueden aprovecharse para el discriminante
  • Zig permite especificar exactamente cuántos bits se necesitan, y otras partes del código pueden aprovechar esa información de manera natural, ofreciendo eficiencia de memoria componible
  • Gracias a la coerción implícita de enteros con widening, la usabilidad se mantiene incluso al trabajar con APIs de distinto ancho de bits
  • Si te interesan la eficiencia y las zero-cost abstractions en lenguajes de programación de sistemas, vale la pena volver a mirar el staged programming, especialmente comptime de Zig

Lecturas relacionadas

1 comentarios

 
GN⁺ 2023-09-19
Opiniones de Hacker News

  • Hay otra estrategia con buena eficiencia de almacenamiento que también conserva la iteración de los elementos. La primera vector contiene la lista de tags, el segundo vector contiene el offset en bytes de cada elemento, y el tercero no es tanto un vector sino los datos comprimidos de las variantes a los que apunta el segundo vector.
    Así se usa la mitad de vectores que la solución final del autor (6 vs. 3), no se desperdician bytes de padding salvo cuando sea necesario por alineación, y los datos quedan en memoria en orden independientemente del tipo, lo que permite recorrerlos de forma amigable con la caché. También permite acceso O(1) a los elementos por índice. En general, para datos heterogéneos tiene características de rendimiento parecidas a Vec.

    • Guardar los offsets en línea es una buena idea. Sin embargo, si los offsets se almacenan en memoria, durante el recorrido aparece una dependencia de datos, y aunque sea amigable con la caché, puede provocar una congestión seria de memoria en el pipeline del procesador.
    • Si hay que modificar una colección así, probablemente se termine usando directamente un asignador de memoria para manejar eliminaciones, cambios a variantes más grandes y fragmentación.
    • Si no se optimiza el tamaño de los offsets, pueden ocupar bastante espacio en comparación con un T pequeño. Por ejemplo, con una combinación de size_t de 64 bits y uint8_t T; si se tiene cuidado solo con el tamaño del offset, parece un enfoque razonable.

  • Me da curiosidad cómo funciona realmente esta estructura de datos AoVA. Desde el punto de vista de un arreglo, como la aritmética de índices quizá ya no tenga sentido, ¿no se pierde el acceso basado en índices? La iteración tampoco parece conservar el orden de inserción.
    En este contexto, creo que es más común usar TLV(tag-length-value), que tiene mejores características de caché. La longitud puede estar implícita en el tag, y al menos ofrece un recorrido hacia adelante con sentido. Véanse getdents, inotify y la mensajería Netlink.

    • Si se mira la leyenda de la figura 4, hay que asumir que el patrón AoVA no encaja bien cuando se necesita conservar el orden total de los elementos insertados.
      Comparado con el diseño SoA anterior, se obtiene un orden parcial, no un orden total. Al insertar, la estructura devuelve un índice etiquetado que contiene tanto el tag del enum como el índice dentro del arreglo de esa variante. Por eso el acceso ordenado parece quedar fuera del alcance aquí. Si se guarda un índice global en cada elemento, se podría recuperar la iteración ordenada, pero seguiría sin ayudar al acceso aleatorio ordenado y probablemente produciría código con bastantes ramas.
    • El “tag del enum y el índice dentro del arreglo de esa variante” que se devuelve al insertar es, en esencia, un puntero. Si se quiere iterar, basta con guardar punteros en un arreglo según el orden de uso deseado. Es lo mismo que hace un programa que asigna memoria en el heap.
      La estrategia de guardar objetos por tamaño también se usa en recolectores de basura y asignadores generales. Se puede ganar eficiencia porque se conocen todos los tamaños posibles de los objetos, y también mediante una liberación más simple, como en un arena.
    • Que AoVA no tenga su propio orden total para los índices puede ser un problema en algunos casos de uso, pero no necesariamente para los nodos AST propuestos aquí.
      En este tipo de caso, los arreglos pueden verse como un componente de una estructura parecida al heap, es decir, como un arena. El costo es que el índice debe volverse bidimensional, algo como (tag_idx, va_for_tag_idx). Pero como la cantidad de tags se conoce en tiempo de compilación, se puede optimizar el almacenamiento empaquetando tag_idx en los 4 o 5 bits superiores y dejando que va_for_tag_idx use el resto. Referencia: https://www.cs.cornell.edu/~asampson/blog/flattening.html
    • Escribir en un arreglo cambiando el tipo del índice parece algo enormemente costoso.

  • Es algo lamentable que el pattern matching de Rust no pueda expresarse más como una especie de trait del sistema de tipos que puedan seguir estructuras arbitrarias, en lugar de ser un tipo de objeto explícito, de primera clase y hardcodeado con su propia estructura de almacenamiento.
    Hace poco implementé un AST como en este artículo y también un intérprete de opcodes/bytecode, y sentí que los enums de Rust no eran del todo ideales para ninguno de los dos. En el AST quería agregar propiedades de línea/columna a todos los nodos de statement, pero poner línea/columna en todos los casos del enum Stmt ensuciaba mucho con boilerplate; y envolver el enum en una nueva estructura Stmt que contuviera el enum original junto con las propiedades de línea/columna implicaba mucho refactoring y no era elegante. Del lado de los opcodes, tampoco diría que un enum de Rust con pattern matching sea la codificación ideal para el rendimiento de un intérprete de opcodes de VM, pero el lenguaje empuja en esa dirección y las capacidades de patrones de destructuración son muy atractivas. Parece haber margen para mejorar el sistema de tipos de modo que se pueda obtener pattern matching usando la implementación de bajo nivel que uno quiera.

    • Sería bueno tener un ejemplo más concreto. Lo primero que se me viene a la mente es una especie de sistema de tipos estructural, aunque no estoy seguro de haberlo entendido exactamente así.
      https://en.wikipedia.org/wiki/Structural_type_system
    • Una técnica antigua en intérpretes de bytecode es usar saltos indirectos para pasar a la implementación del siguiente opcode. gcc tenía la extensión computed goto para esto, y en Rust probablemente haría falta algo que obligue a usar punteros a función y optimización de llamadas de cola.
      Si se pone un salto indirecto de este tipo al principio de la implementación de cada opcode, el predictor de saltos indirectos que la CPU tiene por OOP puede tener modelos separados para los finales de distintos opcodes, aumentando la tasa de aciertos de predicción. Aunque la siguiente instrucción en sí sea difícil de predecir, por ejemplo, después de un test puede ser mucho más probable que venga un branch. Dicho esto, otras técnicas, como guardar el tope de la pila en un registro en una máquina de pila, probablemente sean más importantes, y no estoy seguro de que la técnica anterior siga siendo relevante hoy.
    • Varios lenguajes tienen funcionalidades parecidas a lo que buscas. Mira los extractors de Scala o las active views de F#.

  • Que “el AST de clang parseado consume regularmente 50 veces más memoria que el código fuente original” suena bastante grande, pero el contexto que falta es cuánto puede mejorar. Si hay que preservar la ubicación en el fuente de cada token y codificar suficiente información para poder reconstruirla correctamente desde el AST, me pregunto si el aumento ideal respecto del original sería de 1.5 veces o de 15 veces

    • Por ejemplo, si se pudiera lograr una reducción de memoria del 30%, sería una noticia bastante importante. Pero si para reducir solo 30% se vuelve más difícil mantener el compilador en el futuro, quizá no valga mucho la pena. En cambio, si tratando al compilador con algo menos de cuidado se logra una reducción del 80%, valdría la pena intentarlo
      Es difícil decir cuál sería la tasa ideal de expansión de fuente→AST en un lenguaje que sea amigable tanto para usuarios como para desarrolladores del compilador, pero 50x funciona. El texto original usa esa tasa de expansión de 50x como motivación para automatizar una optimización específica. Sería interesante si los vectores de enum de Rust pudieran descomponer automáticamente los valores enum en tags y valores opacos, para almacenarlos como una estructura de arreglos, como hace el texto original en Zig. Tampoco parece que hubiera muchos lugares donde esconder el uso de unsafe
    • Como comparación, simdjson tape es apenas unas 3 veces más grande que el documento original. Gran parte de eso podría reducirse si se pusieran los números en un solo slot de la tape, o si las cadenas sin secuencias de escape no se copiaran y en su lugar se referenciaran posiciones del documento original
      En documentos compuestos mayormente por caracteres [] o por caracteres 0,, el overhead máximo parece rondar las 8 veces
    • El código fuente es sorprendentemente denso. Como un dato sobre hasta dónde se puede mejorar, está el resultado de que el parser propio de Zig parsee el parser propio de Zig
      Bytes del fuente: 139 KiB, tokens: 24646 (120 KiB), nodos AST: 10998 (140 KiB). Cada token está bastante minimizado, con 5 bytes (tag de 1 byte + offset de archivo de 4 bytes), y los nodos AST también están codificados de forma densa y no uniforme, en este caso alrededor de 13 bytes por nodo. Incluso con esta codificación mínima, el parse tree queda en casi el doble del tamaño del archivo fuente. Aun así, 2x es mucho mejor que 50x. Fuente: zig ast-check -t lib/std/zig/Parse.zig | head -n7
    • Lo mejor es simplemente ver la charla enlazada. Es una gran charla. Si no recuerdo mal, no daba números exactos, y quizá todavía era demasiado temprano para estar seguros. Podrían faltar datos que aún no se habían dado cuenta de que necesitaban, haciendo que los números salieran más bajos

  • Este espacio de problemas se siente como una variante de un problema de empaquetamiento
    Sería bueno poder empezar desde la estructura final, cómoda para humanos, y generar recomendaciones de estructuras de datos que reduzcan el desperdicio de memoria, respeten las reglas de alineación y aumenten la localidad espacial. https://en.wikipedia.org/wiki/Packing_problems

  • Ojalá proc macro evolucione para poder consultar información al compilador. Para eso haría falta un diseño cuidadoso respecto de agregar fases de compilación, pero cosas como “¿esta estructura implementa este trait?” o “dame la lista completa de traits implementados concretamente” suelen ser muy útiles en una proc macro

    • Si no recuerdo mal, el compilador ejecuta los plugins en dos fases. La primera recibe el AST antes del type checking y puede modificarlo; ahí se ejecutan las macros y algunos lints de clippy. La segunda fase es después del type checking, así que recibe información de tipos pero no puede modificar, y ahí se ejecutan otros lints de clippy

  • Entendí solo una parte del artículo, pero desde la perspectiva de alguien que quiere escribir un motor de hojas de cálculo en Rust, parece muy relevante. Los valores de las celdas necesitan algo de esta forma
    pub enum Expr { Number(i32), String(String), Reference(Position), Binary(Box, char, Box), Function(String, Vec), }
    Pienso seguir leyendo y estudiando, y cualquier material de referencia es bienvenido

    • El problema que se destaca aquí es que las variants tienen tamaños muy distintos y, si hay muchos de estos valores en un arreglo, el rendimiento empeora por el espacio desperdiciado en padding
      Una técnica común que viene del mundo de los juegos es dividir un arreglo de estructuras (AoS) en una estructura de arreglos (SoA). Por ejemplo, con algo como struct Humans { healths: Vec, ammo: Vec, … }, el índice i de cada vector corresponde al i-ésimo Human del layout AoS. Estos vectores paralelos son solo un ejemplo y no son óptimos en eficiencia, porque duplican el libro de contabilidad de longitud y capacidad para cada campo, lo que desperdicia espacio. Este artículo básicamente intenta aplicar automáticamente una idea similar a los enum, y en Rust es difícil hacerlo tal cual. Es posible que la magnitud real de este problema esté algo exagerada. En una hoja de cálculo, conviene tenerlo presente solo como una posible optimización y primero decidir si se está construyendo para velocidad o para simplicidad y comprensión
    • Suena como un proyecto interesante. Si apunta a usuarios generales, hay que asumir que alguien pondrá contenido en las cuatro esquinas extremas de la hoja y verá si el motor se cae
      Si se permiten 1 millón × 1 millón de celdas y se almacena null en todas las celdas vacías, la memoria se agota. Por lo tanto, se podría considerar almacenar el contenido de las celdas de forma dispersa. Una forma es usar una implementación de hash map como hashbrown. El punto de este artículo son detalles de bajo nivel, así que si se empieza con un hash map desde el principio y se evitan las restricciones iniciales de memoria, no hace falta pensarlo demasiado por ahora
    • De hecho, alguna vez hice un motor de hojas de cálculo en Rust. No es de código abierto, pero puedo dar algunos consejos. Antes de beneficiarse con el método de este artículo, aparecerán muchos otros problemas de rendimiento
      El problema individual más difícil es la estrategia de evaluación

  • ¿Qué tal https://doc.rust-lang.org/reference/type-layout.html#the-alignment-modifiers?

  • Creo que hay un bug en el código de ejemplo
    field_map[idx] = svec.len - 1;
    Si svec ya contiene size en una entrada que no es la última, eso estará mal