Então eu tenho esta função:
import numpy as np
import numba as nb
@nb.njit(cache=True, parallel=True, nogil=True)
def triangle_half_UR_LL(size: int, swap: bool = False) -> tuple[np.ndarray, np.ndarray]:
total = (size + 1) * size // 2
x_coords = np.full(total, 0, dtype=np.uint16)
y_coords = np.full(total, 0, dtype=np.uint16)
offset = 0
side = np.arange(size, dtype=np.uint16)
for i in nb.prange(size):
offset = i * size - (i - 1) * i // 2
end = offset + size - i
x_coords[offset:end] = i
y_coords[offset:end] = side[i:]
return (x_coords, y_coords) if not swap else (y_coords, x_coords)
O que ele faz não é importante, o importante é que ele é compilado com JIT e Numba e, portanto, muito rápido:
In [2]: triangle_half_UR_LL(10)
Out[2]:
(array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2,
2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5,
5, 6, 6, 6, 6, 7, 7, 7, 8, 8, 9], dtype=uint16),
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3, 4, 5, 6, 7, 8, 9, 2, 3, 4,
5, 6, 7, 8, 9, 3, 4, 5, 6, 7, 8, 9, 4, 5, 6, 7, 8, 9, 5, 6, 7, 8,
9, 6, 7, 8, 9, 7, 8, 9, 8, 9, 9], dtype=uint16))
In [3]: %timeit triangle_half_UR_LL(1000)
166 μs ± 489 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
In [4]: %timeit triangle_half_UR_LL(1000)
166 μs ± 270 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
In [5]: %timeit triangle_half_UR_LL(1000)
166 μs ± 506 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
Agora, se eu definir outra função e compilá-la JIT com o Numba, o desempenho da função rápida cai inexplicavelmente:
In [6]: @nb.njit(cache=True)
...: def dummy():
...: pass
In [7]: dummy()
In [8]: %timeit triangle_half_UR_LL(1000)
980 μs ± 20 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
In [9]: %timeit triangle_half_UR_LL(1000)
976 μs ± 9.9 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
In [10]: %timeit triangle_half_UR_LL(1000)
974 μs ± 3.11 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
Isso é real, já reproduzi esse problema com sucesso muitas vezes, sem falhas. Inicio uma nova sessão do interpretador, colo o código e ele roda rápido. Defino a função fictícia, chamo a função fictícia e a função rápida fica inexplicavelmente mais lenta.
Captura de tela como prova:
Estou usando o Windows 11 e não tenho a mínima ideia do que está acontecendo.
Existe alguma explicação para isso? E como posso evitar esse problema?
Curiosamente, se eu me livrar do nogil
parâmetro e sem alterar mais nada, o problema desaparece magicamente:
In [1]: import numpy as np
...: import numba as nb
...:
...:
...: @nb.njit(cache=True, parallel=True)
...: def triangle_half_UR_LL(size: int, swap: bool = False) -> tuple[np.ndarray, np.ndarray]:
...: total = (size + 1) * size // 2
...: x_coords = np.full(total, 0, dtype=np.uint16)
...: y_coords = np.full(total, 0, dtype=np.uint16)
...: offset = 0
...: side = np.arange(size, dtype=np.uint16)
...: for i in nb.prange(size):
...: offset = i * size - (i - 1) * i // 2
...: end = offset + size - i
...: x_coords[offset:end] = i
...: y_coords[offset:end] = side[i:]
...:
...: return (x_coords, y_coords) if not swap else (y_coords, x_coords)
In [2]: %timeit triangle_half_UR_LL(1000)
186 μs ± 47.9 μs per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [3]: %timeit triangle_half_UR_LL(1000)
167 μs ± 1.61 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
In [4]: %timeit triangle_half_UR_LL(1000)
166 μs ± 109 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
In [5]: @nb.njit(cache=True)
...: def dummy():
...: pass
In [6]: dummy()
In [7]: %timeit triangle_half_UR_LL(1000)
167 μs ± 308 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
In [8]: %timeit triangle_half_UR_LL(1000)
166 μs ± 312 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
In [9]: %timeit triangle_half_UR_LL(1000)
167 μs ± 624 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
Por que isso acontece?
Mas não, se eu definir outras funções, de alguma forma a primeira função fica lenta novamente. A maneira mais simples de reproduzir o problema é simplesmente redefinindo-a:
In [7]: dummy()
In [8]: %timeit triangle_half_UR_LL(1000)
168 μs ± 750 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
In [9]: import numpy as np
In [10]: %timeit triangle_half_UR_LL(1000)
167 μs ± 958 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
In [11]: import numba as nb
In [12]: %timeit triangle_half_UR_LL(1000)
167 μs ± 311 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
In [13]: @nb.njit(cache=True, parallel=True)
...: def triangle_half_UR_LL(size: int, swap: bool = False) -> tuple[np.ndarray, np.ndarray]:
...: total = (size + 1) * size // 2
...: x_coords = np.full(total, 0, dtype=np.uint16)
...: y_coords = np.full(total, 0, dtype=np.uint16)
...: offset = 0
...: side = np.arange(size, dtype=np.uint16)
...: for i in nb.prange(size):
...: offset = i * size - (i - 1) * i // 2
...: end = offset + size - i
...: x_coords[offset:end] = i
...: y_coords[offset:end] = side[i:]
...:
...: return (x_coords, y_coords) if not swap else (y_coords, x_coords)
In [14]: %timeit triangle_half_UR_LL(1000)
1.01 ms ± 94.3 μs per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [15]: %timeit triangle_half_UR_LL(1000)
964 μs ± 2.02 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
A lentidão também acontece se eu definir a seguinte função e chamá-la:
@nb.njit(cache=True)
def Farey_sequence(n: int) -> np.ndarray:
a, b, c, d = 0, 1, 1, n
result = [(a, b)]
while 0 <= c <= n:
k = (n + b) // d
a, b, c, d = c, d, k * c - a, k * d - b
result.append((a, b))
return np.array(result, dtype=np.uint64)
In [6]: %timeit triangle_half_UR_LL(1000)
166 μs ± 296 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
In [7]: %timeit Farey_sequence(16)
The slowest run took 6.25 times longer than the fastest. This could mean that an intermediate result is being cached.
6.03 μs ± 5.72 μs per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [8]: %timeit Farey_sequence(16)
2.77 μs ± 50.8 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
In [9]: %timeit triangle_half_UR_LL(1000)
966 μs ± 6.48 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
TL;DR: Isso se deve principalmente ao alocador do sistema , que não se comporta da mesma maneira em relação ao estado atual da memória (difícil de prever). Quando a função é rápida, não há falhas de página , enquanto quando a função é lenta, parece haver muitas falhas de página, deixando a thread mestre lenta.
Análise
Quando um array Numpy é criado, ele solicita alguma memória para a biblioteca C padrão (também conhecida como libc), normalmente usando a
malloc
função . Para grandes alocações, a libc solicita algum espaço de memória para o sistema operacional. Para pequenas alocações, ela pode pré-alocar algum espaço de memória e reciclar para evitar solicitar memória para o sistema operacional. De fato, esse sistema operacional é bastante caro. Além disso, as páginas de memória virtual solicitadas pelo sistema operacional não são fisicamente mapeadas na RAM quando o sistema operacional fornece o espaço de memória. Isso é feito preguiçosamente: quando a memória é lida ou gravada (página por página). A propósito, é por isso quenp.zeros
enp.empty
pode ser mais rápido do quenp.full
(veja esta postagem para mais informações sobre isso). Este princípio é chamado de overcommit. Quando um programa lê/grava uma página que ainda não está mapeada na RAM, a MMU aciona uma interrupção para que o sistema operacional faça o mapeamento real da página virtual para uma página física na RAM. Isso é chamado de falha de página e é muito caro. De fato, não apenas as exceções são muito caras, mas também todos os sistemas operacionais modernos escrevem zeros nas páginas mapeadas durante uma falha de página por motivos de segurança (porque a página pode conter dados possivelmente confidenciais de outros processos, como senhas de navegadores).O alocador padrão da libc tende a não ser conservador. Isso significa que ele tende a devolver memória ao sistema operacional e solicitá-la novamente. Isso significa que outros processos podem usar essa memória disponível, mas torna os processos que frequentemente solicitam memória significativamente mais lentos. Considerando a quantidade de dados alocados até o momento e o estado geral das páginas em cache, o alocador pode ou não solicitar mais memória ao sistema operacional. Portanto, falhas de página podem ou não ocorrer.
Além disso, quando há espaço contíguo suficiente nos buffers de memória pré-alocados da libc, o alocador padrão pode ou não reciclar um buffer que já foi usado recentemente e armazenado no cache da CPU. Quando isso acontece, não há necessidade de buscar dados na DRAM lenta, pois eles já estão lá. Esse efeito geralmente ocorre em loops que criam matrizes temporárias e as gravam/lêem antes de (implicitamente) excluí-las.
Aqui está uma explicação prática:
Caso 1) Após a primeira chamada da função, os 2 arrays são armazenados no cache, assumindo que ele seja pequeno o suficiente para caber nele (o que é o caso aqui no meu PC). Assim que a função é executada, o array é liberado automaticamente, mas o alocador padrão não o libera de volta para o sistema operacional. Em vez disso, o espaço de memória é reciclado para a próxima chamada de função e esta nova pode ser muito mais rápida: sem falha de página e sem necessidade de buscar dados da DRAM! Isso pode acontecer muitas vezes, desde que o estado da memória alocada não mude repentinamente de forma que não haja espaço de memória disponível para os arrays de memória solicitados. Se este caso acontecer, devemos ver quase nenhuma leitura/escrita na DRAM nem nenhuma chamada de sistema relacionada a falhas de página (difícil de rastrear no Windows).
Caso 2) Ao final da primeira função, a memória é liberada de volta para o sistema operacional e, em seguida, a próxima chamada de função precisa pagar a enorme sobrecarga de falhas de página e também de buscas na DRAM. Na verdade, até onde sei, os dados deveriam ser gravados de volta na DRAM e buscados novamente (já que a chance de a página virtual/física ser a mesma é pequena e também devido a liberações de TLB durante o desmapeamento). Se isso acontecer, devemos ver chamadas de sistema significativas relacionadas a falhas e, certamente, também um número significativamente maior de leituras/gravações na DRAM.
Caso 3) A libc sempre recicla memória, mas os endereços dos arrays nem sempre são os mesmos. Nesse caso, falhas de página não ocorrem, mas a quantidade de espaço de memória de leitura/escrita pode ser muito maior do que o estritamente necessário. Esse espaço de memória pode ser tão grande que não cabe no cache. Nesse caso, ocorre a destruição do cache, causando leituras/escritas de DRAM dispendiosas. Se isso acontecer, não deveremos ver nenhuma chamada de função relacionada a falhas de página ou sobrecarga associada, mas ainda assim leituras/escritas de DRAM significativas.
Com base nisso, seríamos capazes de tornar o problema mais provável de acontecer (não garantido) se alocássemos muitos arrays de tamanho variável para alterar o estado do espaço de memória alocado.
Resultados experimentais
Na minha máquina (Windows 10), aqui está o que obtenho quando consigo reproduzir o efeito:
Podemos ver que a criação de um array temporário impacta o desempenho da função, como esperado. Compilar a função Numba aloca um pouco de memória (bastante, pelo que sei) e, portanto, pode produzir um impacto semelhante.
Quando a função é rápida, a taxa de transferência de DRAM é de 1 GiB/s (~30% são gravações), enquanto é de 6 GiB/s (50~60% são gravações) quando a função é significativamente mais lenta. Observe que tenho alguns processos em execução em segundo plano, consumindo em média cerca de ~0,1 GiB/s (bastante variável).
Quando a função é rápida, também posso ver que 20 a 35% do tempo é gasto no kernel, o que não é ótimo (mas, estranhamente, quase nenhum tempo de kernel na thread mestre). Isso normalmente se deve a falhas de página para esse tipo de código, mas também possivelmente a operações de E/S devido à
cache=True
opção (improvável, já que isso só deveria acontecer na thread mestre) ou trocas de contexto (vejo algumas delas ocorrendo, embora não pareça um problema grave), especialmente devido à execução paralela ou ao tempo de espera das threads de trabalho (certamente devido a algum desequilíbrio de carga). Este último parece ser um problema significativo, mas me parece normal para códigos com limitação de memória como o seu.Quando a função é lenta, a fração de tempo gasto no kernel é um pouco maior (20~40%) nos workers e muito maior no master thread (30~85%).
Infelizmente, não consigo criar o perfil do kernel do Windows (ao contrário do Linux) para rastrear as chamadas de função relacionadas a falhas de página... A pilha de chamadas relatada no profiler de baixo nível que usei (VTune) não parece fornecer tal informação. Posso afirmar com segurança que isso aparentemente não tem nada a ver com operações de E/S, visto que há pouquíssimas E/S durante a computação.
Experimentalmente, a maior diferença clara é o tempo de kernel gasto no thread principal (em leitura no gráfico de agendamento):
Portanto, acredito que o tempo de kernel gasto na thread principal é provavelmente o que aumenta significativamente o tempo de execução. Isso significa que a parte do código não está no
prange
loop, a menos que falhas de página possam ocorrer nela.O culpado são estas linhas:
Conclusão
Isso confirma a teoria explicada anteriormente:
Correção / Solução alternativa
No Linux, você pode facilmente usar outro alocador mais conservador (por exemplo, TC-malloc) ajustando a
LD_PRELOAD
variável. No Windows, isso é aparentemente mais complicado ( pode funcionar com TC-malloc)...Uma solução alternativa é simplesmente executar falhas de página o máximo possível em paralelo para reduzir a sobrecarga, criando um array Numpy com
np.empty
ounp.zeros
. Dito isso, falhas de página tendem a escalar mal no Windows (ao contrário do Linux, onde costumam ser significativamente melhores).A melhor solução é simplesmente pré-alocar e pré-processar matrizes de falhas no momento da inicialização e evitar matrizes temporárias o máximo possível em partes críticas do seu programa (mesmo em benchmarks para fins de reprodutibilidade), especialmente matrizes temporárias grandes. Em aplicações HPC, a pré-alocação é muito comum (também para evitar problemas de falta de memória após >10h de cálculos em >100 máquinas). Em programas nativos que se preocupam com desempenho, como jogos, aplicações HPC (e até mesmo navegadores), alocadores personalizados são frequentemente usados para evitar tais problemas (e a sobrecarga de alocações). Infelizmente, em Python e especialmente em Numpy, é muito comum criar muitos novos objetos e matrizes temporárias (e até mesmo promovê-los), embora isso seja bastante ineficiente (e torne o benchmark mais complicado), e, pelo que sei, não podemos escrever facilmente alocadores personalizados para Numpy/CPython. Às vezes é difícil pré-alocar/processar matrizes de falhas porque nem sempre sabemos o tamanho antes de alguns cálculos. Isso também costuma tornar o código Python mais complexo. O mesmo vale para operações Numpy no local (que podem ser usadas para evitar matrizes temporárias).