Simulando Um Array Sistólico Com SIMD (Mais ou Menos)

Decidi tentar entender como uma TPU funciona. Como não tenho nenhuma experiência com hardware, decidi tentar simular um array sistólico usando operações SIMD. O resultado não é uma simulação “real”, porque não tentei simular o hardware em si, mas a ideia por trás dele e o fluxo de dados (parcialmente).

Eu queria tentar fazer minha CPU realizar multiplicações de matrizes mais rápido usando uma estratégia que vagamente se assemelhava a como um array sistólico funciona.

O que é um array sistólico?

Simplificando, é um dispositivo que organiza várias unidades de processamento de dados de forma homogênea para realizar uma operação predefinida. O que vamos focar é multiplicação de matrizes, mas poderia fazer outras coisas como correlação, convolução, ou até ordenação de dados.

Um array sistólico é o bloco de construção primário de uma TPU (unidade de processamento de tensores), que hoje em dia se tornou muito mais interessante para mim por causa da inteligência artificial. IA pode ultimamente ser reduzida a uma tonelada de multiplicações de matrizes, e a CPU é terrível em realizar multiplicação de matrizes.

No caso de operações de multiplicação de matrizes, cada PE recebe dados de “cima” e da “esquerda”, então realiza uma multiplicação com os dois valores, armazena o resultado em um registrador acumulador, e passa os valores adiante. Os valores continuam em cascata para baixo e para a direita até que o último PE (na borda inferior direita) realize o último cálculo.

Aqui está um diagrama que demonstra o fluxo de dados

Para comparação, um array sistólico 256x256 é capaz de realizar até 65.536 operações por tick de clock, enquanto uma CPU (assumindo execução escalar single-core) pode fazer 1 operação por tick de clock. Esta comparação não é 100% justa, mas serve para ilustrar a discrepância nas capacidades de operação paralela.

Agora, quando esse tipo de operação precisa ser feita, uma GPU é usada (daí as ações da Nvidia 👀), mas imagine um cenário onde você está jogando um jogo e um dos inimigos no jogo é controlado por um modelo de IA mais sofisticado. A GPU pode já estar fazendo muito trabalho para processar os gráficos, a CPU está ocupada servindo dados para a GPU, rodando o SO, e, como já sabemos, CPUs não são conhecidas por serem ótimas com multiplicação de matrizes de qualquer forma. Esse é o cenário onde uma TPU é útil e é uma das principais razões pelas quais acho que hardware de edge computing se tornará mais relevante no futuro.

O que é SIMD

CPUs modernas geralmente contêm instruções especiais comumente chamadas SIMD que podem realizar uma operação com múltiplos valores ao mesmo tempo (no mesmo core), e é mais fácil entendê-las se você tentar pensar em “modo Assembly”. Um exemplo usando pseudo-código poderia ser assim (não sei Assembly):

A = [1, 2, 3, 4, 5, 6, 7, 8]
B = [1, 1, 3, 5, 8, 13, 21, 34]
C = [0; 8]

C[0] = A[0] * B[0]
C[1] = A[1] * B[1]
C[2] = A[2] * B[2]
C[3] = A[3] * B[3]
C[4] = A[4] * B[4]
C[5] = A[5] * B[5]
C[6] = A[6] * B[6]
C[7] = A[7] * B[7]

Cada operação de multiplicação é realizada uma após a outra.

Usando SIMD, um pseudo-código que pode ilustrar o que acontece poderia ser assim:

A = [1, 2, 3, 4, 5, 6, 7, 8]
B = [1, 1, 3, 5, 8, 13, 21, 34]
C = A * B

Infelizmente, não consigo pensar em uma forma melhor de representar isso, mas o ponto principal é que no primeiro exemplo, a CPU vai carregar os valores de A e B um por um em seus registradores, realizar a multiplicação, e armazenar o valor de volta. Enquanto no segundo caso (usando SIMD), todos os valores de A são carregados de uma vez em um “registrador muito grande”, o mesmo para B, a multiplicação é aplicada elemento a elemento, e o resultado é armazenado de volta na memória.

Isso começa a fazer diferença quando começamos a fazer multiplicação de matrizes. Por exemplo, multiplicar uma matriz 256x256 por outra matriz 256x256 (ignorando possíveis otimizações) exigiria quase 17 milhões de multiplicações individuais sozinhas. Mas também há adições, carregamento de valores da memória, e outras operações envolvidas. Pode ficar bem lento (e vai).

Ao longo deste artigo, estarei usando o termo SIMD (que é mais genérico) para me referir ao conjunto de instruções disponível na CPU do meu PC, que é AVX2. Ele pode carregar até 256 bits por lane, o que é equivalente a 8 números de ponto flutuante de 32 bits. Os exemplos de código mostrados provavelmente não funcionarão com outros tipos de conjuntos de instruções SIMD.

Implementando

Para a implementação decidi usar Rust.

O primeiro passo é implementar uma multiplicação de matrizes tradicional usando a CPU de forma escalar com execução single-core para ter algo para comparar. Mas antes disso precisamos de uma forma de criar matrizes, decidi criar uma função que recebe número de colunas e número de linhas como entrada e gera um vetor de vetores com valores aleatórios usando o crate rand.

use rand::Rng;

fn gen_rand_matrix(rows: usize, cols: usize) -> Vec<Vec<f32>> {
    let mut rng = rand::rng();
    let mut matrix = vec![vec![0.0f32; cols]; rows];
    for i in 0..rows {
        for j in 0..cols {
            matrix[i][j] = rng.random_range(0.0..1.0);
        }
    }
    matrix
}

Agora a função que multiplica duas matrizes e a função main para executar e coletar o tempo de execução

use std::time::Instant;
use rand::Rng;

fn gen_rand_matrix(rows: usize, cols: usize) -> Vec<Vec<f32>> {...}

fn scalar_matrix_mult(a: &[Vec<f32>], b: &[Vec<f32>]) -> Vec<Vec<f32>> {
    let rows = a.len();
    let inner = a[0].len();
    let cols = b[0].len();
    assert_eq!(
        inner,
        b.len(),
        "Matrix dimensions do not match for multiplication"
    );
    let mut result = vec![vec![0.0f32; cols]; rows];
    for i in 0..rows {
        for j in 0..cols {
            let mut sum = 0.0f32;
            for k in 0..inner {
                sum += a[i][k] * b[k][j];
            }
            result[i][j] = sum;
        }
    }
    result
}

fn main() {
    let a = gen_rand_matrix(256, 256);
    let b = gen_rand_matrix(256, 256);
    let start_classical = Instant::now();
    let _ = scalar_matrix_mult(&a, &b);
    let duration_classical = start_classical.elapsed().as_millis();
    println!("multiplicação escalar de matrizes: {duration_classical} ms");
}

Na minha máquina leva cerca de 350ms para multiplicar essas duas matrizes quadradas.

Agora, precisamos implementar isso para funcionar com SIMD, para fazer isso precisamos usar o build nightly do cargo e usar a feature portable_simd

#![feature(portable_simd)]
use std::simd::num::SimdFloat;
use std::simd::Simd;
use std::time::Instant;
use rand::Rng;

Também defini uma constante que representa a quantidade de valores float32 que cabem em cada lane SIMD const N: usize = 8; e um tipo para representar uma lane type Lane = Simd<f32, 8>;

Mas antes de implementar a função para a multiplicação de matrizes precisamos considerar o “layout” dos dados, já que estamos tentando simular um array sistólico, faz sentido organizar os dados de cada matriz para serem alimentados para a função de forma similar ao que aconteceria em um array sistólico, com colunas da primeira matriz vindo da “esquerda” e linhas da segunda matriz vindo de “cima”. As funções pack_rows e pack_cols recebem uma matriz (vetor de vetores) como entrada e retornam a mesma matriz mas organizada em Lanes (vetor de vetor de lanes). Não muito eficiente em memória, mas direto.

fn pack_rows(a: &[Vec<f32>]) -> Vec<Vec<Lane>> {
    a.iter()
        .map(|row| {
            row.chunks(N)
                .map(|chunk| {
                    let mut buf = [0.0; N];
                    buf[..chunk.len()].copy_from_slice(chunk);
                    Lane::from_array(buf)
                })
                .collect()
        })
        .collect()
}

fn pack_cols(b: &[Vec<f32>]) -> Vec<Vec<Lane>> {
    let rows = b.len();
    let cols = b[0].len();
    (0..cols)
        .map(|j| {
            let mut col = Vec::new();
            for chunk_start in (0..rows).step_by(N) {
                let mut buf = [0.0f32; N];
                for i in 0..N {
                    let idx = chunk_start + i;
                    if idx < rows {
                        buf[i] = b[idx][j];
                    }
                }
                col.push(Lane::from_array(buf));
            }
            col
        })
        .collect()
}

Com isso em vigor, agora podemos criar a função de multiplicação de matrizes que recebe uma matriz “empacotada” e usa SIMD para realizar as operações

fn simd_matrix_mult(a_rows: &[Vec<Lane>], b_cols: &[Vec<Lane>]) -> Vec<Vec<f32>> {
    let mut result = vec![vec![0.0f32; b_cols.len()]; a_rows.len()];

    for i in 0..a_rows.len() {
        for j in 0..b_cols.len() {
            let a_vecs = &a_rows[i];
            let b_vecs = &b_cols[j];
            let mut acc = 0.0f32;

            for idx in 0..a_vecs.len() {
                acc += (a_vecs[idx] * b_vecs[idx]).reduce_sum();
            }

            result[i][j] = acc;
        }
    }

    result
}

Também decidi criar uma função que compara duas matrizes para ver se são equivalentes, considerando um valor de tolerância porque estamos lidando com aritmética de ponto flutuante

fn matrices_equivalent(a: &[Vec<f32>], b: &[Vec<f32>], tol: f32) -> bool {
    if a.len() != b.len() || a[0].len() != b[0].len() {
        return false;
    }

    for i in 0..a.len() {
        for j in 0..a[0].len() {
            if (a[i][j] - b[i][j]).abs() > tol {
                println!("Diferença em ({}, {}): {} vs {}", i, j, a[i][j], b[i][j]);
                return false;
            }
        }
    }

    true
}

Também alterei a função main para executar ambas as funções, cronometrar e comparar os resultados

fn main() {
    let a = gen_rand_matrix(256, 256);
    let b = gen_rand_matrix(256, 256);

    let a_rows = pack_rows(&a);
    let b_cols = pack_cols(&b);

    let start_classical = Instant::now();
    let scalar_result = scalar_matrix_mult(&a, &b);
    let duration_classical = start_classical.elapsed().as_millis();
    println!("multiplicação escalar de matrizes: {duration_classical} ms");

    let start_simd = Instant::now();
    let simd_result = simd_matrix_mult(&a_rows, &b_cols);
    let duration_simd = start_simd.elapsed().as_millis();
    println!("multiplicação de matrizes SIMD: {duration_simd} ms");

    if matrices_equivalent(&scalar_result, &simd_result, 0.01) {
        println!("Resultados são equivalentes.");
    } else {
        println!("Resultados diferem significativamente.");
    }
}

Como reorganizar os dados leva tempo e isso não aconteceria em um hardware real, decidi não considerar o tempo que leva para empacotar a matriz ao comparar o tempo de execução, mas mostrarei separadamente.

Aqui está a saída

Empacotamento levou: 2 ms
multiplicação escalar de matrizes: 354 ms
multiplicação de matrizes SIMD: 74 ms
Resultados são equivalentes.

Os valores absolutos não significam muito, porque cada hardware vai produzir tempos de execução diferentes, mas a diferença entre os valores é substancial. Isso fica mais evidente se aumentarmos o tamanho da matriz, aqui está uma comparação indo de 256 até 2048.

Podemos observar um aumento no desempenho comparativo que é quase linear, e claro que isso eventualmente vai saturar por causa de gargalo de memória, mas mesmo assim, acho isso incrível.

Usando múltiplos cores

Como já fizemos tudo isso, por que não dar um passo adiante e usar todos os cores disponíveis, cada um usando SIMD para ver quais resultados obtemos?

Para fazer isso decidi usar o crate rayon e a ideia é simples, cada core de CPU (no meu caso 8 no total) vai usar SIMD para realizar múltiplas operações em paralelo e a matriz será dividida em chunks maiores de 8 para que cada chunk maior possa ser processado por um core individual. Como cada lane SIMD pode trabalhar com até 8 valores float32 de uma vez, cada chunk maior será de tamanho 64 (oito ao quadrado).

Aqui está a função:

fn parallel_simd_matrix_mult(
    a_rows: &[Vec<Lane>],
    b_cols: &[Vec<Lane>],
) -> Vec<Vec<f32>> {
    let m = a_rows.len();
    let n = b_cols.len();
    let tile_size = 64;

    let mut result = vec![vec![0.0f32; n]; m];

    result
        .par_chunks_mut(tile_size)
        .enumerate()
        .for_each(|(tile_row_idx, tile_chunk)| {
            let i0 = tile_row_idx * tile_size;
            let i_max = (i0 + tile_chunk.len()).min(m);

            for (ii, row) in tile_chunk.iter_mut().enumerate() {
                let i = i0 + ii;

                for j0 in (0..n).step_by(tile_size) {
                    let j_max = (j0 + tile_size).min(n);
                    for j in j0..j_max {
                        let a_vecs = &a_rows[i];
                        let b_vecs = &b_cols[j];

                        let mut acc = 0.0f32;
                        for (a_lane, b_lane) in a_vecs.iter().zip(b_vecs.iter()) {
                            acc += (*a_lane * *b_lane).reduce_sum();
                        }

                        row[j] = acc;
                    }
                }
            }
        });

    result
}

Para ser 100% honesto, essa é a parte onde eu me perdi e apenas usei o chatGPT para chegar na função final, nesse ponto eu estava apenas tentando obter números ainda mais impressionantes e não estava me importando tanto com como eu cheguei lá.

De qualquer forma, aqui está a comparação

Não comparei todos os três de uma vez porque fazer a última multiplicação levaria muito tempo rodando em CPU escalar

Conclusão

Para mim tudo isso foi uma ótima experiência de aprendizado, não acho que nenhuma dessas coisas seja “útil” no contexto de desenvolvimento de software profissional, mas me fez apreciar mais o trabalho que entra em resolver um “problema simples” como fazer multiplicação de matrizes ir mais rápido. Levei muito tempo para entender como as coisas estavam funcionando nesta “simulação” de software, imagine o que entra em criar o hardware real.

No final estou animado para ver quais outras tecnologias incríveis serão criadas, aprimoradas ou reutilizadas para suportar o uso amplo de modelos de IA.