vm / app_cpppy /app_cpppy.cpp
bitterapricot's picture
add python ext.
bc08ced
Raw
History Blame Contribute Delete
1.95 kB
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <pybind11/stl.h>
namespace py = pybind11;
// 示例C++函数:向量加法
std::vector<double> vector_add(const std::vector<double>& a, const std::vector<double>& b) {
if (a.size() != b.size()) {
throw std::runtime_error("Vectors must have the same size");
}
std::vector<double> result(a.size());
for (size_t i = 0; i < a.size(); ++i) {
result[i] = a[i] + b[i];
}
return result;
}
// 高性能矩阵乘法(示例)
py::array_t<double> matrix_multiply(py::array_t<double> a, py::array_t<double> b) {
py::buffer_info a_buf = a.request();
py::buffer_info b_buf = b.request();
if (a_buf.ndim != 2 || b_buf.ndim != 2) {
throw std::runtime_error("Input must be 2D arrays");
}
if (a_buf.shape[1] != b_buf.shape[0]) {
throw std::runtime_error("Matrix dimensions do not match");
}
auto result = py::array_t<double>({a_buf.shape[0], b_buf.shape[1]});
py::buffer_info r_buf = result.request();
double* a_ptr = static_cast<double*>(a_buf.ptr);
double* b_ptr = static_cast<double*>(b_buf.ptr);
double* r_ptr = static_cast<double*>(r_buf.ptr);
size_t m = a_buf.shape[0];
size_t n = a_buf.shape[1];
size_t p = b_buf.shape[1];
// 简单矩阵乘法(可优化)
for (size_t i = 0; i < m; ++i) {
for (size_t j = 0; j < p; ++j) {
double sum = 0.0;
for (size_t k = 0; k < n; ++k) {
sum += a_ptr[i * n + k] * b_ptr[k * p + j];
}
r_ptr[i * p + j] = sum;
}
}
return result;
}
// 定义Python模块
PYBIND11_MODULE(app_cpppy, m) {
m.doc() = "C++ extension for FastAPI";
m.def("vector_add", &vector_add, "Add two vectors");
m.def("matrix_multiply", &matrix_multiply, "Multiply two matrices");
// 添加更多C++函数...
}