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

src/backends/gpu/vulkan/tensorimpl.hpp
99 строк3 KB
kolkir
Device checks and exchange
06 апр 2025, 15:28
06 апр 2025, 15:2873b23e8
#pragma once

#include <adept/backends/memory_pool.h>
#include <adept/device_tensor.hpp>
#include <adept/exception.hpp>
#include <adept/index_utils.hpp>
#include <adept/print.hpp>
#include <adept/shape.hpp>
#include <adept/tensor_props.hpp>
#include <adept/types.hpp>

#include <memory>

namespace adept::gpu::vulkan {

class TensorImpl : public DeviceTensor {
 public:
  TensorImpl(const TensorProperties& props, pool_buffer_ptr_t data_buffer);

  ~TensorImpl() override;
  TensorImpl(const TensorImpl&) = delete;
  TensorImpl& operator=(const TensorImpl&) = delete;
  TensorImpl(TensorImpl&& other);
  TensorImpl& operator=(TensorImpl&& other);

  const TensorProperties& properties() const override { return props_; }
  bool is_shape_aligned() const override { return true; }

  float32_t get_float32_at(const coords_t& coords) const override;
  float64_t get_float64_at(const coords_t& coords) const override;
  int32_t get_int32_at(const coords_t& coords) const override;
  int8_t get_int8_at(const coords_t& coords) const override;

  void add(const DeviceTensor& other) override;
  void add(float32_t scalar) override;
  void sub(const DeviceTensor& other) override;
  void sub(float32_t scalar) override;
  void mul(const DeviceTensor& other) override;
  void mul(float32_t scalar) override;
  void div(const DeviceTensor& other) override;
  void div(float32_t scalar) override;
  void neg() override;
  void exp() override;
  void sqrt() override;

  std::shared_ptr<DeviceTensor> gt(float32_t value) const override;
  std::shared_ptr<DeviceTensor> ge(float32_t value) const override;

  std::shared_ptr<DeviceTensor> lt(float32_t value) const override;
  std::shared_ptr<DeviceTensor> le(float32_t value) const override;

  std::shared_ptr<DeviceTensor> max() const override;
  std::shared_ptr<DeviceTensor> sum() const override;
  std::shared_ptr<DeviceTensor> sum_dim0() const override;
  std::shared_ptr<DeviceTensor> mean() const override;

  std::shared_ptr<DeviceTensor> transpose2d() const override;

  std::shared_ptr<DeviceTensor> matmul(const DeviceTensor& other) const override;

  const void* data() const override;

  void* mutable_data() override;

  const void* host_data() const override;

  bool is_same_backend(backend_t backend, index_t device_id) const override;

  const Indexer indexer() const override { return indexer_; }

  std::shared_ptr<DeviceTensor> squeeze(index_t dim) const override;
  std::shared_ptr<DeviceTensor> unsqueeze(index_t dim) const override;

  void copy_from(const TensorImpl& other);

 private:
  std::shared_ptr<DeviceTensor> sum_div(index_t divider) const;

  void check_dev_type_compatibility(const DeviceTensor& other) const;

  void check_data_type() const;

  template <typename DataType>
  DataType get(const coords_t& coords) const {
    auto expected_dtype = to_dtype<DataType>();
    if (props_.dtype != expected_dtype) {
      THROW_ERROR("Tensor::at() failed: incompatible data type requested ", props_.dtype,
                  " != ", expected_dtype);
    } else {
      return static_cast<const DataType*>(data())[indexer_.idxravel(coords)];
    }
  }

 private:
  TensorProperties props_;
  Indexer indexer_;
  pool_buffer_ptr_t data_buffer_;
};
}  // namespace adept::gpu::vulkan