Aceleración de Operaciones Tensoriales en Java con HAT y Project Babylon para GPUs Heterogéneas

La creciente demanda de rendimiento en cargas de trabajo intensivas en matrices, como la inferencia y el entrenamiento de LLMs, ha impulsado la integración de unidades de multiplicación-acumulación de matrices (MMA) dedicadas, como los Tensor Cores de NVIDIA, en el hardware de GPU. Sin embargo, el acceso a estas capacidades de hardware desde lenguajes de alto nivel como Java presenta un desafío de portabilidad significativo, ya que las APIs nativas (CUDA, Metal, oneAPI) son específicas de cada proveedor y de bajo nivel. Este trabajo aborda el problema fundamental de la computación de cómo abstraer y portar operaciones tensoriales de alto rendimiento a través de arquitecturas de aceleradores heterogéneas, permitiendo a los desarrolladores de Java aprovechar estas unidades especializadas sin sacrificar la portabilidad del código.

Históricamente, la programación de GPU ha requerido un conocimiento profundo de arquitecturas específicas y APIs de bajo nivel, lo que limita la adopción en ecosistemas como Java. La tesis central es que, mediante el uso de técnicas de reflexión de código y transformación en tiempo de compilación (JIT), es posible construir un marco que ofrezca una API de tensores de alto nivel y portable en Java, que se optimice dinámicamente para el hardware subyacente. Esto permite que el mismo código Java se ejecute eficientemente en Tensor Cores de NVIDIA, mapeando a las instrucciones HMMA, o se degrade elegantemente a un procesamiento de tiles explícito en plataformas OpenCL 1.2 sin soporte directo de tensores, como las GPUs Apple M4 Max.

Arquitectura del Sistema

La solución se articula alrededor de la extensión del Heterogeneous Accelerator Toolkit (HAT) con una API de tensores y transformaciones de código que utilizan la API de reflexión de OpenJDK Project Babylon. El sistema permite a los desarrolladores de Java definir operaciones tensoriales de multiplicación-acumulación (MMA) utilizando una API de alto nivel (Tensor.mma, Tensor.loadF16, Tensor.store).

Cuando el backend de HAT es CUDA y el hardware subyacente es una GPU NVIDIA con Tensor Cores, el compilador JIT de HAT, asistido por Project Babylon, transforma las llamadas a la API de tensores de Java en invocaciones directas a la API wmma (Warp Matrix Multiply Accumulate) de CUDA. Esto incluye la declaración de wmma::fragment para los tensores de entrada y el acumulador, y el uso de wmma::load_matrix_sync, wmma::mma_sync y wmma::store_matrix_sync. La alineación de memoria a 32 bytes para los arrays de entrada F16Array es crítica y se maneja a nivel de Schema en HAT para optimizar el acceso a memoria global y evitar penalizaciones de rendimiento.

Para backends como OpenCL 1.2 (ej. Apple M4 Max), donde no hay soporte explícito para Tensor Cores o APIs wmma, el compilador de HAT implementa una estrategia de "loop-tiles". Las operaciones tensoriales se bajan a bucles anidados explícitos que realizan la multiplicación y acumulación de matrices en tiles pequeños, mapeando estos tiles a la memoria privada del acelerador. La portabilidad del tamaño de warp (kc.wrs) y el tamaño de tile se gestiona a través de una extensión de la API NDRange de HAT, permitiendo que el runtime recalcule el número total de hilos a desplegar para cada backend (OpenCL vs. CUDA) sin requerir cambios en el código fuente Java. Esto se logra mediante la especialización del código en tiempo de compilación, donde kc.wrs se convierte en una constante (ej., 32 para NVIDIA) o un valor apropiado para el backend de OpenCL (ej., 1 para OpenCL 1.2 si no hay subgroups).

@Reflect
public static void hatTensors(@RO KernelContext kc,
@RO F16Array matrixA,
@RO F16Array matrixB,
@WO F32Array matrixC,
int size) {
final int shapeSize = 16;
final int WMMA_M = shapeSize;
final int WMMA_N = shapeSize;
final int WMMA_K = shapeSize;

int warpM = kc.gix / kc.wrs;
int warpN = kc.giy;

final int lda = size;
final int ldb = size;
final int ldc = size;

var shape = Tensor.shape(WMMA_M, WMMA_N, WMMA_K);
Tensor acc = Tensor.zeros(shape, float.class);

int aRow = warpM * WMMA_M;
int bCol = warpN * WMMA_N;

for (int i = 0;
i < size && aRow < size && bCol < size;
i += WMMA_K) {
Tensor tensorA = Tensor.loadF16(matrixA,
aRow, i, lda, shape);
Tensor tensorB = Tensor.loadF16(matrixB,
i, bCol, ldb, shape);
acc = Tensor.mma(tensorA, tensorB, acc);
}

int cRow = warpM * WMMA_M;
int cCol = warpN * WMMA_N;
if (cRow < size && cCol < size) {
Tensor.store(matrixC, cRow, cCol, acc, ldc);
}
}

HAT_KERNEL void hatTensors(
HAT_GLOBAL_MEM KernelContext_t* kc,
HAT_GLOBAL_MEM F16Array_t* matrixA,
HAT_GLOBAL_MEM F16Array_t* matrixB,
HAT_GLOBAL_MEM F32Array_t* matrixC,
int size
){
int sizeShape = 16;
int WMMA_M = sizeShape;
int WMMA_N = sizeShape;
int WMMA_K = sizeShape;
int warpM = HAT_GIX/32;
int warpN = HAT_GIY;
int lda = size;
int ldb = size;
int ldc = size;
nvcuda::wmma::fragment<
nvcuda::wmma::accumulator, 16, 16, 16, float> acc;
nvcuda::wmma::fill_fragment(acc,0.0);
nvcuda::wmma::fragment<
nvcuda::wmma::matrix_a, 16, 16, 16, half,
nvcuda::wmma::row_major> tensorA;
nvcuda::wmma::fragment
<nvcuda::wmma::matrix_b, 16, 16, 16, half,
nvcuda::wmma::row_major> tensorB;
int aRow = warpM*WMMA_M;
int bCol = warpN*WMMA_N;
for(int i = 0;
i < size && aRow < size && bCol < size;
i=i+WMMA_K) {
nvcuda::wmma::load_matrix_sync(tensorA,
(half*)matrixA->array + i+(aRow*lda),lda);
nvcuda::wmma::load_matrix_sync(tensorB,
(half*)matrixB->array + bCol+(i*ldb),ldb);
nvcuda::wmma::mma_sync(acc,tensorA,tensorB,acc);
}
int cRow = warpM*WMMA_M;
int cCol = warpN*WMMA_N;
if (cRow < size && cCol < size) {
nvcuda::wmma::store_matrix_sync(
matrixC->array + cCol+(cRow*ldc),
acc,ldc,nvcuda::wmma::mem_row_major);
}
return;
}

Fundamentos Teóricos

Este trabajo se conecta con los principios fundamentales de la computación paralela y la optimización de compiladores para arquitecturas heterogéneas. La idea de mapear operaciones de alto nivel a instrucciones de hardware especializadas resuena con el concepto de "domain-specific languages" (DSLs) y la "meta-programación", donde el código se genera o transforma en tiempo de ejecución o compilación para adaptarse mejor al hardware. Project Babylon de OpenJDK, con su API de reflexión de código, es una implementación moderna de estos principios, permitiendo la introspección y manipulación del código Java para generar código nativo optimizado.

La estrategia de "loop-tiles" para la portabilidad en OpenCL 1.2 se basa en técnicas clásicas de "tiling" o "blocking" en la optimización de compiladores para álgebra lineal, como se describe en trabajos seminales sobre optimización de BLAS (Basic Linear Algebra Subprograms). Estas técnicas buscan mejorar la localidad de los datos (cache locality) y el paralelismo explotando la jerarquía de memoria y las unidades de procesamiento vectorial o matricial. Aunque no se cita un paper específico, la filosofía de la transformación de código para adaptarse a las capacidades del hardware (Tensor Cores vs. bucles explícitos) es un tema recurrente en la investigación de compiladores para GPUs y sistemas heterogéneos, buscando un equilibrio entre rendimiento y portabilidad, un trade-off fundamental en el diseño de sistemas distribuidos y paralelos.