多种代码范式详解

C++侧给python写接口

mirage/python/mirage/kernel.py

#include <Python.h>        // CPython C API，用来写 Python 扩展模块
#include <cuda_runtime.h>  // CUDA runtime API，主要用 cudaStream_t
#include <vector>          // std::vector，原代码最好补上这个 include

// Python 层调用的入口函数：__mirage_launcher.launch(...)
static PyObject *launch(PyObject *self, PyObject *args) {
  // Python 传进来的 5 个对象
  PyObject *input_list, *output_list, *py_buffer, *py_stream, *py_profiler_buffer;

  // C++/CUDA 侧真正需要的裸指针
  void *buffer;
  std::vector<void const *> input_tensors; // 输入 tensor 指针，只读
  std::vector<void*> output_tensors;       // 输出 tensor 指针，可写
  void *profiler_buffer;

  // 解析 Python 参数：期望传入 5 个 Python object
  // Python 调用大概是：
  // launch(input_ptrs, output_ptrs, workspace_ptr, stream_ptr, profiler_ptr)
  if (!PyArg_ParseTuple(args, "OOOOO",
                        &input_list,
                        &output_list,
                        &py_buffer,
                        &py_stream,
                        &py_profiler_buffer)) {
    PyErr_SetString(PyExc_TypeError, "Invalid parameters");
    return NULL;
  }

  // 检查 input/output 是否是 Python list
  // 注意这里传的不是 tensor object，而是 tensor.data_ptr() 组成的 list
  if(!PyList_Check(input_list) || !PyList_Check(output_list)) {
    PyErr_SetString(PyExc_TypeError, "Both arg1 and arg2 must be lists.");
    return NULL;
  }

  Py_ssize_t input_size = PyList_Size(input_list);
  Py_ssize_t output_size = PyList_Size(output_list);

  // 遍历 input_list，把 Python int 地址转成 C++ void const*
  for(Py_ssize_t i = 0; i < input_size; i++) {
    PyObject *item = PyList_GetItem(input_list, i); // borrowed reference
    void* tensor = PyLong_AsVoidPtr(item);          // Python int -> void*

    if(!tensor) {
      PyErr_Format(PyExc_TypeError,
                   "Failed to convert item %d (input) to void pointer", i);
      return NULL;
    }

    input_tensors.push_back(tensor);
  }

  // 遍历 output_list，把 Python int 地址转成 C++ void*
  for(Py_ssize_t i = 0; i < output_size; i++) {
    PyObject *item = PyList_GetItem(output_list, i);
    void* tensor = PyLong_AsVoidPtr(item);

    if(!tensor) {
      PyErr_Format(PyExc_TypeError,
                   "Failed to convert item %d (output) to void pointer", i);
      return NULL;
    }

    output_tensors.push_back(tensor);
  }

  // workspace buffer 指针
  buffer = PyLong_AsVoidPtr(py_buffer);

  // profiler buffer 指针，用来写 profiling 信息
  profiler_buffer = PyLong_AsVoidPtr(py_profiler_buffer);

  // CUDA stream 指针，决定 kernel 在哪个 stream 上 launch
  cudaStream_t stream = (cudaStream_t)PyLong_AsVoidPtr(py_stream);

  // 真正执行 Mirage 生成的 graph/kernel
  // 前面的代码只是 Python -> C++/CUDA 的 thin wrapper
  execute_mugraph(input_tensors,
                  output_tensors,
                  buffer,
                  stream,
                  profiler_buffer);

  Py_RETURN_NONE;
}

// Python 模块中的函数表
// 这里把 C++ 函数 launch 注册成 Python 函数 "launch"
static PyMethodDef ModuleMethods[] = {
  {
    "launch",        // Python 侧函数名：__mirage_launcher.launch
    launch,          // 对应的 C/C++ 函数
    METH_VARARGS,    // 使用普通 tuple 参数传递
    "Entry point for all kernels with this signature"
  },
  {NULL, NULL, 0, NULL} // sentinel，表示函数表结束
};

// Python 扩展模块定义
// 描述模块名、文档、函数表等信息
static struct PyModuleDef ModuleDef = {
  PyModuleDef_HEAD_INIT,
  "__mirage_launcher", // Python 模块名：import __mirage_launcher
  NULL,                // 模块文档
  -1,                  // 模块状态大小；-1 表示全局状态
  ModuleMethods,       // 模块函数表
  nullptr,             // m_slots，多阶段初始化用，这里不用
  nullptr,             // m_traverse，GC 相关，这里不用
  nullptr,             // m_clear，GC 相关，这里不用
  nullptr,             // m_free，模块释放 hook，这里不用
};

// Python import 这个 .so 时会调用这个初始化函数
// 模块名是 __mirage_launcher，所以函数名必须是 PyInit___mirage_launcher
PyMODINIT_FUNC PyInit___mirage_launcher(void) {
  // 根据 ModuleDef 创建 Python module object
  PyObject *m = PyModule_Create(&ModuleDef);
  if(m == NULL) {
    return NULL;
  }

  // 这里理论上有点冗余：
  // 因为 PyModule_Create(&ModuleDef) 已经会根据 ModuleMethods 注册函数
  PyModule_AddFunctions(m, ModuleMethods);

  return m;
}

.pyx

.pyx 的开发范式可以概括为：Python 负责易用接口，Cython 负责跨语言 binding，C++ 负责真正逻辑。在这类代码里，Python 用户看到的是 DTensor、CyKNGraph 这种高级对象；但这些对象内部实际保存的是 CppDTensor*、CppKNGraph* 之类的 C++ 指针。每个 .pyx 方法通常都遵循同一个模式：先检查 Python 参数类型，把 list/tuple/str/dtype 等 Python 对象转换成 C++ 能接受的 vector/char*/enum/raw pointer，然后调用 C++ core 的方法，最后把 C++ 返回的指针重新包装成 Python 对象返回。因此 .pyx 层不是主要算法实现层，而是 marshalling + thin wrapper + API facade：它让 Python 能自然地调用底层 C++ 编译器/IR/search/transpiler，同时避免把核心系统逻辑写在 Python 里。

def fuse_tensors(self, list[DTensor] inputs, int fused_dim, int num_groups, str name):
    # Python list[DTensor] -> C++ vector<const CppDTensor*>
    # 注意：这里传的是 graph tensor 的 C++ 指针，不是 tensor 数据本身。
    cdef vector[const CppDTensor*] cinputs
    cinputs.resize(len(inputs))

    cdef DTensor t
    for i in range(len(inputs)):
        assert(type(inputs[i]) == DTensor)
        t = inputs[i]
        cinputs[i] = t.c_ptr   # 取出 Python wrapper 内部保存的 CppDTensor*

    # Python str -> C char*
    # 这个 char* 只在本次 C++ 调用期间安全；如果 C++ 长期保存 name，C++ 侧应自行拷贝。
    cdef char* cname = NULL
    if name is not None:
        py_byte_string = name.encode('UTF-8')
        cname = py_byte_string

    # 调用真正的 C++ graph API，在 C++ IR 中创建 fused tensor。
    cdef CppDTensor* ptr = self.p_kgraph.fuse_tensors(
        cinputs, fused_dim, num_groups, cname
    )

    # C++ 返回 CppDTensor*，这里重新包装成 Python 侧 DTensor 句柄。
    output = ctypes.cast(<unsigned long long>ptr, ctypes.c_void_p)
    return DTensor(output)

C和python转换

kernel.py 里的 Python/C API 只解决 运行时 launch 的裸指针传递：
torch.Tensor.data_ptr()
-> Python int
-> PyLong_AsVoidPtr()
-> void*
-> execute_mugraph(...)

core.pyx 解决的是 编译器控制面 的复杂类型转换：
Python tuple/list/str/dtype/object
-> C++ vector/string/enum/pointer
-> Mirage C++ IR / search / transpiler