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

src/backends/gpu/vulkan/manager.cpp
120 строк4 KB
kolkir
Make device selection by id for Vulkan
06 апр 2025, 16:33
06 апр 2025, 16:332d071a4
#include <adept/backends/gpu/vulkan/manager.hpp>
#include <adept/exception.hpp>
#include <adept/types_dispatch.hpp>

#include "tensorimpl.hpp"
#include "vulkanbuffer.hpp"

#include <flat_map/flat_hash_map.hpp>
#include <kompute/Kompute.hpp>
#include <kompute/operations/OpSyncDevice.hpp>

namespace adept::gpu::vulkan {
namespace {
template <typename T>
requires std::integral<T> || std::floating_point<T>
class VulkanAllocator : public Allocator {
 public:
  VulkanAllocator(kp::Manager& mgr) : mgr_(mgr) {}

  buffer_ptr_t allocate(size_t size) override {
    auto kp_tensor = mgr_.tensorT<T>(size);
    return std::make_shared<VulkanBuffer>(std::move(kp_tensor));
  }

  void free(buffer_ptr_t&& buffer) override { std::move(buffer).reset(); }
  size_t to_mem_size(size_t numel) override { return numel * sizeof(T); }

 private:
  kp::Manager& mgr_;
};
}  // namespace

struct Manager::ManagerInstance {
  std::shared_ptr<kp::Manager> mgr;
  std::shared_ptr<MemoryPool> mem_pool;
  std::recursive_mutex guard;
};

Manager::Manager() {
  kp::Manager mgr;
  auto devices = mgr.listDevices();
  if (devices.empty()) {
    THROW_ERROR("Failed to initialize Vulkan - compatible devices not found");
  }
  index_t idx = 0;
  for (auto& device : devices) {
    std::cout << "Vulkan device found: " << device.getProperties().deviceName << std::endl;

    auto dev_mgr = std::make_shared<ManagerInstance>();
    dev_mgr->mgr = std::make_shared<kp::Manager>(idx);
    dev_mgr->mem_pool =
        std::make_shared<MemoryPool>(std::make_shared<VulkanAllocator<float32_t>>(*dev_mgr->mgr));

    managers_.push_back(std::move(dev_mgr));
  }
}

Tensor Manager::make_empty_tensor(const TensorProperties& props) {
  if (props.shape.empty()) {
    THROW_ERROR("Can't create tensor with empty shape ", props);
  }
  auto devmgr = managers_.at(props.device_id);
  std::scoped_lock lock(devmgr->guard);
  auto buffer = devmgr->mem_pool->allocate(props.shape.numel());
  return Tensor(std::make_shared<TensorImpl>(props, std::move(buffer)));
}

Tensor Manager::make_tensor_from_blob(const void* data, const TensorProperties& props) {
  auto devmgr = managers_.at(props.device_id);
  std::scoped_lock lock(devmgr->guard);
  if (props.shape.empty()) {
    THROW_ERROR("Can't create tensor with empty shape ", props);
  }

  auto numel = props.shape.numel();
  auto buffer = devmgr->mem_pool->allocate(numel);
  auto* vlk_buffer = static_cast<VulkanBuffer*>(buffer->raw_buffer());
  vlk_buffer->kp_tensor->setData(data, vlk_buffer->size());
  devmgr->mgr->sequence()->eval<kp::OpSyncDevice>({vlk_buffer->kp_tensor})->eval();
  return Tensor(std::make_shared<TensorImpl>(props, std::move(buffer)));
}

Tensor Manager::clone_tensor(const Tensor& tensor) {
  auto devmgr = managers_.at(tensor.properties().device_id);
  std::scoped_lock lock(devmgr->guard);
  auto numel = tensor.properties().shape.numel();
  auto buffer = devmgr->mem_pool->allocate(numel);
  auto new_tensor = std::make_shared<TensorImpl>(tensor.properties(), std::move(buffer));
  const auto& in_tensor = static_cast<const TensorImpl&>(*tensor.impl());
  new_tensor->copy_from(in_tensor);
  return Tensor(new_tensor);
}

std::shared_ptr<kp::Sequence> Manager::make_sequence(index_t device_id) {
  auto devmgr = managers_.at(device_id);
  std::scoped_lock lock(devmgr->guard);
  return devmgr->mgr->sequence();
}

std::shared_ptr<kp::Algorithm> Manager::make_algorithm(
    index_t device_id,
    const std::vector<std::shared_ptr<kp::Memory>>& params,
    const std::vector<uint32_t>& shader) {
  auto devmgr = managers_.at(device_id);
  std::scoped_lock lock(devmgr->guard);
  return devmgr->mgr->algorithm(params, shader);
}

Tensor make_empty_tensor(const TensorProperties& props) {
  return Manager::instance().make_empty_tensor(props);
}

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

Tensor clone_tensor(const Tensor& tensor) {
  return Manager::instance().clone_tensor(tensor);
}
}  // namespace adept::gpu::vulkan