使用C语言编写extension为Kynetix大幅提速

使用C语言和swig为Kynetix统计部分代码写了C扩展，提升计算效率20多倍！

之前Kynetix代码几乎全部是使用Python写的，使用了KMCLib(Python + C++)之后，在进行主体KMC循环的时候效率得到了明显的提升，但是在刚刚完成on-the-fly分析部分代码以后开始跑程序，明显感觉程序异常的慢，因为kmc写到现在真心感受到kMC程序真心是个compute intense的活，其中没有用到矩阵运算，没有用到什么传统的优化算法，

所以尝试只用一种解释性语言如matlab和python去完成整个kMC程序的最终肯定会使程序慢得一塌糊涂！跑蒙特卡洛程序慢成狗，那这程序也几乎没啥用了感觉。

于是我也想学习KMCLib那样使用静态语言为我的python程序编写扩展模块，虽然我也知道针对Fortran有f2py这种工具存在，但是由于我个人的喜好，不是历史原因必须使用fortran的话，我还是选择C或者C++。
至于python的C扩展，我选择使用SWIG。

需要使用C重写的几个核心统计函数

在kmc_functions.py中有三个函数用于在KMC主循环运行当中对表面的configuration进行on-the-fly统计:

  • collect_coverge(), 扫描整个网格，统计每个可能的中间物种的覆盖度。
  • match_elements(), 扫描整个网格，针对某一特定的局部configuration进行匹配并统计，其目的就是为了统计在当前网格中有多少局部的构型能够发生某一特定的反应。
  • match_elements_list(), 针对一系列的局部构型列表，循环调用match_elements()，统计能够匹配列表中所有构型的总数。

这里由于针对NxN的网格，匹配单独一个local configuration的复杂度在$\mathcal{O}(N^{2})$， 匹配一列表的复杂度同扫面一次网格遍历所有构型的复杂度是相同的, 都是$\mathcal{O}(MN^{2})$

编写C扩展

单纯使用C语言写功能相同的函数还是比较简单的，最麻烦的是将其封装成python的接口，并实现C语言的数据类型同python内置的数据类型以及numpy的数组等数据类型之间进行转换，使得C语言编写的函数能够让python调用（下一部分进行简单介绍）。
C扩展的源代码在pynetics/extension/src/kmc_function.c中，我针对kmc_functions.py中的函数编写了相应的C语言版本，参数类型有所区别，主要是针对矩阵和C数组形状的参数传递，此问题在编写swig接口文件中进行了统一解决。
代码就不贴上来了，有兴趣的可以直接到项目的GitHub主页去看，传送 $\rightarrow$ kmc_functions.c

SWIG接口文件

这一部分是写C扩展让我最费劲的地方，因为涉及到数据类型的转换(自定义的类型映射customized typemap)，例如python的string list和C语言的char * 数组(或者char **)的转换等。为了这个看了不少stackoverflow和SWIG的文档，最后还是在官方文档里面找到了解决办法。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
/* tell SWIG to treat char ** as a list of strings */
%typemap(in) char ** {
    // check if is a list
    if(PyList_Check($input))
    {
        int size = PyList_Size($input);
        int i = 0;
        $1 = (char **)malloc((size + 1)*sizeof(char *));
        for(i = 0; i < size; ++i)
        {
            PyObject * o = PyList_GetItem($input, i);
            if(PyString_Check(o))
                $1[i] = PyString_AsString(o);
            else
            {
                PyErr_SetString(PyExc_TypeError, "list must contain strings");
                free($1);
                return NULL;
            }
        }
    }
    else
    {
        PyErr_SetString(PyExc_TypeError, "not a list");
        return NULL;
    }
}

大致原理就是根据输入的类型判断是否是python的list类型，并动态分配内存，生成相应的char **数组，这样完成了list -> char **的映射。

我在pynetics/extensions/wrap/中包含了numpy.i，此文件是numpy的swig接口文件，主要是利用里面关于numpy array的类型映射，来映射numpy的N维数组到C语言的N维数组：

1
2
3
int match_elements(char ** types, char ** elements,
                   int DIM1, int DIM2, double * IN_ARRAY2,
                   const int grid_shape[2]);

另外关于python中的string二维list，至今我是没找到如何映射到C语言相应的数组，囧（如果看到这篇博文的大神知道映射方法，请务必告诉我，小的感激不尽。
于是我想到了另外一种方法：
就是在python代码中将二维string数组转换成相应的一维数组将其降维，类似numpy中的reshape函数的功能，然后再用上文得到的string一维数组的方法将其映射。

另外我还编写了一份Makefile来编译并使用SWIG封装扩展，内容如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
CC = gcc
SWIG = swig
SWIGFLAGS = -python 
CCFLAGS = -O3 -fPIC -c 
OMPFLAG = # -fopenmp
INC = # -I /usr/local/lib/python2.7/site-packages/numpy/core/include 
LDFLAGS = -shared
EXTNAME = kmc_functions
EXTDIR = ../pynetics/solvers/plugin_backends

_kmc_functions.so: build/kmc_functions.o build/kmc_functions_wrap.o
    $(CC) $(OMPFLAG) $(LDFLAGS) $^ -o $@

build/kmc_functions.o: src/kmc_functions.c
    $(CC) $(OMPFLAG) $(CCFLAGS) $(INC) -I include $< -o $@

build/kmc_functions_wrap.o: wrap/kmc_functions_wrap.c
    $(CC) $(OMPFLAG) $(CCFLAGS) $(INC) -I include $< -o $@

wrap/kmc_functions_wrap.c: wrap/kmc_functions.i
    $(SWIG) $(SWIGFLAGS) $^
    mv wrap/*py ./

install:
    cp _$(EXTNAME).so $(EXTNAME).py $(EXTDIR)

clean:
    rm -rf build/* wrap/$(EXTNAME)_wrap.c $(EXTNAME).py *.so

uninstall:
    rm -f $(EXTDIR)/$(EXTNAME).py $(EXTDIR)/_$(EXTNAME).so

通过调用封装好的C扩展模块进行相同的计算，我在我的CentOS虚拟机中相同的环境下跑相同的输入参数，比纯python快了20倍！

如果算上使用KMCLib的C++提速，保守估计，比纯解释性语言写的KMC程序要快50~60倍！

Comments

▲