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

Numba (1)

---

Numba 타입 컨테이너

Numba에서 가장 중요한 권장 사항은 최적의 컴파일을 제공하는 "nopython" 모드를 사용하는 것일 겁니다. (오브젝트 모드 (object mode)는 최소한의 최적화만 진행하게 됩니다) 이는 @jit 데코레이터에 "nopython=True" 옵션을 주거나 @njit 데코레이터를 사용하면 되는데요, 하지만 지난 포스트에서 살펴봤듯이 이 모드는 제한이 많아 컴파일이 성공하려면 Numba가 작상한 함수 내의 모든 변수 타입을 추론할 수 있어야 합니다. 

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

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

from numba import njit
from numba.typed import List

@njit
def foo(x):
    result = x.copy()
    result.append(11)
    return result

a = List()
for i in (2,3,5,7):
    a.append(i)
b = foo(a)

• a=List() 로부터 Numba typed-list 를 만들고 일반 파이썬 코드 안에서 내용을 typed-list에 추가합니다. 이때 typed-list의 타입은 맨 처음 추가된 원소에 따라 추론됩니다.
• b=foo(a) 에서는 Numba로 컴파일한 foo 함수를 호출하고 11을 추가합니다.

 

Numba와 for 루프

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

import numpy as np

@njit
def numpy_func(a):
    return a.sum()

@njit
def for_loop(a):
    acc = 0
    for i in a:
        acc += i
    return acc

a = np.arange(1000000, dtype=np.int64)
>>> %timeit numpy_func(a)
473 µs ± 134 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
>>> %timeit for_loop(a)
411 µs ± 2.77 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
>>> %timeit numpy_func.py_func(a)
758 µs ± 51.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
>>> %timeit for_loop.py_func(a)
114 ms ± 4.16 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

• Numba 컴파일된 함수의 py_func 속성을 참조하면 원래 순수 파이썬 함수도 사용할 수 있습니다.

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

 

루프 퓨전 (loop fusion)

따라서 넘파이 배열식을 for 루프로 다시 작성하는 대신 Numba 컴파일을 사용해 성능 최적화를 얻을 수 있습니다. 바로 루프 퓨전 (loop fusion) 이라는 최적화죠.

@njit
def loop_fused(a, b):
    return a*b - 4.1*a > 2.5*b

a, b = np.arange(1e6), np.arange(1e6)

>>> %timeit loop_fused(a,b)
1.12 ms ± 297 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
>>> %timeit loop_fused.py_func(a,b)
6.09 ms ± 123 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

결과를 보면 Numba 컴파일 버전이 순수 넘파이 버전보다 5배 정도 빠릅니다. Numba가 없으면 배열식이 for 루프 여럿으로 바뀌고 메모리에서 임시 배열이 생성됩니다. 배열식의 산술 연산마다 임시 배열에 대해 for 루프를 수행하고 각 루프의 결과가 임시 배열에 저장됩니다. 하지만 Numba 루프 퓨전 최적화는 여러 루프에 걸쳐 있는 산술 연산은 하나의 루프로 합쳐주어 결과를 계산하는데 필요한 메모리 크기외 메모리 검색 횟수를 줄일 수 있습니다. 실제로 루프 퓨전이 일어난 버전은 다음과 비슷한 모습입니다.

@njit
def manual_loop_fused(a, b):
    N = len(a)
    result = np.empty(N, dtype=np.bool_)
    for i in range(N):
        a_i, b_i = a[i], b[i]
        result[i] = a_i * b_i - 4.1 * a_i > 2.5 * b_i
    return result

>>> %timeit manual_loop_fused(a,b)
868 µs ± 12.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)