src/backends/gpu/vulkan/tensorimpl.cpp - master - vuron/adept

src/backends/gpu/vulkan/tensorimpl.cpp
222 строки6 KB
kolkir
Make device selection by id for Vulkan
06 апр 2025, 16:33
06 апр 2025, 16:332d071a4
#include "tensorimpl.hpp"

#include <adept/backends/gpu/vulkan/manager.hpp>
#include "vulkanbuffer.hpp"

#include <kompute/operations/OpCopy.hpp>
#include <kompute/operations/OpSyncDevice.hpp>

#include "shaders/addf.hpp"

namespace adept::gpu::vulkan {

TensorImpl::TensorImpl(const TensorProperties& props, pool_buffer_ptr_t data_buffer)
    : props_(props), indexer_{&props_.shape}, data_buffer_(std::move(data_buffer)) {
  check_data_type();
}

TensorImpl::~TensorImpl() {}

TensorImpl::TensorImpl(TensorImpl&& other)
    : props_(std::move(other.props_)),
      indexer_(std::move(other.indexer_)),
      data_buffer_(std::move(other.data_buffer_)) {
  check_data_type();
}

TensorImpl& TensorImpl::operator=(TensorImpl&& other) {
  if (&other != this) {
    props_ = std::move(other.props_);
    indexer_ = std::move(other.indexer_);
    data_buffer_ = std::move(other.data_buffer_);
  }
  check_data_type();
  return *this;
}

void TensorImpl::copy_from(const TensorImpl& other) {
  check_dev_type_compatibility(other);
  auto* vlk_buffer = static_cast<VulkanBuffer*>(data_buffer_->raw_buffer());
  auto* other_vlk_buffer = static_cast<VulkanBuffer*>(other.data_buffer_->raw_buffer());
  auto seq = Manager::instance().make_sequence(props_.device_id);
  const std::vector<std::shared_ptr<kp::Memory>> params = {other_vlk_buffer->kp_tensor,
                                                           vlk_buffer->kp_tensor};
  seq->eval<kp::OpCopy>(params)->eval();
}

float32_t TensorImpl::get_float32_at(const coords_t& coords) const {
  return get<float32_t>(coords);
}
float64_t TensorImpl::get_float64_at(const coords_t& coords) const {
  return get<float64_t>(coords);
}
int32_t TensorImpl::get_int32_at(const coords_t& coords) const {
  return get<int32_t>(coords);
}
int8_t TensorImpl::get_int8_at(const coords_t& coords) const {
  return get<int8_t>(coords);
}

void TensorImpl::add(const DeviceTensor& other) {
  check_dev_type_compatibility(other);
  auto* vlk_buffer = static_cast<VulkanBuffer*>(data_buffer_->raw_buffer());
  auto* other_vlk_buffer =
      static_cast<VulkanBuffer*>(static_cast<const TensorImpl&>(other).data_buffer_->raw_buffer());

  const std::vector<std::shared_ptr<kp::Memory>> params = {
      vlk_buffer->kp_tensor,
      other_vlk_buffer->kp_tensor,
      vlk_buffer->kp_tensor,
  };

  std::shared_ptr<kp::Algorithm> algo =
      Manager::instance().make_algorithm(props_.device_id, params, ADDF_SPV);
  Manager::instance().make_sequence(props_.device_id)->eval<kp::OpAlgoDispatch>(algo);
}

void TensorImpl::add(float32_t scalar) {
  THROW_ERROR("Method is not implemented");
}

void TensorImpl::sub(const DeviceTensor& other) {
  check_dev_type_compatibility(other);
  THROW_ERROR("Method is not implemented");
}

void TensorImpl::sub(float32_t scalar) {
  THROW_ERROR("Method is not implemented");
}

void TensorImpl::mul(const DeviceTensor& other) {
  check_dev_type_compatibility(other);
  THROW_ERROR("Method is not implemented");
}

void TensorImpl::mul(float32_t scalar) {
  THROW_ERROR("Method is not implemented");
}

void TensorImpl::div(const DeviceTensor& other) {
  check_dev_type_compatibility(other);
  THROW_ERROR("Method is not implemented");
}

void TensorImpl::div(float32_t scalar) {
  THROW_ERROR("Method is not implemented");
}

void TensorImpl::neg() {
  THROW_ERROR("Method is not implemented");
}

void TensorImpl::exp() {
  THROW_ERROR("Method is not implemented");
}

void TensorImpl::sqrt() {
  THROW_ERROR("Method is not implemented");
}

std::shared_ptr<DeviceTensor> TensorImpl::gt(float32_t value) const {
  THROW_ERROR("Method is not implemented");
  return {};
}

std::shared_ptr<DeviceTensor> TensorImpl::ge(float32_t value) const {
  THROW_ERROR("Method is not implemented");
  return {};
}

std::shared_ptr<DeviceTensor> TensorImpl::lt(float32_t value) const {
  THROW_ERROR("Method is not implemented");
  return {};
}

std::shared_ptr<DeviceTensor> TensorImpl::le(float32_t value) const {
  THROW_ERROR("Method is not implemented");
  return {};
}

std::shared_ptr<DeviceTensor> TensorImpl::max() const {
  THROW_ERROR("Method is not implemented");
  return {};
}

std::shared_ptr<DeviceTensor> TensorImpl::sum() const {
  THROW_ERROR("Method is not implemented");
  return {};
}

std::shared_ptr<DeviceTensor> TensorImpl::sum_dim0() const {
  THROW_ERROR("Method is not implemented");
  return {};
}

std::shared_ptr<DeviceTensor> TensorImpl::mean() const {
  THROW_ERROR("Method is not implemented");
  return {};
}

std::shared_ptr<DeviceTensor> TensorImpl::transpose2d() const {
  THROW_ERROR("Method is not implemented");
  return {};
}

std::shared_ptr<DeviceTensor> TensorImpl::matmul(const DeviceTensor& other) const {
  check_dev_type_compatibility(other);
  THROW_ERROR("Method is not implemented");
  return {};
}

const void* TensorImpl::data() const {
  auto* vlk_buffer = static_cast<VulkanBuffer*>(data_buffer_->raw_buffer());
  return vlk_buffer->kp_tensor->rawData();
}

void* TensorImpl::mutable_data() {
  THROW_ERROR("Method is not implemented");
  return {};
}

const void* TensorImpl::host_data() const {
  auto* vlk_buffer = static_cast<VulkanBuffer*>(data_buffer_->raw_buffer());
  auto seq = Manager::instance().make_sequence(props_.device_id);
  seq->eval<kp::OpSyncLocal>({vlk_buffer->kp_tensor})->eval();
  return data();
}

bool TensorImpl::is_same_backend(backend_t backend, index_t device_id) const {
  return backend == backend_t::Vulkan && props_.device_id == device_id;
}

std::shared_ptr<DeviceTensor> TensorImpl::squeeze(index_t dim) const {
  THROW_ERROR("Method is not implemented");
  return {};
}

std::shared_ptr<DeviceTensor> TensorImpl::unsqueeze(index_t dim) const {
  THROW_ERROR("Method is not implemented");
  return {};
}

void TensorImpl::check_dev_type_compatibility(const DeviceTensor& other) const {
  if (props_.device != other.properties().device)
    THROW_ERROR("Vulkan Tensor op failed: incompatible devices ", props_.device,
                " != ", other.properties().device);

  if (props_.device_id != other.properties().device_id)
    THROW_ERROR("Vulkan Tensor op failed: incompatible devices ids ", props_.device_id,
                " != ", other.properties().device_id);

  if (props_.dtype != other.properties().dtype)
    THROW_ERROR("Vulkan Tensor op failed:: incompatible data types!");
}

void TensorImpl::check_data_type() const {
  if (props_.dtype == dtype_t::Float64 || props_.dtype == dtype_t::Int32 ||
      props_.dtype == dtype_t::Int8) {
    THROW_ERROR("Vulkan Tensor usupported data types: ", props_.dtype);
  }
}

}  // namespace adept::gpu::vulkan