src/tensor.cpp - master - glebestraikh/adept

src/tensor.cpp
393 строки10 KB
kolkir
Fix tensor data moving between GPUs
23 ноя 2025, 14:14
23 ноя 2025, 14:14997bd1b
#include <adept/backends/cpu/tensorimpl.hpp>
#include <adept/nn/init.hpp>
#include <adept/tensor.hpp>
#include <adept/tensor_print.hpp>
#include <adept/types.hpp>
#include <dispatch/log_softmax.hpp>
#include <dispatch/tensor_creation.hpp>

namespace adept {

Tensor Tensor::empty(const TensorProperties& props) {
  return make_empty_tensor(props);
}

Tensor Tensor::zero(const TensorProperties& props) {
  auto t = make_empty_tensor(props);
  fill_zero(t);
  return t;
}

Tensor Tensor::ones(const TensorProperties& props) {
  auto t = make_empty_tensor(props);
  for_dtype(props.dtype, [&]<typename T>() { fill(t, T{1}); });
  return t;
}

Tensor Tensor::uniform(float32_t start, float32_t end, const TensorProperties& props) {
  auto t = make_empty_tensor(props);
  fill_uniform(start, end, t);
  return t;
}

Tensor Tensor::normal(float32_t mean, float32_t std, const TensorProperties& props) {
  auto t = make_empty_tensor(props);
  fill_normal(mean, std, t);
  return t;
}

Tensor Tensor::from_blob(const void* data, const TensorProperties& props) {
  return make_tensor_from_blob(data, props);
}

Tensor Tensor::from_shared_mem_file(const std::string& filename, const TensorProperties& props) {
  return make_tensor_from_shared_mem_file(filename, props);
}

Tensor Tensor::stack(const std::vector<Tensor>& tensors) {
  if (tensors.empty()) {
    THROW_ERROR("Can't stack empty container");
  }
  return stack_tensors(tensors);
}

Tensor::Tensor(std::shared_ptr<DeviceTensor> impl) : impl_{impl} {};

Tensor::Tensor(const Tensor& other) : impl_(other.impl_) {}
Tensor& Tensor::operator=(const Tensor& other) {
  if (&other != this) {
    impl_ = other.impl_;
  }
  return *this;
}
Tensor::Tensor(Tensor&& other) : impl_(std::move(other.impl_)) {}
Tensor& Tensor::operator=(Tensor&& other) {
  if (&other != this) {
    impl_ = std::move(other.impl_);
  }
  return *this;
}

std::shared_ptr<DeviceTensor> Tensor::impl() const {
  return impl_;
}

const TensorProperties& Tensor::properties() const {
  CHECK(defined(), "Tensor is not defined!");
  return impl_->properties();
}

const Shape& Tensor::shape() const {
  return properties().shape;
}

dtype_t Tensor::dtype() const {
  return properties().dtype;
}

device_t Tensor::device() const {
  return properties().device;
}

bool Tensor::is_shape_aligned() const {
  CHECK(defined(), "Tensor is not defined!");
  return impl_->is_shape_aligned();
}

Tensor Tensor::reshape(const Shape& new_shape) const {
  CHECK(defined(), "Tensor is not defined!");
  return Tensor(impl_->reshape(new_shape));
}

Tensor Tensor::squeeze(index_t dim) const {
  CHECK(defined(), "Tensor is not defined!");
  return Tensor(impl_->squeeze(dim));
}

Tensor Tensor::unsqueeze(index_t dim) const {
  CHECK(defined(), "Tensor is not defined!");
  return Tensor(impl_->unsqueeze(dim));
}

Tensor Tensor::clone() const {
  return clone_tensor(*this);
}

void Tensor::copy_from(const Tensor& other) {
  CHECK(defined(), "Tensor is not defined!");
  CHECK(other.defined(), "operand Tensor is not defined!");
  bool device_id_same = properties().device_id == other.properties().device_id;
  if (device() == other.device() && device_id_same) {
    impl_->copy_from(*other.impl_);
  } else if (device() == device_t::GPU && other.device() == device_t::CPU) {
    CHECK(dtype() == other.dtype(), "Tensor::copy_from failed:: incompatible data types!");
    CHECK(shape() == other.shape(), "Tensor::copy_from failed with different shapes");
    impl_->copy_from(*other.impl_);
  } else {
    THROW_ERROR("Tensor::copy_from failed for ", properties(), " and ", other.properties());
  }
}

bool Tensor::defined() const {
  return impl_.get() != nullptr;
}

Tensor Tensor::add_(const Tensor& other) {
  CHECK(defined(), "Tensor is not defined!");
  CHECK(other.defined(), "operand Tensor is not defined!");
  impl_->add(*other.impl_);
  return {impl_};
}
Tensor Tensor::operator+=(const Tensor& other) {
  return add_(other);
}

Tensor Tensor::add_(float32_t scalar) {
  CHECK(defined(), "Tensor is not defined!");
  impl_->add(scalar);
  return {impl_};
}
Tensor Tensor::operator+=(float32_t scalar) {
  return add_(scalar);
}

Tensor Tensor::sub_(const Tensor& other) {
  CHECK(defined(), "Tensor is not defined!");
  CHECK(other.defined(), "operand Tensor is not defined!");
  impl_->sub(*other.impl_);
  return {impl_};
}
Tensor Tensor::operator-=(const Tensor& other) {
  return sub_(other);
}

Tensor Tensor::sub_(float32_t scalar) {
  CHECK(defined(), "Tensor is not defined!");
  impl_->sub(scalar);
  return {impl_};
}
Tensor Tensor::operator-=(float32_t scalar) {
  return sub_(scalar);
}

Tensor Tensor::mul_(const Tensor& other) {
  CHECK(defined(), "Tensor is not defined!");
  CHECK(other.defined(), "operand Tensor is not defined!");
  impl_->mul(*other.impl_);
  return {impl_};
}
Tensor Tensor::operator*=(const Tensor& other) {
  return mul_(other);
}

Tensor Tensor::mul_(float32_t scalar) {
  CHECK(defined(), "Tensor is not defined!");
  impl_->mul(scalar);
  return {impl_};
}
Tensor Tensor::operator*=(float32_t scalar) {
  return mul_(scalar);
}

Tensor Tensor::div_(const Tensor& other) {
  CHECK(defined(), "Tensor is not defined!");
  CHECK(other.defined(), "operand Tensor is not defined!");
  impl_->div(*other.impl_);
  return {impl_};
}
Tensor Tensor::operator/=(const Tensor& other) {
  return div_(other);
}

Tensor Tensor::div_(float32_t value) {
  CHECK(defined(), "Tensor is not defined!");
  impl_->div(value);
  return {impl_};
}
Tensor Tensor::operator/=(float32_t scalar) {
  return div_(scalar);
}

Tensor Tensor::eq_(const Tensor& other) const {
  CHECK(defined(), "Tensor is not defined!");
  CHECK(other.defined(), "operand Tensor is not defined!");
  return impl_->eq(*other.impl_);
}

Tensor Tensor::gt_(float32_t value) const {
  CHECK(defined(), "Tensor is not defined!");
  return impl_->gt(value);
}
Tensor Tensor::ge_(float32_t value) const {
  CHECK(defined(), "Tensor is not defined!");
  return impl_->ge(value);
}
Tensor Tensor::lt_(float32_t value) const {
  CHECK(defined(), "Tensor is not defined!");
  return impl_->lt(value);
}
Tensor Tensor::le_(float32_t value) const {
  CHECK(defined(), "Tensor is not defined!");
  return impl_->le(value);
}

Tensor Tensor::softmax_last_dim() const {
  CHECK(defined(), "Tensor is not defined!");
  return ::adept::softmax_last_dim(*this);
}

Tensor Tensor::neg_() {
  CHECK(defined(), "Tensor is not defined!");
  impl_->neg();
  return {impl_};
}

Tensor Tensor::exp_() {
  CHECK(defined(), "Tensor is not defined!");
  impl_->exp();
  return {impl_};
}

Tensor Tensor::sqrt_() {
  CHECK(defined(), "Tensor is not defined!");
  impl_->sqrt();
  return {impl_};
}

Tensor Tensor::invsqrt_() {
  CHECK(defined(), "Tensor is not defined!");
  impl_->invsqrt();
  return {impl_};
}

Tensor Tensor::sqrt() const {
  CHECK(defined(), "Tensor is not defined!");
  return clone_tensor(*this).sqrt_();
}

Tensor Tensor::invsqrt() const {
  CHECK(defined(), "Tensor is not defined!");
  return clone_tensor(*this).invsqrt_();
}

Tensor Tensor::operator-() {
  return neg_();
}

Tensor Tensor::max() const {
  CHECK(defined(), "Tensor is not defined!");
  return (impl_->max());
}

Tensor Tensor::sum() const {
  CHECK(defined(), "Tensor is not defined!");
  return (impl_->sum());
}

Tensor Tensor::sum_dim0() const {
  CHECK(defined(), "Tensor is not defined!");
  return (impl_->sum_dim0());
}

std::tuple<Tensor, Tensor> Tensor::max_last_dim() const {
  CHECK(defined(), "Tensor is not defined!");
  return impl_->max_last_dim();
}

Tensor Tensor::mean() const {
  CHECK(defined(), "Tensor is not defined!");
  return (impl_->mean());
}

Tensor Tensor::transpose2d() const {
  CHECK(defined(), "Tensor is not defined!");
  return impl_->transpose2d();
}

Tensor Tensor::matmul(const Tensor& other) const {
  CHECK(defined(), "Tensor is not defined!");
  CHECK(other.defined(), "operand Tensor is not defined!");
  return impl_->matmul(*other.impl_);
}

std::string Tensor::to_string() const {
  std::stringstream out;
  out << *this;
  return out.str();
}

Tensor Tensor::cpu() const {
  CHECK(defined(), "Tensor is not defined!");
  if (device() != device_t::CPU || is_shared() ||
      !impl_->is_same_backend(Dispatcher::instance().get_current_backend(device_t::CPU), 0)) {
    auto new_tensor =
        Tensor::empty(TensorProperties{properties().shape, device_t::CPU, properties().dtype, 0});
    auto& cpu_impl = new_tensor.impl<cpu::TensorImpl>();
    impl_->copy_to(cpu_impl.mutable_data());
    return new_tensor;
  } else {
    return {*this};
  }
}

Tensor Tensor::gpu(index_t device_id) const {
  CHECK(defined(), "Tensor is not defined!");
  if (device() != device_t::GPU ||
      !impl_->is_same_backend(Dispatcher::instance().get_current_backend(device_t::GPU),
                              device_id)) {
    // TODO: later improve GPU to GPU transfer
    auto host_tensor = cpu();
    return Tensor::from_blob(
        host_tensor.impl<cpu::TensorImpl>().data(),
        TensorProperties{properties().shape, device_t::GPU, properties().dtype, device_id});
  } else {
    return {*this};
  }
}

Tensor Tensor::to(device_t device, index_t device_id) const {
  CHECK(defined(), "Tensor is not defined!");
  if (device == device_t::CPU) {
    return cpu();
  } else if (device == device_t::GPU) {
    return gpu(device_id);
  } else {
    THROW_ERROR("Tensor::to got incorrect device type!");
  }
}

Tensor Tensor::to(dtype_t new_type) const {
  CHECK(defined(), "Tensor is not defined!");
  if (new_type != dtype()) {
    return impl_->type_cast(new_type);
  }
  return {*this};
}

void Tensor::share_memory() {
  CHECK(defined(), "Tensor is not defined!");
  CHECK(device() == device_t::CPU, "Only CPU tensors can have shared memory!");
  impl_->share_memory();
}

bool Tensor::is_shared() const {
  CHECK(defined(), "Tensor is not defined!");
  return impl_->is_shared();
}

const std::string& Tensor::share_filename() const {
  CHECK(is_shared(), "Tensor is not shared!");
  return impl_->share_filename();
}

void Tensor::inc_shared_refcounter() {
  CHECK(is_shared(), "Tensor is not shared!");
  impl_->inc_shared_refcounter();
}

void Tensor::dec_shared_refcounter() {
  CHECK(is_shared(), "Tensor is not shared!");
  impl_->dec_shared_refcounter();
}

}  // namespace adept