python/cpp_bindings/tensor_py.cpp - master - glebestraikh/adept

python/cpp_bindings/tensor_py.cpp
202 строки9 KB
kolkir
Release GIL in Python bindings for long memory operations
30 авг 2025, 18:34
30 авг 2025, 18:348186655
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include <adept/backends/cpu/tensorimpl.hpp>
#include <adept/exception.hpp>
#include <adept/index_utils.hpp>
#include <adept/tensor.hpp>
#include <adept/tensor_print.hpp>
#include <adept/types_dispatch.hpp>

#include <sstream>

namespace py = pybind11;

using namespace adept;

using ScalarValue = std::variant<float32_t, float64_t, int32_t, int64_t, int8_t>;

void bind_tensor(py::module& m) {
  py::class_<Tensor>(m, "Tensor", py::buffer_protocol())
      .def(py::init([](const std::vector<float32_t>& data, const TensorProperties& props) {
        if (props.dtype != dtype_t::Float32)
          THROW_ERROR("Inconsistent type ", props.dtype, " for Tensor creation!");
        py::gil_scoped_release no_gil;
        return Tensor::from_blob(data.data(), props);
      }))
      .def(py::init([](const std::vector<float64_t>& data, const TensorProperties& props) {
        if (props.dtype != dtype_t::Float64)
          THROW_ERROR("Inconsistent type ", props.dtype, " for Tensor creation!");
        py::gil_scoped_release no_gil;
        return Tensor::from_blob(data.data(), props);
      }))
      .def(py::init([](const std::vector<int32_t>& data, const TensorProperties& props) {
        if (props.dtype != dtype_t::Int32)
          THROW_ERROR("Inconsistent type ", props.dtype, " for Tensor creation!");
        py::gil_scoped_release no_gil;
        return Tensor::from_blob(data.data(), props);
      }))
      .def(py::init([](const std::vector<int8_t>& data, const TensorProperties& props) {
        if (props.dtype != dtype_t::Int8)
          THROW_ERROR("Inconsistent type ", props.dtype, " for Tensor creation!");
        py::gil_scoped_release no_gil;
        return Tensor::from_blob(data.data(), props);
      }))
      .def_property_readonly_static("empty",
                                    [](py::object) {
                                      return py::cpp_function([](const TensorProperties& props) {
                                        py::gil_scoped_release no_gil;
                                        return Tensor::empty(props);
                                      });
                                    })
      .def_property_readonly_static("zero",
                                    [](py::object) {
                                      return py::cpp_function([](const TensorProperties& props) {
                                        py::gil_scoped_release no_gil;
                                        return Tensor::zero(props);
                                      });
                                    })
      .def_property_readonly_static("uniform",
                                    [](py::object) {
                                      return py::cpp_function([](float32_t start, float32_t end,
                                                                 const TensorProperties& props) {
                                        py::gil_scoped_release no_gil;
                                        return Tensor::uniform(start, end, props);
                                      });
                                    })
      .def_property_readonly_static("normal",
                                    [](py::object) {
                                      return py::cpp_function([](float32_t mean, float32_t std,
                                                                 const TensorProperties& props) {
                                        py::gil_scoped_release no_gil;
                                        return Tensor::normal(mean, std, props);
                                      });
                                    })
      .def_property_readonly_static(
          "from_shared_mem_file",
          [](py::object) {
            return py::cpp_function([](const std::string& filename, const TensorProperties& props) {
              py::gil_scoped_release no_gil;
              return Tensor::from_shared_mem_file(filename, props);
            });
          })
      .def_property_readonly_static(
          "stack",
          [](py::object) {
            return py::cpp_function([](const std::vector<Tensor>& tensors) {
              py::gil_scoped_release no_gil;
              return Tensor::stack(tensors);
            });
          })
      .def(py::init([](py::buffer b) {
        py::buffer_info info = b.request();
        dtype_t dtype = dtype_t::Float32;
        if (info.item_type_is_equivalent_to<float32_t>()) {
          CHECK(sizeof(float32_t) == info.itemsize, "Invlaid python buffer type size");
          dtype = dtype_t::Float32;
        } else if (info.item_type_is_equivalent_to<float64_t>()) {
          CHECK(sizeof(float64_t) == info.itemsize, "Invlaid python buffer type size");
          dtype = dtype_t::Float64;
        } else if (info.item_type_is_equivalent_to<int32_t>()) {
          CHECK(sizeof(int32_t) == info.itemsize, "Invlaid python buffer type size");
          dtype = dtype_t::Int32;
        } else {
          THROW_ERROR("Tensor can't be constructed with the ", info.format, " type");
        }
        TensorProperties props(Shape(info.shape), device_t::CPU, dtype);
        py::gil_scoped_release no_gil;
        return Tensor::from_blob(info.ptr, props);
      }))
      .def_buffer([](Tensor& t) -> py::buffer_info {
        if (t.device() != device_t::CPU) {
          THROW_ERROR("Can't convert GPU tensor to numpy");
        }
        py::buffer_info info;
        DISPATCH_TYPE(t.dtype(), [&]() {
          info = py::buffer_info(
              t.impl<cpu::TensorImpl>().mutable_data_ptr<scalar_t>(), /* Pointer to buffer */
              t.shape().dims(),                                       /* shape_in */
              make_continuous_byte_strides<scalar_t>(t.shape()),      /* stride_in */
              true /* Readonly */);
        });
        return info;
      })
      .def("cpu",
           [](const Tensor& t) {
             py::gil_scoped_release no_gil;
             return t.cpu();
           })
      .def(
          "gpu",
          [](const Tensor& t, index_t device_id) {
            py::gil_scoped_release no_gil;
            return t.gpu(device_id);
          },
          py::arg("device_id") = 0)
      .def(
          "to",
          [](const Tensor& t, device_t device, index_t device_id) {
            py::gil_scoped_release no_gil;
            return t.to(device, device_id);
          },
          py::arg("device"), py::arg("device_id") = 0)
      .def("clone",
           [](Tensor& t) {
             py::gil_scoped_release no_gil;
             return t.clone();
           })
      .def("copy_from",
           [](Tensor& t, const Tensor& other) {
             py::gil_scoped_release no_gil;
             t.copy_from(other);
           })
      .def("properties", &Tensor::properties)
      .def("sum", &Tensor::sum)
      .def("mean", &Tensor::mean)
      .def("shape", [](const Tensor& self) { return self.properties().shape; })
      .def("__repr__",
           [](const Tensor& t) {
             std::stringstream buf;
             buf << t;
             return buf.str();
           })
      .def("__add__", py::overload_cast<const Tensor&>(&Tensor::add_))
      .def("__add__", py::overload_cast<float32_t>(&Tensor::add_))
      .def("__sub__", py::overload_cast<const Tensor&>(&Tensor::sub_))
      .def("__sub__", py::overload_cast<float32_t>(&Tensor::sub_))
      .def("__mul__", py::overload_cast<const Tensor&>(&Tensor::mul_))
      .def("__mul__", py::overload_cast<float32_t>(&Tensor::mul_))
      .def("__truediv__", py::overload_cast<const Tensor&>(&Tensor::div_))
      .def("__truediv__", py::overload_cast<float32_t>(&Tensor::div_))
      .def("__eq__", &Tensor::eq_)
      .def("__gt__", &Tensor::gt_)
      .def("__ge__", &Tensor::ge_)
      .def("__lt__", &Tensor::lt_)
      .def("__le__", &Tensor::le_)
      .def("__neg__", &Tensor::neg_)
      .def("exp_", &Tensor::exp_)
      .def("sqrt_", &Tensor::sqrt_)
      .def("sqrt", &Tensor::sqrt)
      .def("transpose2d", &Tensor::transpose2d)
      .def("matmul", &Tensor::matmul)
      .def("at",
           [](const Tensor& self, const coords_t& coords) -> ScalarValue {
             ScalarValue value;
             DISPATCH_TYPE(self.dtype(), [&]() { value = self.at<scalar_t>(coords); });
             return value;
           })
      .def("squeeze", &Tensor::squeeze)
      .def("unsqueeze", &Tensor::unsqueeze)
      .def("reshape", &Tensor::reshape)
      .def("softmax_last_dim", &Tensor::softmax_last_dim)
      .def("max_last_dim", &Tensor::max_last_dim)
      .def("share_memory",
           [](Tensor& t) {
             py::gil_scoped_release no_gil;
             t.share_memory();
           })
      .def("is_shared", &Tensor::is_shared)
      .def("share_filename", &Tensor::share_filename)
      .def("inc_shared_refcounter", &Tensor::inc_shared_refcounter)
      .def("dec_shared_refcounter", &Tensor::dec_shared_refcounter);
}