Arrays de enums eficientes em memória em Zig

 (alic.dev)
3 pontos por GN⁺ 2023-09-19 | 1 comentários | Compartilhar no WhatsApp
  • Ao armazenar em massa enums/tagged unions com variants de tamanhos diferentes, o custo de padding e fragmentação em Vec e HashMap aumenta porque o espaço é reservado com base na maior variant
  • Zig consegue verificar tamanho de campos, alinhamento e discriminant por meio de comptime e reflexão de tipos, e transformar genericamente contêineres de enum com base no layout de memória
  • Um Vec<Enum> simples usa, para cada elemento, espaço equivalente ao tamanho da maior variant; SoA reduz o padding das tags, mas ainda deixa fragmentação de variants na área de valores
  • A AoVA densa, que agrupa variants do mesmo tamanho, reduz no enum de exemplo 15 vetores para 3 clusters de 2, 4 e 8 bytes, mas quando várias variants se misturam na mesma allocation fica difícil iterar com segurança de tipos
  • Macros procedurais em Rust têm dificuldade para acessar informações de tamanho e alinhamento de tipos e também sofrem restrições para calcular comprimentos de arrays genéricos, enquanto o staging com consciência de tipos do Zig evidencia melhor a composabilidade de eficiência de memória em código de sistemas

O espaço desperdiçado por arrays de enums em Rust

  • Um enum/tagged union com variants de tamanhos diferentes precisa reservar memória suficiente para conter a maior variant
  • O enum de exemplo Foo tem variants u8, u16, u32 e u64; por causa da tag e do alinhamento, o tamanho do tipo vira 16 bytes
  • Ao colocar muitos desses enums em um Vec ou HashMap, cada elemento ocupa espaço com base na maior variant, aumentando padding e fragmentação
  • A transformação para struct of arrays (SoA), colocando as tags em uma allocation separada, reduz parte do padding, mas não elimina a fragmentação da área de valores causada pela diferença de tamanho entre variants
  • Em Rust, até é possível criar manualmente uma estrutura de dados para um enum específico, mas criar uma estrutura de dados genérica e o mais eficiente possível em memória para qualquer enum é difícil ou praticamente impossível
    • É difícil aplicar #[derive] de uma proc macro a tipos de terceiros ou type aliases, e a composabilidade é baixa
    • Não há consciência de tipos; desvios baseados em generic_const_expr espalham cláusulas where verbosas pelo grafo de chamadas e não combinam bem com generic type parameters

Por que o problema se destaca em ASTs de compiladores

  • Uma das grandes motivações para arrays de enums eficientes é o uso de memória de ASTs de compiladores
  • ASTs grandes causam latência de memória e eviction de cache durante a compilação, criando um custo significativo para a performance do front-end
  • Em um vídeo de Chandler Carruth sobre o compilador Carbon, é dito que a AST parseada do clang frequentemente consome 50 vezes mais memória do que o código-fonte original
  • Um exemplo em Rust para representar nós de expressão é composto por um enum Expr
    • Unit
    • Number
    • Binary(Operation, ExprId, ExprId)
    • Ident(Symbol)
    • Eval(ExprId, ExprSlice)
    • BlockExpression(ExprId, StatementSlice)
  • Em OCaml, como o sistema de runtime e o GC cuidam do gerenciamento de memória, é possível representar tipos de dados recursivos sem indirection explícito
  • O Vec<Expr> de Rust usa, para todos os elementos, espaço igual a sizeof(Enum), incluindo o tamanho da maior variant, a tag e o padding

Reduzindo fragmentação com SoA e AoVA

  • Quando um enum simples de 3 variants tem membros de 8, 16 e 32 bits, um Vec comum reserva espaço grande para todos os elementos para acomodar a variant de 32 bits e os requisitos de alinhamento
  • Uma melhoria comum é usar tagged index etc. para manter a própria enum variant pequena
    • O crate tagged_index do compilador Rust
    • Casos de small-string optimization
    • É uma otimização usada com frequência em código de alta performance, como runtimes de linguagem, GCs, compiladores, game engines e kernels de SO
  • Também é possível trocar o contêiner e armazenar o discriminant e os valores em allocations separadas, no estilo SoA
    • O compilador Zig self-hosted usa essa abordagem
    • O padding causado pela tag diminui, mas a coleção de valores da union ainda fica com fragmentação de variants
  • A staged compilation do Zig permite criar genericamente um contêiner que realiza a transformação SoA para qualquer tipo
  • Rust precisa depender de proc macros como soa_derive, com a limitação de que não é possível adicionar #[derive] sem alterar o código-fonte de tipos de terceiros

Arrays por variant e clustering por tamanho

  • Para reduzir ainda mais a fragmentação da área de valores, é possível manter um vetor por variant
  • Na inserção, retorna-se um tagged index que contém a tag do enum e o índice dentro do array da respectiva variant
  • Esse padrão é chamado de array of variant arrays (AoVA)
  • AoVA pode ser implementado com proc macro em Rust e com comptime em Zig
  • Quando há muitas variants e várias têm o mesmo tamanho, vetores por variant podem se multiplicar demais
    • O enum Foo do exemplo tem 15 variants
    • A abordagem com vetores por variant adiciona 15 vetores
    • Isso pode aumentar realocações e o número de chamadas de sistema, além de exigir mais memória para amortization em comparação com um Vec ingênuo
    • Os vetores podem ficar espalhados arbitrariamente na memória, aumentando a chance de conflitos de cache
    • O próprio contêiner AoVA também usa bastante memória e pode inflar a struct que o contém
  • Agrupando por tamanho, o enum de exemplo se divide em três clusters: 2 bytes, 4 bytes e 8 bytes
    • c_2: Vec<[u8; 2]> armazena de A a D
    • c_4: Vec<[u8; 4]> armazena de E a I
    • c_8: Vec<[u8; 8]> armazena de J a O
  • A abordagem de AoVA densa pode reduzir o número total de vetores em 80%
  • Ao colocar variants diferentes juntas na mesma allocation, fica difícil iterar sobre o vetor com segurança de tipos
    • O acesso só é possível por meio do tagged pointer criado no momento da inserção
    • Em uma estrutura de árvore baseada em índices flattened que não precise de blind iteration, isso pode ser um trade-off aceitável
  • Se for necessária iteração com segurança de tipos, é possível recolocar a tag aceitando o custo de padding
  • Se o padding for grande demais, pode-se aplicar transformação SoA a cada array de variant, mas nesse caso o número de vetores dobra

A composabilidade de layout de memória criada pelo comptime do Zig

  • O protótipo em Zig está implementado em osmium
  • O ponto central é a reflexão em tempo de compilação que inspeciona, por meio de built-ins do compilador, tipo dos campos, tamanho em bytes, tamanho em bits e discriminant
  • O código de exemplo verifica o tipo com @typeInfo(inner) e processa apenas quando ele é uma union
    • Percorre os campos da union
    • Calcula o espaço necessário com @max(field.alignment, @sizeOf(field.type))
    • Armazena as informações de tamanho em um vetor stack-allocated
    • Constrói o mapping de cada campo da union para o índice do cluster
    • Se não for uma union, gera um compile error
  • O trecho exato de código está nesta fonte
  • Criar o mesmo exemplo com proc macro em Rust é basicamente impossível
    • Proc macros não conseguem acessar informações de size ou alignment de um tipo
    • Mesmo que gerem uma const fn de cálculo de cluster para um enum específico, isso não pode ser usado para especificar o comprimento de arrays de generic types
  • Implementações de contêiner genérico em Rust dificilmente conseguem mudar condicionalmente dependendo de o tipo fornecido ser enum ou struct
  • Em Zig, conceitualmente é possível escolher entre EfficientEnumArray<T> e EfficientStructArray<T> de acordo com T.isEnum()
  • A implementação de AoVA também pode ser escolhida conforme as características do enum
    • Por exemplo, é possível especializar só quando o benefício de colocar variants diferentes juntas for considerado significativo, como ao reduzir o número de vetores em mais de 90%
  • Se a capacidade máxima for conhecida em tempo de compilação, a função geradora de tipos pode determinar a bitwidth necessária para o tagged index
  • Quando esse tagged index é incluído em outra estrutura de dados, por exemplo dentro de outro enum, os bits restantes podem ser aproveitados no discriminant
  • Zig permite especificar concretamente o número necessário de bits, oferecendo eficiência de memória composável, na qual outras partes do código aproveitam essa informação naturalmente
  • Graças à coerção implícita de widening integers, a usabilidade se mantém mesmo ao lidar com APIs de bitwidths diferentes
  • Uma linguagem de programação de sistemas que valorize eficiência e zero-cost abstractions deveria reconsiderar staged programming, especialmente o comptime do Zig

Leituras relacionadas

1 comentários

 
GN⁺ 2023-09-19
Opiniões no Hacker News

  • Há outra estratégia que tem boa eficiência de armazenamento e também preserva a iteração pelos elementos. O primeiro vetor seria a lista de tags, o segundo vetor teria os deslocamentos em bytes de cada elemento, e o terceiro, mais do que um vetor, conteria os dados compactados das variantes apontados pelo segundo vetor.
    Assim, há metade do número de vetores em relação à solução final do autor (6 vs. 3), não se desperdiçam bytes de padding exceto quando necessário por causa de alinhamento, e os dados ficam em sequência na memória independentemente do tipo, permitindo uma iteração amigável ao cache. Também é possível acessar elementos por índice em O(1). No geral, para dados heterogêneos, isso tem características de desempenho parecidas com as de um Vec.

    • Armazenar os deslocamentos em bytes inline é uma boa ideia. Mas, se os offsets ficarem armazenados na memória, a iteração passa a ter uma dependência de dados, o que pode causar uma séria estagnação de memória no pipeline do processador, mesmo sendo amigável ao cache.
    • Se for preciso modificar esse tipo de coleção, no fim é bem provável que você acabe usando diretamente um alocador de memória para lidar com remoções, alterações para variantes maiores e fragmentação.
    • Se o tamanho do offset não for otimizado, ele pode ocupar bastante espaço em comparação com um T pequeno. Por exemplo, é o caso da combinação de um size_t de 64 bits com uint8_t T; tomando cuidado apenas com o tamanho do offset, parece uma abordagem razoável.

  • Fico curioso sobre como essa estrutura de dados AoVA funciona na prática. Do ponto de vista de array, a aritmética de índices pode deixar de fazer sentido; então não se perde o acesso baseado em índice? A iteração também não parece preservar a ordem de inserção.
    Nesse contexto, acho que TLV (tag-length-value), com melhores características de cache, é mais comum. O comprimento pode ser implícito pela tag e, no mínimo, oferece uma iteração para frente significativa. Veja getdents, inotify e mensagens Netlink.

    • Pela legenda da Figura 4, parece correto dizer que o padrão AoVA não se encaixa bem quando é preciso manter a ordem total dos elementos inseridos.
      Em comparação com o layout SoA anterior, surge uma ordem parcial, não uma ordem total. Na inserção, a estrutura retorna um índice marcado que contém tanto a tag do enum quanto o índice dentro do array da variante correspondente. Então o acesso ordenado parece ficar fora do escopo aqui. Armazenar um índice global em cada elemento permitiria recuperar a iteração ordenada, mas ainda não ajudaria no acesso aleatório ordenado e provavelmente resultaria em código com bastante branching.
    • A “tag do enum e o índice dentro do array da variante correspondente” retornados na inserção são essencialmente um ponteiro. Se você quiser iterar, basta armazenar ponteiros em um array na ordem de uso desejada. É o mesmo que faz um programa que aloca memória no heap.
      Armazenar objetos por tamanho também é uma técnica usada em garbage collectors e alocadores genéricos. Dá para obter eficiência por conhecer todos os tamanhos possíveis dos objetos, ou por usar uma estratégia de liberação mais simples, como em uma arena.
    • O fato de AoVA não ter uma ordem total própria para índices pode ser um problema em alguns casos de uso, mas não necessariamente para os nós de AST propostos aqui.
      Nesse caso, os arrays podem ser vistos como um componente de uma estrutura parecida com heap, ou seja, como uma arena. O custo é que o índice precisa se tornar bidimensional, algo como (tag_idx, va_for_tag_idx). Mas, como o número de tags é conhecido em tempo de compilação, dá para otimizar o armazenamento empacotando tag_idx nos 4 ou 5 bits superiores e deixando va_for_tag_idx usar o restante. Referência: https://www.cs.cornell.edu/~asampson/blog/flattening.html
    • Escritas em arrays que mudam o tipo do índice parecem absurdamente caras.

  • É um pouco frustrante que o pattern matching do Rust não possa ser expresso mais como uma espécie de trait do sistema de tipos, que qualquer struct arbitrária poderia implementar, em vez de ser um tipo de objeto explícito, de primeira classe e hardcoded, com sua própria estrutura de armazenamento.
    Recentemente implementei uma AST como a deste artigo e também um interpretador de opcode/bytecode, e senti que os enums do Rust não são totalmente ideais para nenhum dos dois. Na AST, eu queria adicionar atributos de número de linha/coluna a todos os nós de statement, mas colocar linha/coluna em todos os casos do enum Stmt deixava o boilerplate feio; envolver o enum em uma nova struct Stmt que contivesse o enum original junto com os atributos de linha/coluna exigia muita refatoração e não era elegante. No lado dos opcodes, também não dá para dizer que um enum do Rust usado com pattern matching seja uma codificação ideal para o desempenho de um interpretador de opcodes de VM, mas a linguagem empurra nessa direção e os recursos de padrões de desestruturação são muito atraentes. Parece haver espaço para melhorias no sistema de tipos que permitam obter o pattern matching enquanto se usa a implementação de baixo nível desejada.

    • Seria bom ter um exemplo mais concreto. A primeira coisa que vem à mente é algum tipo de sistema de tipos estrutural, mas não tenho certeza se é exatamente isso que você quis dizer.
      https://en.wikipedia.org/wiki/Structural_type_system
    • Uma técnica antiga em interpretadores de bytecode é usar um salto indireto ao passar para a implementação do próximo opcode. O gcc tinha a extensão computed goto para isso; em Rust, imagino que seria necessário algo com ponteiros de função e uma forma de forçar otimização de chamada em cauda.
      Se esse salto indireto for colocado no início da implementação de cada opcode, o preditor de saltos indiretos que a CPU tem por causa de OOP passa a ter modelos separados para o fim de diferentes opcodes, aumentando a taxa de acerto da previsão. A próxima instrução em si pode ser difícil de prever, mas, por exemplo, depois de um test, é muito mais provável que venha um branch. Dito isso, outras técnicas, como manter o topo da pilha em um registrador em uma máquina de pilha, provavelmente são mais importantes, e não sei bem se a técnica acima ainda faz sentido hoje.
    • Várias linguagens têm recursos parecidos com o que você quer. Veja os extractors do Scala ou as active views do F#.

  • A afirmação de que “o AST do clang parseado consome regularmente 50 vezes mais memória que o código-fonte original” parece bem grande, mas o contexto que falta é quanto isso pode melhorar. Se for preciso preservar a posição de origem de cada token e codificar informação suficiente para recuperá-la corretamente a partir do AST, fico curioso se o aumento ideal em relação ao original seria 1,5× ou 15×

    • Por exemplo, se for possível uma redução de 30% de memória, isso seria uma notícia e tanto. Mas, se isso tornar o compilador mais difícil de manter daqui para frente e reduzir só 30%, talvez não valha muito a pena. Por outro lado, se der para obter uma redução de 80% mesmo tratando o compilador de forma um pouco mais agressiva, aí parece valer a tentativa
      É difícil dizer qual seria a taxa ideal de expansão de fonte→AST em uma linguagem que seja amigável tanto para usuários quanto para desenvolvedores de compilador, mas 50× ainda funciona. O texto original usa a taxa de expansão de 50× como motivação para automatizar uma otimização específica. Seria interessante se vetores de enums em Rust pudessem decompor automaticamente valores de enum em tags e valores opacos, para armazená-los como uma estrutura de arrays, como o texto original faz em Zig. Também parece que não haveria muitos lugares onde esconder o uso de unsafe
    • Como comparação, a tape do simdjson é só cerca de 3× maior que o documento original. Uma boa parte disso poderia ser reduzida colocando números em apenas um slot da tape ou, para strings sem sequências de escape, referenciando a posição no documento original em vez de copiá-las
      Em documentos compostos majoritariamente por caracteres [] ou por caracteres 0,, o overhead máximo parece ficar em torno de 8×
    • Código-fonte é surpreendentemente denso. Como um dado sobre até onde isso pode melhorar, há o resultado do parser do próprio Zig parseando o próprio parser do Zig
      Bytes do fonte: 139 KiB, tokens: 24646 (120 KiB), nós do AST: 10998 (140 KiB). Cada token é bem minimalista, com 5 bytes (tag de 1 byte + offset de arquivo de 4 bytes), e os nós do AST também são codificados de forma densa e não uniforme, ficando neste caso em cerca de 13 bytes por nó. Mesmo com essa codificação mínima, a parse tree fica com quase 2× o tamanho do arquivo-fonte. Ainda assim, 2× é muito melhor que 50×. Fonte: zig ast-check -t lib/std/zig/Parse.zig | head -n7
    • É melhor simplesmente assistir à apresentação linkada. É uma ótima apresentação. Pelo que me lembro, acho que ela não apresentava números exatos, e talvez ainda fosse cedo demais para ter certeza. Pode haver dados que ainda não perceberam que são necessários, fazendo os números parecerem menores

  • Esse espaço de problema parece uma variação de problemas de empacotamento
    Seria bom poder começar pela estrutura final, fácil para humanos trabalharem, e gerar recomendações de estruturas de dados que reduzam desperdício de memória, respeitem regras de alinhamento e aumentem a localidade espacial. https://en.wikipedia.org/wiki/Packing_problems

  • Seria bom se proc macros evoluíssem para poder consultar informações do compilador. Isso exigiria um design cuidadoso em relação à adição de fases de compilação, mas coisas como “essa struct implementa este trait?” e “me dê a lista de todos os traits implementados concretamente” costumam ser muito úteis em proc macros

    • Se minha memória não falha, o compilador executa plugins em duas etapas. A primeira recebe o AST antes da checagem de tipos e pode modificar o AST; macros e alguns lints do clippy rodam aí. A segunda etapa é depois da checagem de tipos, então recebe informações de tipo, mas não pode modificar nada; outros lints do clippy rodam aí

  • Entendi só parte do artigo, mas, para alguém que quer escrever um motor de planilhas em Rust, isso parece muito relevante. Valores de célula precisam de algo deste tipo
    pub enum Expr { Number(i32), String(String), Reference(Position), Binary(Box, char, Box), Function(String, Vec), }
    Pretendo continuar lendo e estudando, e referências são bem-vindas

    • O problema destacado aqui é que as variantes têm tamanhos muito diferentes e, quando há muitos desses valores em um array, o desempenho piora por causa do espaço desperdiçado com padding
      Uma técnica comum vinda da área de jogos é dividir um array de structs (AoS) em uma struct de arrays (SoA). Por exemplo, com algo como struct Humans { healths: Vec, ammo: Vec, … }, o índice i de cada vetor corresponde ao i-ésimo Human no layout AoS. Esses vetores paralelos são apenas um exemplo, não a eficiência ótima, porque desperdiçam memória duplicando a contabilidade de tamanho e capacidade para cada campo. Este artigo basicamente tenta aplicar automaticamente uma ideia parecida a enums, e em Rust é difícil fazer isso diretamente. Talvez a importância real desse problema esteja um pouco exagerada. Em uma planilha, eu manteria isso em mente apenas como uma otimização possível e primeiro decidiria se você está construindo para velocidade ou para simplicidade e compreensão
    • Parece um projeto interessante. Se for voltado a usuários finais, espere que alguém coloque conteúdo nos quatro cantos extremos da planilha e veja se o motor cai
      Se você permitir 1 milhão × 1 milhão de células e armazenar null em todas as células não preenchidas, a memória vai acabar. Portanto, vale considerar armazenar o conteúdo das células de forma esparsa. Uma forma é usar uma implementação de hash map como hashbrown. O ponto deste artigo está nos detalhes de baixo nível; se você começar com um hash map desde o início e evitar as restrições iniciais de memória, não precisa pensar muito nisso por enquanto
    • Na verdade, já implementei um motor de planilhas em Rust. Não é open source, mas posso dar alguns conselhos. Antes de se beneficiar da abordagem deste artigo, você encontrará muitos outros problemas de desempenho
      O problema isolado mais difícil é a estratégia de avaliação

  • E quanto a https://doc.rust-lang.org/reference/type-layout.html#the-alignment-modifiers?

  • Acho que há um bug no código de exemplo
    field_map[idx] = svec.len - 1;
    Se svec já contiver size em uma posição que não seja a última entrada, isso vai dar errado