Los tensores dispersos son estructuras matemáticas que incluyen vectores, matrices y generalizaciones de dimensiones superiores, caracterizadas por tener muchos ceros. Son cruciales en campos como la computación científica, el procesamiento de señales y el aprendizaje profundo, gracias a su eficiencia en almacenamiento y cálculo. Sin embargo, la manipulación manual de tensores dispersos o mediante bibliotecas existentes puede resultar engorrosa y propensa a errores, además de no escalar adecuadamente con la complejidad de los patrones de dispersión.
La investigación se centra principalmente en formatos de almacenamiento disperso, que son estructuras de datos que almacenan de forma compacta los elementos no nulos y permiten operaciones eficientes. Esto facilita el escalado a tamaños mayores o la resolución de problemas con menos recursos. Sin embargo, no existe un formato disperso óptimo; la mejor elección depende de la distribución de los elementos no nulos, las operaciones y la arquitectura objetivo.
Universal Sparse Tensor (UST)
El Universal Sparse Tensor (UST) separa la dispersión de un tensor de su representación de almacenamiento en memoria. Utiliza un lenguaje específico de dominio (DSL) para describir cómo debe representarse un tensor en memoria, permitiendo a los desarrolladores centrarse únicamente en su dispersión. La inspección en tiempo de compilación o de ejecución del formato elegido para los operandos en operaciones polimórficas agnósticas a la dispersión decide si se despacha a una biblioteca optimizada o se genera automáticamente un código disperso cuando no existe una solución predefinida.
Este artículo se centra en cómo los desarrolladores pueden utilizar el UST para definir formatos de almacenamiento disperso comunes y menos comunes, adaptados a las propiedades específicas de los tensores dispersos en su aplicación. La operatividad con bibliotecas como SciPy, CuPy y PyTorch permite mapear formatos comunes como COO, CSR y DIA al DSL correspondiente del UST.
DSL para formatos de tensor
El DSL de formato de tensor mapea las dimensiones del tensor a niveles de almacenamiento usando una función invertible que define cómo debe almacenarse cada nivel. Incluye:
1. Una secuencia ordenada de especificaciones de dimensión, como (i, j), que proporciona una referencia nombrada a cada dimensión.
2. Una secuencia ordenada de especificaciones de nivel, como (i : denso, j : comprimido), que define lo que se almacena en cada nivel y un tipo de nivel requerido que describe cómo se almacena, incluyendo:
- Propiedades del nivel, como ser único y ordenado.
- Formatos de nivel, tales como:
- denso: Almacenamiento que solo implica indexación similar a un arreglo denso.
- comprimido: Un arreglo de posiciones define las ubicaciones de un nivel almacenado, y un arreglo de coordenadas almacena los índices dentro de cada nivel almacenado.
- singleton: Variante comprimida donde cada coordenada pertenece directamente al padre en el nivel anterior.
- rango: Variante densa donde los valores de coordenada se basan en una compresión en un nivel previo.
Los arreglos de posiciones y coordenadas se almacenan como arreglos irregulares, primero indexados por nivel. Al final de todos los niveles, un solo arreglo de valores contiene los valores numéricos linealizados de todos los elementos almacenados. La universalidad del UST radica en que el DSL puede extenderse para incluir más formatos de almacenamiento.
Ejemplos de UST
Los ejemplos de UST muestran que las posibilidades de almacenamiento de tensores densos y dispersos crecen rápidamente con la dimensionalidad.
Escalares
Un escalar es un tensor de 0 dimensiones, por lo que no hay dimensiones y, por tanto, no hay niveles. En la UST, el arreglo de valores consiste en un solo elemento denso que contiene el valor escalar.
Vectores
Un vector es un tensor de 1 dimensión, que puede interpretarse visualmente como un vector fila o columna. Se presentan dos formas comunes de almacenamiento: denso o comprimido.
La UST permite mapear una dimensión a dos niveles, obteniendo almacenamiento vectorial bloqueado. Esto permite definir bloques de tamaño 4, por ejemplo. Los arreglos de posiciones en el primer nivel indican que se almacenan bloques densos.
Matrices
El almacenamiento de matrices dispersas (2D) ha sido ampliamente estudiado y existen varios formatos de almacenamiento disperso adoptados, cada uno con un nombre único. El UST puede definir la mayoría de ellos, unificando muchas variaciones bajo un solo marco.
Usando solo formatos de nivel denso/comprimido combinados con permutaciones de índices, el DSL del formato de tensor de un tensor de d dimensiones ya da lugar a múltiples combinaciones. Por ejemplo, para d=2, se obtienen ocho formatos de matriz diferentes.
El formato COO se expresa utilizando el DSL, donde el nivel comprimido permite que el mismo índice aparezca varias veces. El formato de rango puede combinarse con expresiones de nivel para definir almacenamiento diagonal y antidiagonal. Esto permite crear variantes que se adaptan a matrices no cuadradas.
Las variaciones de almacenamiento se pueden obtener utilizando bloqueos, donde índices 2D se mapean en niveles 3D o 4D. Esto permite distinguir entre formato por fila y por columna dentro de cada bloque almacenado.
Tensores
Como se mencionó, el uso de tipos de nivel denso/comprimido junto con permutaciones de índice en el DSL del formato de tensor da lugar a diferentes formas de almacenar tensores. Por ejemplo, almacenar un tensor 3D con todos los niveles comprimidos genera el formato CSF (fibra dispersa comprimida).
Este almacenamiento es adecuado para tensores muy dispersos donde los elementos no nulos están distribuidos uniformemente. Se pueden observar configuraciones interesantes al almacenar tensores en lotes, como se muestra en el siguiente ejemplo.
Aprender más
Este artículo ha ilustrado la amplitud de formatos de almacenamiento de tensores que abarca el UST. Se espera la integración del UST con operaciones polimórficas que despachen a bibliotecas optimizadas o generen automáticamente código. Este enfoque establece un ecosistema disperso limpio, fácil de usar y escalable, facilitando la introducción de nuevos esquemas de almacenamiento sin necesidad de codificación explícita.
Para aprender más, regístrate para la sesión GTC 2026 de NVIDIA, Acelerando la computación científica con nvmath-python, donde se presentará una demostración sobre el UST y emocionantes actualizaciones de productos.
