src/backends/cpu/init.cpp - master - vuron/adept

src/backends/cpu/init.cpp
61 строка2 KB
kolkir
Fix tensor filling and module registration
06 мар 2025, 09:54
06 мар 2025, 09:54cea0aaa
#include <adept/backends/cpu/init.hpp>
#include <adept/threading.hpp>
#include <adept/types_dispatch.hpp>

#include <hwy/highway.h>

#include <random>

namespace adept {
namespace cpu {

namespace {
struct Random {
  mutable std::mt19937 engine{/*seed*/ 42};  // TODO:: make configurable
};
thread_local Random rnd;
}  // namespace

void fill_zero_tensor(Tensor& tensor) {
  DISPATCH_TYPE(tensor.properties().dtype, [&]() {
    scalar_t* data_ptr = tensor.mutable_data_ptr<scalar_t>();
    hwy::ZeroBytes(data_ptr, tensor.properties().shape.numel() * sizeof(scalar_t));
  });
}

void fill_tensor(Tensor& tensor, const std::any& scalar) {
  DISPATCH_TYPE(tensor.properties().dtype, [&]() {
    scalar_t* data_ptr = tensor.mutable_data_ptr<scalar_t>();
    auto value = std::any_cast<scalar_t>(scalar);
    std::fill_n(data_ptr, tensor.properties().shape.numel(), value);
  });
}

void fill_uniform_tensor(float32_t start, float32_t end, Tensor& tensor) {
  std::uniform_real_distribution<float32_t> dist(start, end);

  DISPATCH_TYPE(tensor.properties().dtype, [&]() {
    scalar_t* data_ptr = tensor.mutable_data_ptr<scalar_t>();
    parallel_for<scalar_t>(0, tensor.properties().shape.numel(), [&](auto begin, auto end) {
      for (auto i = begin; i < end; ++i) {
        data_ptr[i] = static_cast<scalar_t>(dist(rnd.engine));
      }
    });
  });
}

void fill_normal_tensor(float32_t mean, float32_t std, Tensor& tensor) {
  std::normal_distribution<float32_t> dist(mean, std);

  DISPATCH_TYPE(tensor.properties().dtype, [&]() {
    scalar_t* data_ptr = tensor.mutable_data_ptr<scalar_t>();
    parallel_for<scalar_t>(0, tensor.properties().shape.numel(), [&](auto begin, auto end) {
      for (auto i = begin; i < end; ++i) {
        data_ptr[i] = static_cast<scalar_t>(dist(rnd.engine));
      }
    });
  });
}

}  // namespace cpu
}  // namespace adept