Hardware-software co-design for accelerating sparse tensor algebra

Abstract

(English) Many modern scientific and machine learning workloads exhibit inherently sparse and irregular structure. However, traditional supercomputers and datacenters are optimized for dense, regular computation. As a result, sparse workloads use only a fraction of available system resources, leaving significant performance untapped. Although prior solutions have been proposed, they typically target narrow classes of sparse tensor algebra operations, do not fully leverage modern vector compute cores, or leave much of the programmability burden to the user, limiting widespread adoption. To address this gap, this thesis presents a hardware–software co-design for the efficient execution of general sparse tensor algebra operations. We begin by analyzing the architectural implications of sparsity across a broad range of scientific and machine learning applications. Our analysis reveals that traditional compute cores are fundamentally ill-suited to the irregular tensor traversal and coordinate merging required to operate on sparse tensor formats. Building on these insights, we design a Decoupled Access–Execute (DAE) architecture that offloads these operations to a specialized unit, the Tensor Marshaling Unit (TMU), while leaving computation on the core. We demonstrate that the resulting DAE architecture outperforms traditional CPUs and GPUs by up to an order of magnitude. To deliver these gains without exposing the complexity of DAE programming to end users, we also present Ember, a compiler that lowers sparse machine learning operations to DAE code for architectures such as TMU–CPU systems. Ember is implemented in MLIR, enabling natural integration with PyTorch and TensorFlow. By automating this mapping, Ember unlocks the full performance potential of the TMU without imposing an additional programmability burden. (Català) Moltes càrregues de treball científiques i d’aprenentatge automàtic modernes presenten una estructura inherentment dispersa i irregular. Tanmateix, els superordinadors i centres de dades tradicionals estan optimitzats per a computació densa i regular. Com a resultat, les càrregues de treball disperses només utilitzen una fracció dels recursos disponibles del sistema, deixant sense aprofitar un rendiment significatiu. Tot i que s’han proposat solucions prèvies, aquestes solen adreçar-se a classes estretes d’operacions d’àlgebra de tensors dispersos, no aprofiten plenament els nuclis vectorials moderns, o deixen gran part de la càrrega de programabilitat a l’usuari, cosa que en limita l’adopció generalitzada. Per abordar aquesta mancança, aquesta tesi presenta un codisseny hardware–software per a l’execució eficient d’operacions generals d’àlgebra de tensors dispersos. Comencem analitzant les implicacions arquitectòniques de la dispersió en una àmplia gamma d’aplicacions científiques i d’aprenentatge automàtic. La nostra anàlisi revela que els nuclis de còmput tradicionals són fonamentalment poc adequats per al recorregut irregular de tensors i la fusió de coordenades necessaris per operar sobre formats de tensors dispersos. A partir d’aquestes observacions, dissenyem una arquitectura Decoupled Access–Execute (DAE) que descarrega aquestes operacions a una unitat especialitzada, la Tensor Marshaling Unit (TMU), mentre manté la computació al nucli. Demostrem que l’arquitectura DAE resultant supera les CPU i GPU tradicionals fins a un ordre de magnitud. Per oferir aquests guanys sense exposar la complexitat de la programació DAE als usuaris finals, també presentem Ember, un compilador que tradueix operacions disperses d’aprenentatge automàtic a codi DAE per a arquitectures com ara els sistemes TMU–CPU. Ember està implementat en MLIR, fet que permet una integració natural amb PyTorch i TensorFlow. Mitjançant l’automatització d’aquesta transformació, Ember desbloqueja tot el potencial de rendiment de la TMU sense imposar cap càrrega addicional de programabilitat. (Español) Muchas cargas de trabajo científicas y de aprendizaje automático modernas presentan una estructura inherentemente dispersa e irregular. Sin embargo, los superordenadores y centros de datos tradicionales están optimizados para computación densa y regular. Como resultado, las cargas de trabajo dispersas utilizan solo una fracción de los recursos disponibles del sistema, dejando sin explotar un rendimiento significativo. Aunque se han propuesto soluciones previas, estas suelen dirigirse a clases estrechas de operaciones de álgebra de tensores dispersos, no aprovechan plenamente los núcleos vectoriales modernos, o dejan gran parte de la carga de programabilidad al usuario, lo que limita su adopción generalizada. Para abordar esta carencia, esta tesis presenta un codiseño hardware–software para la ejecución eficiente de operaciones generales de álgebra de tensores dispersos. Comenzamos analizando las implicaciones arquitectónicas de la dispersión en una amplia gama de aplicaciones científicas y de aprendizaje automático. Nuestro análisis revela que los núcleos de cómputo tradicionales son fundamentalmente poco adecuados para el recorrido irregular de tensores y la fusión de coordenadas necesarios para operar sobre formatos de tensores dispersos. A partir de estos resultados, diseñamos una arquitectura Decoupled Access–Execute (DAE) que descarga estas operaciones a una unidad especializada, la Tensor Marshaling Unit (TMU), mientras mantiene la computación en el núcleo. Demostramos que la arquitectura DAE resultante supera a las CPU y GPU tradicionales hasta en un orden de magnitud. Para ofrecer estas mejoras sin exponer la complejidad de la programación DAE a los usuarios finales, también presentamos Ember, un compilador que traduce operaciones dispersas de aprendizaje automático a código DAE para arquitecturas como los sistemas TMU–CPU. Ember está implementado en MLIR, lo que permite una integración natural con PyTorch y TensorFlow. Al automatizar esta transformación, Ember desbloquea todo el potencial de rendimiento de la TMU sin imponer una carga adicional de programabilidad.