为什么 numpy 的 cumsum 比手动 C++ 的循环快得多？

Question

我想了解为什么 numpy 的 cumsum 函数比手动 C++ 循环快得多。我正在使用 pybind 提供从 C++ 代码到 python.

的映射

数字在前：

[nav] In [24]: x = np.arange(100_000, dtype=np.float32)

[nav] In [25]: %timeit np.cumsum(x)
295 µs ± 34.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

[ins] In [26]: %timeit derived.cumsum(x)
9.26 ms ± 137 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

现在derived.cumsum(x)是怎么得到的？

template <typename Sequence,
          typename = std::enable_if_t<std::is_rvalue_reference_v<Sequence&&>>>
inline py::array_t<typename Sequence::value_type> as_pyarray(Sequence&& seq) {
  auto size = seq.size();
  auto data = seq.data();
  auto seq_ptr = std::make_unique<Sequence>(std::move(seq));
  auto capsule = py::capsule(seq_ptr.get(), [](void* p) {
    std::unique_ptr<Sequence>(reinterpret_cast<Sequence*>(p));
  });
  seq_ptr.release();
  return py::array(size, data, capsule);
}

std::vector<float> cumsum(const std::vector<float>& nums) {
  std::vector<float> result(nums.size());
  float sum = 0;
  for (size_t i = 0; i < nums.size(); ++i) {
    sum += nums[i];
    result[i] = sum;
  }
  return result;
}

PYBIND11_MODULE(derived, m) {
  m.def("cumsum", [](const std::vector<float>& nums) {
    auto result = cumsum(nums);
    return as_pyarray(std::move(result));
  });
}

as_pyarray 是从 this pybind 问题中获得的，它确保在 python 和 c++ 之间交互时的零拷贝。

请注意，C++ 代码是使用 -O3 -DNDEBUG -fPIC 编译的。

Answer 1

问题出在cumsum的参数上。 const std::vector<float>& 仍在制作相当昂贵的副本。直接以py::array_t<float>为参数解决了这个问题。

更多详细信息位于：https://github.com/pybind/pybind11/issues/1042

为什么 numpy 的 cumsum 比手动 C++ 的循环快得多？

Why is numpy's cumsum so much faster than manual C++'s loop?

c++

numpy

pybind11