Rust Portable SIMD en GPU: Unificación de Paralelismo a Nivel de Lane

La computación moderna, especialmente en sistemas de alto rendimiento, depende fundamentalmente de la explotación del paralelismo a múltiples niveles. Históricamente, el paralelismo a nivel de instrucción única y múltiples datos (SIMD) en CPUs y el paralelismo a nivel de thread en GPUs (SIMT) han sido abordados con abstracciones y herramientas distintas, lo que complica la portabilidad y la reutilización de código. Este trabajo aborda el problema de unificar la programación SIMD/SIMT al demostrar que la abstracción core::simd de Rust, diseñada para CPUs, puede mapearse directamente a la ejecución de warps en GPUs.

La relevancia de esta unificación radica en la creciente importancia de las GPUs para cargas de trabajo que van desde la inteligencia artificial hasta la simulación científica. Al permitir que el código SIMD portable de Rust se ejecute sin modificaciones en GPUs, se reduce la barrera de entrada para los desarrolladores de Rust que buscan aprovechar el rendimiento de la GPU, y se promueve un modelo de programación más coherente y menos propenso a errores. Este enfoque se alinea con la tendencia de abstracciones de hardware agnósticas que buscan maximizar la eficiencia del desarrollador sin sacrificar el rendimiento nativo del hardware.

Arquitectura del Sistema

La arquitectura central de esta implementación se basa en la observación de que el modelo de ejecución SIMT de las GPUs (Single Instruction, Multiple Thread) es conceptualmente equivalente a SIMD (Single Instruction, Multiple Data). Un 'warp' de GPU, que ejecuta una instrucción en 32 lanes de forma simultánea, se trata como una unidad vectorial ancha. La abstracción Simd<T, N> de Rust, que representa un vector de N elementos de tipo T, se mapea directamente a esta unidad.

Las operaciones element-wise (suma, multiplicación, comparación) en Simd se traducen directamente a las operaciones nativas del warp. Las reducciones SIMD (como reduce_sum, reduce_max) se implementan utilizando las instrucciones de 'warp shuffle' de la GPU, que permiten el intercambio y combinación eficiente de valores entre lanes. Las operaciones de 'cross-lane shuffles' (como simd_swizzle!) también se mapean a estas primitivas de 'warp shuffle'. Las máscaras SIMD (Mask<T, N>) utilizan instrucciones de 'vote' y 'ballot' de la GPU para consultas horizontales. Los valores escalares se replican en todas las lanes del warp, manteniendo la distinción entre datos uniformes y variables. Para manejar anchos de vector N que no coinciden con el ancho fijo del warp (32 en NVIDIA), se utiliza un 'Intermediate Representation' (IR) codificado en el sistema de tipos de Rust. Este IR, compuesto por operaciones tipadas como 'ballots', 'shuffles', 'reductions', 'scans', 'gathers', 'scatters', 'atomics' y 'strip mining', permite que el compilador genere instrucciones PTX (Parallel Thread Execution) de GPU de costo cero, sin necesidad de un intérprete en tiempo de ejecución. Este IR también se ejecuta en CPU para simulación y pruebas diferenciales.

use core::simd::{Simd, Mask};

fn main() {
    let a = Simd::<f32, 32>::splat(1.0);
    let b = Simd::<f32, 32>::splat(2.0);
    let sum = a + b; // Element-wise addition

    let mask = sum.simd_gt(Simd::splat(2.5)); // Comparison, produces a mask
    let selected = mask.select(sum, Simd::splat(0.0)); // Select based on mask

    let total_sum = selected.reduce_sum(); // Horizontal reduction
    println!("Total sum: {}", total_sum);
}

Fundamentos Teóricos

La idea de unificar modelos de programación paralela no es nueva y tiene raíces profundas en la investigación de arquitecturas de computadoras y lenguajes de programación. El concepto de SIMD se formalizó en la taxonomía de Flynn (1966) para clasificar arquitecturas de computadoras. La programación SIMT en GPUs, popularizada por NVIDIA con CUDA, es una evolución de este principio, adaptada a la masiva paralelización de datos. El desafío de mapear abstracciones de alto nivel a hardware SIMD/SIMT ha sido un tema recurrente en la investigación de compiladores y lenguajes de programación paralela.

Este trabajo se conecta con principios de diseño de lenguajes como el 'zero-cost abstraction' de C++ y Rust, donde las características de alto nivel no imponen una sobrecarga de rendimiento en comparación con el código escrito a mano. La utilización de un IR tipado en el sistema de tipos de Rust para la generación de código específico de hardware recuerda a técnicas de 'metaprogramming' y 'domain-specific languages' (DSLs) embebidos, donde las propiedades del programa se verifican en tiempo de compilación. La gestión de la coherencia entre diferentes anchos de vector y warp se relaciona con problemas de 'vectorización' y 'strip mining' en compiladores, que buscan optimizar bucles escalares para ejecución vectorial.