python原始碼閱讀筆記

zwlwf發表於2024-11-10
Python原始碼筆記

SWIG的python介面使用demo

主要準備三個檔案，example.h, example.i, example.c,
example.i,

%module example
%{
#include "example.h"
%}
%include "example.h"

//exampe.h
int kernelProbe(long long x,long long y, int len);

//example.c
#include <stdint.h>
#include <stdio.h>

int kernelProbe(uint64_t x_gm, uint64_t y_gm, int len) {
  printf("real value in x tensor:\n");
  uint8_t* x = (uint8_t*)x_gm;
  uint8_t* y = (uint8_t*)y_gm;
  for(int i=0; i<len; i++) {
    printf("%d\n", x[i]);
  }
  return 0;
}

首先利用swig根據example.i生成example_wrap.c和example.py,

swig -python example.i

然後利用setup.py完成extension so的編譯，

#setup.py
from distutils.core import setup, Extension
example_module = Extension('_example',
                            sources=['example_wrap.c', 'example.c'], extra_compile_args=['-g'])
setup(name='example',
      version='0.1',
      author='Your Name',
      description='Example swig',
      ext_modules=[example_module],
      py_modules=["example"])

python setup.py build_ext --inplace

inplace表示替換原本的so

最後可以執行a.py來比較C中列印的tensor內容和python中print列印的tensor內容

#a.py
import example as ee
import torch
x = torch.randint(0, 10, [20,], dtype=torch.uint8)
y = torch.randint(0, 10, [20,], dtype=torch.uint8)
c = ee.kernelProbe(x.data_ptr(), y.data_ptr(), 10)
print(x)


python a.py

tensor.data_ptr()本質是個地址，介面中可以申明為long long, C裡強轉為不同的指標，介面中定義為int8_t* Py轉的時候有問題。

example_wrap.c中解讀

example_wrap.c主要做兩個事，一是將C定義的函式進行掛載到模組中，二是將執行時python傳遞來的各種PyObject*的引數parse為C對應的引數。

python掛載C函式用的PyMethodDef定義如下，

其中

在自定義的extension的C中使用示例如下，定義一個PyMethodDef的表，

這個C函式介面表被掛到PyModuleDef的結構體中，

這個結構體物件被傳給Python中的介面PyModule_Create(PyMethodDef*);註冊到python的主程序中，從而可以被呼叫
m = PyModule_Create(&SWIG_module);

然後在PyInit_soName函式中完成m = PyModule_Create(&SWIG_module);操作，並返回生產的m這個PyObject*。

相關文章