포트란만큼 빠른 C++ 벡터 라이브러리를 파이썬에서 사용하기

GMES의 기하 관련 모듈은 원래 파이썬으로 구현되어 있었는데, 커다란 3차원 구조를 다루기에는 속도가 너무 느렸다. 속도 개선을 위해서 기하 관련 부분을 C++로 포팅하기로 결심했다. 그리곤, 기하 루틴에 많이 들어 있는 벡터 관련 함수를 어떻게 구현해야 할지 고민에 빠졌다. 파이썬에선 벡터의 자료형으로 numpy의 ndarray를 쓰면 되지만, C++에선 어떻게 구현하는 게 가장 좋을까? 시스템 호환성과 관리의 편리성을 고려해 널리 쓰이는 방법 중에서 선택할 수 있는 것들은 아래와 같다. 사용할 벡터는 실수를 원소로 갖는 3차원 벡터로 한정한다.

1. double *v, int size
2. std::array<double, 3>
3. boost::numeric::ublas::bounded_array<double, 3>
4. blitz::TinyVector<double, 3>

1번은 가장 쉽고 가장 빠른 방법이다. SWIG를 쓴다면, numpy.i를 이용해서 numpy와 연동하기도 쉽다. 하지만 벡터 연산을 일일이 직접 구현해야 하고, 함수의 인자 부분이 못생긴 형태가 되며, 왠지 C++답지 않은 방법이다.

2번은 표준 템플릿 라이브러리에 포함된 방법이며, 훨씬 C++스럽다. 하지만 1번처럼 벡터 관련 연산을 일일이 구현해야 하며 1번보다 느리다는 문제가 있다.

3번은 boost 라이브러리에 들어있는 컨테이너로 사실상 표준이라 할 수 있으며, 벡터 관련 연산도 모두 구현되어 있다. 하지만 속도가 조금 느리다는 단점이 있다. 간단한 벡터 연산뿐만 아니라 다양한 수학연산을 사용해야 한다면 boost 라이브러리를 선택하는 것이 좋다.

4번은 blitz++ 라이브러리에 들어있는 컨테이너로 포트란에 견줄 만한 속도가 장점이다. 과학계산 분야에선 거의 표준이라 할 수 있다.

보는 바와 같이 3번이나 4번이 좋은 선택이다. GMES 구현에는 4번을 선택했는데, 그 이유는 GMES가 SWIG를 이용해서 파이썬과 연동을 하고 있기 때문이다. boost는 자체적인 파이썬 연동 방식을 제공하고 있어서 한 프로그램에서 SWIG와 boost가 동시에 사용하는 것은 뭔가 자연스럽지 못한 느낌이 들었다. 그뿐만 아니라, blitz++가 아직은 더 빠르다.

blitz::TinyVector<double, 3>을 입력이나 출력으로 사용하는 함수를 파이썬에 연결하려면 어떻게 해야 할까? 우선 아래와 같은 간단한 C++ 함수를 생각해보자.

// util.hh

#include <blitz/tinyvec.h>

namespace util {
  typedef blitz::TinyVector<double, 3> Vector3;
  
  double accumulate(const Vector3& v);
  void weight(Vector3& v, double s = 1);
}

util::accumulate는 와 벡터 성분의 합을 구하는 함수이고, util::weight는 벡터에 상수를 곱하는 함수이다. 이들 함수는 아래와 같이 별로 어렵지 않게 구현해볼 수 있다.

// util.cc

#include "util.hh"

using namespace util;

double
util::accumulate(const Vector3& v)
{
  return v[0] + v[1] + v[2];
}

void 
util::weight(Vector3& v, double s)
{
  v *= s;
}

util::accumulate는 상수 인자를 받으니, 파이썬 측의 시퀀스(tuple, list, numpy.ndarray 등)를 C++ 쪽으로 변환해 주기만 하면 된다. 반면, util::weight의 인자는 값이 변경되므로 util::Vector 형을 다시 파이썬의 시퀀스로 변경해 줘야 한다. 일단 파이썬 시퀀스 형태만 만족하면 tuple, list형은 문제가 되지 않는다. numpy.ndarray 형도 numpy가 알아서 처리한다. 우선 blitz::TinyVector<>를 고려하지 않고 기본적인 인터페이스 파일을 만들어 보자.

// util.i

%module util

%{
#include "util.hh"
%}

%include "util.hh"

이렇게만 해도 컴파일이 되고, 파이썬 모듈을 읽어들일 수 있다. 하지만 함수를 실행할 수는 없다.

>>> import util
>>> a=[1,2,3]
>>> util.accumulate(a)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: in method 'accumulate', argument 1 of type 'util::Vector const &'

역시나 함수 인자의 형태가 맞지 않다는 에러를 내보낼 뿐이다. 즉, 파이썬의 시퀀스 형을 util::Vector 형으로 변환하는 방식을 알려줘야 한다. 문제는 SWIG가 blitz::TinyVector<>를 자동으로 처리해 주지 않는다는 점이다. 이건 %typemap을 써서 직접 변환 규칙을 만들어 줘야 한다.

%typemap(in) util::Vector3&
{
  $1 = convert_to_vector3($input);
  if (!$1) return NULL;
}

%typemap(argout) util::Vector3&
{
  for (int i = 0; i < 3; i++) {
    PyObject *o = PyFloat_FromDouble((*$1)[i]);
    PySequence_SetItem($input, i, o);
  }
}

%typemap(freearg) util::Vector3&
{
  if ($1) delete $1;
}

기본적인 역할은 %typemap(in)이 맡는다. 함수 인자의 util::Vector& 부분을 찾아 이 규칙을 적용한다. $1은 C++ 함수의 첫 번째 인자를 뜻하고, $input은 파이썬 측 함수의 인자를 가리킨다. convert_to_vector에 실제 변환 코드가 들어가는데, 이 코드는 아래에서 다룬다. util::weight 함수가 제대로 작동하려면 %typemap(argout)을 써야 한다. %typemap(argout)은 출력 역할을 맡은 인자를 처리할 수 있도록 해준다. 마지막으로 %typemap(freearg)는 형 변환 중에 힙에 생성한 자료 영역을 지울 수 있도록 해준다. 아래의 convert_to_vector를 보면 %typemap(freearg)를 사용한 이유를 알 수 있다.

%{
static util::Vector3*
convert_to_vector3(PyObject* input)
{
  if (!PySequence_Check(input)) {
    PyErr_SetString(PyExc_ValueError, "Expected a sequence");
      return NULL;
  }
  if (PySequence_Length(input) != 3) {
    PyErr_SetString(PyExc_ValueError, "Size mismatch. Expected 3 elements");
    return NULL;
  }

  util::Vector* v = new util::Vector3();
  for (int i = 0; i < 3; i++) {
    PyObject *o = PySequence_GetItem(input, i);
    if (PyNumber_Check(o))
      (*v)[i] = PyFloat_AsDouble(o);
    else {
      PyErr_SetString(PyExc_ValueError, "Sequence elements must be number");
      return NULL;
    }
  }
  return v;
}
%}

파이썬측의 입력이 시퀀스인지, 길이가 3인지를 확인한 후에 util::Vector를 임시로 생성하고 여기에 인자를 하나씩 복사하도록 했다. 복사하기 전에 각 원소가 숫자인지 확인하도록 했다.

아직 조금 더 남아있다. 형식 변환 규칙이 하나 더 필요한데, 바로 %typemap(typecheck)이다. util::weight 함수의 두 번째 인자가 기본값을 가지고 있는데, SWIG은 이럴 경우, util::weight(util::Vector3&)와 util::weight(util::Vector3&, double)을 각각 만들고, 동적 자료형 검사를 통해서 어느 함수를 사용할지 결정하도록 코드를 생성한다. 그러므로 다음과 같이 간단한 자료형 비교 코드도 만들어 줘야 한다.

%typemap(typecheck, precedence=SWIG_TYPECHECK_POINTER) util::Vector3&
{
  $1 = false;
  if (PySequence_Check($input) && PySequence_Length($input) == 3)
    $1 = true;
}

이제 distutils를 이용해서 컴파일하려면, 다음과 같은 setup.py를 만들고, python setup.py build_ext -i라고 입력하면 파이썬에서 불러올 수 있는 모듈을 만들어준다.

from distutils.core import setup, Extension

util_module = Extension(name = '_util',
                        sources = ['util.i', 'util.cc'],
                        depends = ['util.hh'],
                        swig_opts = ['-c++'])

setup(ext_modules = [util_module])

이제 파이썬 해석기를 다시 시작하고 모듈을 읽어오면(reload로는 안 된다.) 잘 작동하는 것을 확인할 수 있다.

>>> import numpy, util
>>> v1 = [1,2,3]
>>> v2 = numpy.array(v1, float)
>>> util.accumulate(v1)
6.0
>>> util.weight(v1, 2)
>>> v1
[2.0, 4.0, 6.0]
>>> util.weight(v2, 3)
>>> v2
array([ 3.,  6.,  9.])
>>>