Numba

파이썬은 생산성이 높은 언어지만, 느리다는 단점이 있습니다.

 

Python은 인터프리터 언어로서 C/C++/Fortran과 같은 컴파일 언어에 비해 느리지만 Python code를 LLVM 컴파일러를 이용해 머신 코드로 바꾸어 수치연산을 가속화해주는 Numba compiler가 존재합니다. 

 

Numba는 수치 계산에 초점을 맞춘 파이썬을 위한 오픈 소스 JIT(Just-In Time) compiler로 2012년 컨티넘 애널리틱스(현 Anaconda)에서 처음 만들었습니다.

 

Numba 라이브러리의 주요 특징으로는,

• 전체 애플리케이션, 프로그램이 아닌 @jit, @njit decorator로 장식된 함수에 대해서만 별도로 컴파일합니다. Python은 보통 인터프리터 CPython을 사용하는데, Numba는 별도의 인터프리터를 사용하지 않고 최적화된 별도의 빠른 함수를 만들어줍니다.
• 함수에는 입력(argument), 출력(return)이 존재하는데 Numba는 특정 수치형 데이터 타입에 (int, float, complex) 대해서 빠른 연산을 제공합니다. 따라서 numpy array와 일반적으로 같이 사용합니다. 
• 함수가 처음 호출되었을 때 컴파일합니다. (Just In Time, JIT) 그 이후로 타입이 같은 인자를 전달해 함수를 실행하면 컴파일된 버전을 사용해 원래 버전보다 빠르게 작동합니다.
• Numba는 Python뿐만 아니라 Numpy를 인식해서 Numpy를 사용하는 코드도 컴파일합니다. 하지만 완전한 compiler는 아니다보니 Python과 Numpy의 일부분만 컴파일할 수 있습니다.
• Numba는 내부적으로 LLVM 프로젝트에 의존합니다. LLVM 프로젝트는 재사용할 수 있고, 모듈화된 compiler와 툴체인 기술의 모음입니다.

Numba 라이브러리의 동작 과정은 그림 1과 같습니다. Numba 라이브러리의 jit 모듈을 가속화하고자 하는 함수에 decorating 하면 함수의 argument의 타입을 파악하고 (type inference), LLVM compiler의 입력으로 변환해주는 과정을 (lowering) 거쳐 머신 코드로 컴파일하게 됩니다.

 

그림 1: Numba 작동 원리

• IR Intermediate Representations
• Bytecode Analysis Intermediate code more abstract than machine code
• LLVM Low Level Virtual Machine, infrastructure to develop compilers
• NVVM It is an IR compiler based on LLVM, it is designed to represent GPU kernels.

 

 

jit, njit

Numba 라이브러리에서는 jit 혹은 njit을 이용하여 함수를 데코레이팅할 수 있습니다. 가장 기본이 되는 decorator는 jit으로 njit은 jit(nopython=True)와 같습니다. Numba에서는 object mode / nopython mode가 존재하는데 nopython mode는 CPython interpreter를 사용하지 않겠다는 것으로 가장 최적의 컴파일 코드를 제공합니다. 하지만 이것은 int, float, complex와 같은 지정된 타입에 대해서만 동작합니다. 예를 들어 다음 예와 같이 딕셔너리 형태에 대해서는 nopython mode가 동작하지 않습니다.

@jit(nopython=True, cache=True)
def factorial_with_numba(n):
    res = 1
    for i in range(1, n+1):
        res *= i
    return res

jit에 nopython argument를 지정하지 않으면 디폴트로 object mode로 컴파일됩니다. 이것은 모든 형태의 파이썬 객체를 컴파일할 수 있는 모드이지만 수치형 타입이 아닌 데이터에 대해서 최적의 속도를 보장하지 않습니다. 다만, nopython을 지정하지 않더라도 loop-jitting이라는 컴파일 특성에 의해 수치형 타입에 대해서는 자동으로 nopython mode로 컴파일됩니다. 

 

또한 jit을 이용하면 다음과 같이 함수의 argument, 출력의 타입을 미리 지정할 수 있어 그림 1의 inference 단계를 건너뛸 수 있습니다. 미리 지정하지 않는 경우에는 첫번쨰 함수를 실행할 때까지 컴파일이 연기됩니다.

 

Numba가 가진 장점들은 확실한 것 같다.

• Decorator만 붙이면 된다!
• C, C++ 컴파일러를 다운 받을 필요가 없다.
• Numpy 코드 그대로 사용하면 된다.

물론 명확한 단점도 있는데, Numba는 Numpy 함수만 이해한다는 것이다.

 

Numba 타입 컨테이너

Numba에서 가장 중요한 권장 사항은 최적의 컴파일을 제공하는 nopython 모드를 사용하는 것입니다. 

Object mode는 최소한의 최적화만 진행하게 됩니다. 이는 @jit decorator에 nopython=True 옵션을 주거나 @njit decorator를 사용하면 됩니다. 하지만 지난 포스트에서 살펴봤듯이 이 mode는 제한이 많아 컴파일이 성공하려면 Numba가 작성한 함수 내의 모든 변수 타입을 추론할 수 있어야 합니다.

 

Numba는 튜플, 문자열, 정수, 실수 등의 간단한 스칼라 타입과 Numpy array 데이터 타입을 지원하지만 대표적으로 리스트, 딕셔너리 argument에 대해서는 동작하지 않습니다. 이는 Python 리스트나 딕셔너리는 서로 다른 타입의 원소를 저장할 수 있는데, Numba가 컴파일하려면 컨테이너의 원소가 모두 같은 타입이어야해서 문제가 생기는 것입니다. 이럴때 사용할 수 있는 것은 Numba의 타입 컨테이너입니다.

 

Numba에서는 typed-list, typed-dict 타입 컨테이너가 존재하는데, 파이썬 리스트, 사전 중에 균일한 타입의 version만을 뜻합니다. 즉, 이런 타입 컨테이너 내에서 오직 한 가지  타입의 원소만 들어갈 수 있습니다. 이런 제약을 제외하면 타입 컨테이너는 파이썬의 리스트나 사전과 비슷하고 사용할 수 있는 API도 비슷합니다. 또한 타입 컨테이너를 일반 파이썬 코드나 Numba로 컴파일한 함수 내에서 사용할 수 있고, Numba로 컴파일한 함수에 인자로 넘기거나, Numba로 컴파일한 함수에서 반환할 수도 있습니다.

 

Numba와 for loop

Numba는 for loop의 제한을 받지 않습니다. 즉, 컴파일한 Numba 함수 내에서 for loop를 사용해도 아무 문제가 없습니다. 다음의 두 함수가 있습니다. numpy_func 함수는 Numpy의 sum을 Numba로 컴파일한 것을 사용하고 for_loop 함수는 합을 계산하기 위해 for loop를 사용합니다.

 

결과를 보면 Numpy 배열을 Numba 컴파일한 numpy_func는 컴파일하지 않은 것과 2배 정도 크게 차이가 나지는 않습니다. 반면 순수 Python for loop 구현은 컴파일한 쪽보다 매우 느립니다.

 

loop fusion

따라서 Numpy 배열식을 for loop로 다시 작성하는 대신 Numba 컴파일을 사용해 성능 최적화를 얻을 수 있습니다. 바로 loop fusion이라는 최적화입니다.

 

 

A ~5 minute guide to Numba — Numba 0.50.1 documentation

 

Numba: “weapon of mass optimization”

 

Numba (2) - 타입 컨테이너, 루프 퓨전