Desarrolladores crean una versión minimalista del TPU de Google completamente open source

Un equipo de cuatro ingenieros ha desarrollado tiny-tpu, una implementación de código abierto inspirada en las unidades de procesamiento tensorial (TPU) V1 y V2 de Google. Este proyecto, disponible en GitHub, busca democratizar el conocimiento sobre el diseño de aceleradores de hardware para inteligencia artificial.

¿Qué es tiny-tpu?

tiny-tpu es una unidad de procesamiento tensorial minimalista que replica la arquitectura fundamental de los TPU de Google, tradicionalmente mantenida como código cerrado. El proyecto incluye todas las especificaciones técnicas, código fuente y documentación necesaria para comprender y construir aceleradores de hardware desde cero.

Arquitectura y componentes principales

Elemento de procesamiento (PE)

Cada PE ejecuta una operación de multiplicación y acumulación (MAC) por ciclo de reloj, multiplicando datos entrantes por pesos almacenados y sumando el resultado a sumas parciales.

Array sistólico

La arquitectura utiliza una cuadrícula de elementos de procesamiento que inicia con un tamaño de 2×2. Los datos de entrada fluyen horizontalmente, las sumas parciales se mueven verticalmente, y los pesos permanecen fijos durante la computación.

Unidad de procesamiento vectorial (VPU)

Realiza operaciones elemento por elemento después del array sistólico, incluyendo:

• Adición de bias
• Función de activación Leaky ReLU
• Pérdida MSE (Mean Squared Error)
• Derivada de Leaky ReLU

Buffer unificado (UB)

Memoria de doble puerto que almacena matrices de entrada, pesos, vectores de bias y valores de activación para retropropagación.

Unidad de control

Maneja un conjunto de instrucciones de 94 bits que controla directamente todos los subsistemas del TPU.

Conjunto de instrucciones detallado

El ISA (Instruction Set Architecture) de 94 bits incluye campos específicos para:

• Bits 0-4: Señales de control de 1 bit (activación del sistema, inicio de lectura UB, modo transpuesto)
• Bits 5-6: Tamaño de columnas de lectura UB
• Bits 7-14: Tamaño de filas de lectura UB (0-255)
• Bits 15-22: Dirección de lectura UB
• Bits 23-25: Selección de puntero UB
• Bits 26-41 y 42-57: Datos de escritura host (punto fijo de 16 bits)
• Bits 58-61: Ruta de datos VPU
• Bits 62-77: Factor de escala de tamaño de lote inverso
• Bits 78-93: Factor de fuga VPU para activación

Configuración y desarrollo

Requisitos del sistema

• Python: Entorno virtual con cocotb
• iverilog: Para simulación
• gtkwave: Para visualización de formas de onda (construido desde código fuente en macOS)

Estructura del proyecto

tiny-tpu/ ├── src/ # Módulos SystemVerilog ├── test/ # Archivos de prueba Python/SystemVerilog ├── waveforms/ # Archivos VCD para visualización ├── docs/ # Documentación └── jupyter/ # Notebooks explicativos

Impacto y objetivos futuros

El proyecto representa un avance significativo en la democratización del conocimiento sobre diseño de chips. Según los desarrolladores, ninguno tenía experiencia profesional previa en arquitectura de hardware antes de iniciar este ambicioso proyecto.

Próximos pasos

1. Desarrollo de un compilador para el conjunto de instrucciones
2. Escalamiento del TPU a dimensiones mayores (256×256 o 512×512)

Filosofía open source

El equipo detrás de tiny-tpu enfatiza que los detalles de la arquitectura TPU son tradicionalmente código cerrado, al igual que la mayoría del diseño de chips. Su objetivo es crear la guía definitiva para ingresar al mundo de los aceleradores de hardware, accesible para todos los niveles de experiencia técnica.

El proyecto está disponible en GitHub bajo el repositorio tiny-tpu-v2/tiny-tpu y ha recibido 97 estrellas y 3 forks desde su publicación, demostrando el creciente interés de la comunidad en hardware especializado para IA.

Esta iniciativa marca un hito importante en la educación sobre diseño de hardware, proporcionando recursos prácticos y teóricos que anteriormente solo estaban disponibles dentro de grandes corporaciones tecnológicas.