Python使用ctypes模塊調用C/C++

最近在做圖卷積相關的實驗，里面涉及到圖采樣，該過程可以抽象為：從一個包含n個節點,m條邊的圖中根據一定規則采樣一個連通圖。由於實驗使用的是FB15k-237數據集，共包含14541個節點，272115條邊，每次采樣30000條邊，采樣一次需要8s，這對於深度學習實驗來說是難以接受的，會導致GPU長時間空閑。因此我開始嘗試使用C/C++優化代碼，雖然最后優化效果不行，但是也是對python調用C代碼的一次學習，因此在此紀錄一下。

Python原代碼

 def get_adj_and_degrees(num_nodes, triplets):
    """ Get adjacency list and degrees of the graph"""
    adj_list = [[] for _ in range(num_nodes)]
    for i, triplet in enumerate(triplets):
        adj_list[triplet[0]].append([i, triplet[2]])
        adj_list[triplet[2]].append([i, triplet[0]])

    degrees = np.array([len(a) for a in adj_list])
    adj_list = [np.array(a) for a in adj_list]
    return adj_list, degrees

這里以get_adj_and_degrees函數為例，我們使用C/C++優化該函數。該函數只是起演示作用，具體細節不重要。

C/C++實現代碼

我們在sampler.hpp中對該函數進行優化，該文件的定義如下：

#ifndef SAMPLER_H
#define SAMPLER_H

#include <vector>
#include "utility.hpp"

using namespace std;

// global graph data
int num_node = 0;
int num_edge = 0;
vector<int> degrees; // shape=[N]
vector<vector<vector<int>>> adj_list; // shape=[N, variable_size, 2]


void build_graph(int* src, int* rel, int* dst, int num_node_m, int num_edge_m) {
    num_node = num_node_m;
    num_edge = num_edge_m;

    // resize the vectors
    degrees.resize(num_node);
    adj_list.resize(num_node);

    for (int i = 0; i < num_edge; i++) {
        int s = src[i];
        int r = rel[i];
        int d = dst[i];

        vector<int> p = {i, d};
        vector<int> q = {i, s};
        adj_list[s].push_back(p);
        adj_list[d].push_back(q);
    }

    for (int i = 0; i < num_node; i++) {
        degrees[i] = adj_list[i].size();
    }
}

#endif

這里C/C++函數把結果作為全局變量進行存儲，是為了后一步使用。具體的函數細節也不在講述，因為我們的重點是如何用python調用。

生成so庫

ctypes只能調用C函數，因此我們需要把上述C++函數導出為C函數。因此我們在lib.cpp中做出如下定義：

#ifndef LIB_H
#define LIB_H

#include "sampler.hpp"

extern "C" {
    void build_graph_c(int* src, int* rel, int* dst, int num_node, int num_edge) {
        build_graph(src, rel, dst, num_node, num_edge);
    }
}

#endif

然后使用如下命令進行編譯，為了優化代碼，加上了O3，march=native選項：

g++ lib.cpp -fPIC -shared -o libsampler.so -O3 -march=native

Python調用C/C++函數

編譯完之后，在當前目錄下生成了libsampler.so庫，我們就可以編寫python代碼調用C/C++函數了，Python代碼如下：

import numpy as np
import time
from ctypes import cdll, POINTER, Array, cast
from ctypes import c_int


class CPPLib:
    """Class for operating CPP library

    Attributes:
        lib_path: (str) the path of a library, e.g. 'lib.so.6'
    """
    def __init__(self, lib_path):
        self.lib = cdll.LoadLibrary(lib_path)

        IntArray = IntArrayType()
        self.lib.build_graph_c.argtypes = (IntArray, IntArray, IntArray, c_int, c_int)
        self.lib.build_graph_c.restype = None

    def build_graph(self, src, rel, dst, num_node, num_edge):
        self.lib.build_graph_c(src, rel, dst, num_node, num_edge)

class IntArrayType:
    # Define a special type for the 'int *' argument
    def from_param(self, param):
        typename = type(param).__name__
        if hasattr(self, 'from_' + typename):
            return getattr(self, 'from_' + typename)(param)
        elif isinstance(param, Array):
            return param
        else:
            raise TypeError("Can't convert %s" % typename)

    # Cast from array.array objects
    def from_array(self, param):
        if param.typecode != 'i':
            raise TypeError('must be an array of doubles')
        ptr, _ = param.buffer_info()
        return cast(ptr, POINTER(c_int))

    # Cast from lists/tuples
    def from_list(self, param):
        val = ((c_int) * len(param))(*param)
        return val

    from_tuple = from_list

    # Cast from a numpy array
    def from_ndarray(self, param):
        return param.ctypes.data_as(POINTER(c_int))

總結

python使用ctypes庫調用C/C++函數本身不難，但是優化代碼確是一個深坑，尤其是優化Numpy等科學計算庫時。因為這些庫本身已經進行了大量優化，自己使用C++實現的話，很有可能就比優化前還更差。