Python é conhecido por sua simplicidade e versatilidade. Mas essas qualidades cobram um preço alto em termos de desempenho. Felizmente, existem várias estratégias para acelerar códigos Python. A vetorização é certamente uma das principais. Também existem truques simples envolvendo operações com strings que podem garantir uma melhora na performance dos códigos. Outra opção é o uso de compiladores Just-in-Time (JIT).
JIT é um conceito fascinante que pode melhorar significativamente o desempenho de algumas linguagens de programação, incluindo Python. Porém, eles são pouco conhecidos.
Por que Códigos Python são Lentos?
Para entender o conceito de JIT, primeiro é importante compreender algumas diferenças entre linguagens de programação compiladas e interpretadas.
Linguagens de programação compiladas como C traduzem todo o código-fonte em código de máquina (instruções específicas que a CPU entende) antes de executar um programa. Isso torna o programa significativamente mais rápido. Porém, essa abordagem requer uma etapa de compilação separada e limita consideravelmente a flexibilidade dos programas.
Em comparação, linguagens de programação interpretadas como Python executam diretamente o código-fonte linha por linha sem criar código de máquina. Ou seja, sem pré-compilação. A vantagem dessa abordagem é que ela permite desenvolvimentos rápidos, flexíveis e interativos. Mas isso pode causar desempenhos mais lentos em comparação com as linguagens compiladas.
Compiladores JIT
Ao contrário dos compiladores tradicionais, um compilador JIT compila o código apenas quando é necessário. Isso reduz o tempo de inicialização dos programas. Em vez de compilar todo o programa, um compilador JIT se concentra em partes usadas com frequência, onde as melhorias de velocidade são mais impactantes. À medida que o programa é executado, o compilador JIT analisa e otimiza o código, adaptando-o a padrões específicos de hardware e execução. Com isso, compiladores JIT podem proporcionar ganhos significativos de desempenho ao longo do tempo.
Benefícios dos Compiladores JIT para Códigos Python
O Python não é conhecido por sua velocidade, mas compiladores JIT podem melhorar drasticamente o desempenho de seus códigos. Um compilador JIT pode ser particularmente importante para acelerar loops e funções usadas com frequência.
Compiladores JIT injetam uma dose de velocidade nos códigos sem comprometer a sintaxe amigável do Python. Isso é uma grande vantagem em comparação, por exemplo, com o uso de vetorizações que aceleram códigos, mas requerem modificações extensas em suas sintaxes.
Outra vantagem dos compiladores JIT é a redução do uso de memória. Um compilador JIT ajuda a evitar a compilação de seções raramente usadas, mantendo o espaço ocupado pela memória otimizado.
Compiladores JIT também apresentam otimização dinâmica. Eles se adaptam facilmente a padrões específicos de hardware e permitem aumentos de desempenho personalizados.
Alguns Compiladores JIT para Python
Nem todos os compiladores JIT são iguais. Alguns são adaptados para casos específicos, enquanto outros oferecem uma abordagem mais ampla. Portanto, para usá-los é importante conhecer algumas opções de compiladores JIT, e definir aquele que se alinha com as necessidades do seu projeto.
O PyPy é uma implementação alternativa do Python com seu próprio compilador JIT. Para inúmeros projetos, ele proporciona melhorias significativas de velocidade. Além disso, ele funciona perfeitamente com a maioria dos códigos Python sem precisar de modificações.
Jython é uma implementação de Python para a Java Virtual Machine (JVM). Como o nome sugere, Jython combina a elegância e simplicidade do Python com a robustez e versatilidade do Java. Ele utiliza o compilador JIT da JVM.
Numba é um compilador JIT para Python focado em computação numérica e científica. Ele permite que os usuários criem funções numéricas rápidas e eficientes usando a sintaxe Python. Numba é particularmente útil para códigos que envolvem cálculos de matriz e arrays.
JAX é uma biblioteca de alta-performance da Google para computação numérica e científica em CPU, GPU e TPU. Ela conta com uma versão atualizada do Autograd para diferenciar automaticamente códigos nativos de Python e NumPy. Portanto, JAX é excelente para implementações de algoritmos de machine learning, especialmente redes neurais. Embora JAX não seja um compilador JIT para código Python em geral, ela incorpora a compilação JIT para certas operações numéricas. Ela usa XLA (Accelerated Linear Algebra), um compilador para operações de álgebra linear. JAX aproveita a compilação JIT para operações específicas para alcançar alto desempenho em várias arquiteturas de hardware.
JIT na Prática
Para ilustrar um pouco o poder de aceleração de um compilador JIT, faremos uma comparação entre uma operação vetorizada usando NumPy e uma versão vetorizada usando o compilador JIT da biblioteca JAX. Lembre-se que a vetorização com NumPy já oferece um ganho de velocidade considerável em comparação com loops tradicionais. Veja o trecho abaixo:
import numpy as np
import time
import jax
import jax.numpy as jnp
from jax import jit
def f(x):
return -4*x*x*x + 9*x*x + 6*x - 3
x = np.random.randn(10000, 10000)
init_numpy = time.time()
f(x)
fim_numpy = time.time()
print('tempo de execução NumPy: ' + str(round(fim_numpy - init_numpy, 5)))
init_jax = time.time()
jax.jit(f)(jnp.array(x)) # JAX + JIT
fim_jax = time.time()
print('tempo de execução JAX: ' + str(round(fim_jax - init_jax, 5)))
Nesse código, aplicamos uma equação (definida nas linhas 7-8) a uma matriz de 10000 por 10000 elementos (linha 10). Primeiro resolvemos essa equação usando apenas NumPy (linha 12), que já é uma estratégia de aceleração para Python. Depois, usamos o compilador JIT da biblioteca JAX (linha 17). Note que para as duas situações estamos computando os tempos de execução. Eles são mostrados abaixo:
tempo de execução NumPy: 1.06333
tempo de execução JAX: 0.63822
Esses resultados variam muito de código para código e também com o hardware usado. Mas, sem dúvidas, JIT pode proporciar uma aceleração considerável em inúmeras situações.
Conclusões
Códigos Python são lentos, mas podem ser acelerados com várias estratégias. Neste post, o foco foi nos compiladores JIT. Eles são uma alternativa bem interessante para acelerar códigos Python sem precisar modificá-los. Se interessou? Busque se informar sobre o PyPy, JAX e outras opções e entenda qual é mais apropriada para seus projetos.
