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

src/backends/cpu/cpu_backend_init.cpp
105 строк5 KB
kolkir
Multidevice test utils
06 апр 2025, 19:28
06 апр 2025, 19:2806a1d99
#include <adept/backends/cpu/activations.hpp>
#include <adept/backends/cpu/avgpool2d.hpp>
#include <adept/backends/cpu/batchnorm2d.hpp>
#include <adept/backends/cpu/conv2d.hpp>
#include <adept/backends/cpu/init.hpp>
#include <adept/backends/cpu/log_softmax.hpp>
#include <adept/backends/cpu/maxpool2d.hpp>
#include <adept/backends/cpu/nll_loss.hpp>
#include <adept/backends/cpu/stack.hpp>
#include <adept/backends/cpu/tensor_factory.hpp>
#include <adept/dispatch/dispatcher.hpp>
#include "../../init_order.hpp"

#include <cblas.h>

#include <tbb/tbb.h>

namespace adept {

__attribute__((__constructor__(BACKEND_INIT_ORDER))) void cpu_init(void) {
  int num_threads = oneapi::tbb::info::default_concurrency();
  // TODO: the following global_limit should be configured somewhere else
  // oneapi::tbb::global_control global_limit(oneapi::tbb::global_control::max_allowed_parallelism,
  //                                          num_threads / 2);
  openblas_set_num_threads(num_threads / 2);

  Dispatcher::instance().register_backend(device_t::CPU, backend_t::SIMD);

  // tensor creation
  Dispatcher::instance().add_func("make_empty_tensor", {device_t::CPU, backend_t::SIMD},
                                  cpu::make_empty_tensor);
  Dispatcher::instance().add_func("make_tensor_from_blob", {device_t::CPU, backend_t::SIMD},
                                  cpu::make_tensor_from_blob);
  Dispatcher::instance().add_func("clone_tensor", {device_t::CPU, backend_t::SIMD},
                                  cpu::clone_tensor);
  Dispatcher::instance().add_func("stack_tensors", {device_t::CPU, backend_t::SIMD},
                                  cpu::stack_tensors);
  // tensor creation end

  // init
  Dispatcher::instance().add_func("fill_zero_tensor", {device_t::CPU, backend_t::SIMD},
                                  cpu::fill_zero_tensor);
  Dispatcher::instance().add_func("fill_tensor", {device_t::CPU, backend_t::SIMD},
                                  cpu::fill_tensor);
  Dispatcher::instance().add_func("fill_uniform_tensor", {device_t::CPU, backend_t::SIMD},
                                  cpu::fill_uniform_tensor);
  Dispatcher::instance().add_func("fill_normal_tensor", {device_t::CPU, backend_t::SIMD},
                                  cpu::fill_normal_tensor);
  // init end

  // losses
  Dispatcher::instance().add_func("log_softmax_last_dim_fwd", {device_t::CPU, backend_t::SIMD},
                                  cpu::log_softmax_last_dim_fwd);
  Dispatcher::instance().add_func("log_softmax_last_dim_bwd", {device_t::CPU, backend_t::SIMD},
                                  cpu::log_softmax_last_dim_bwd);
  Dispatcher::instance().add_func("nll_loss_fwd", {device_t::CPU, backend_t::SIMD},
                                  cpu::nll_loss_fwd);
  Dispatcher::instance().add_func("nll_loss_bwd", {device_t::CPU, backend_t::SIMD},
                                  cpu::nll_loss_bwd);
  // losses end

  // conv
  Dispatcher::instance().add_func("conv2d_fwd", {device_t::CPU, backend_t::SIMD}, cpu::conv2d_fwd);
  Dispatcher::instance().add_func("conv2d_bwd", {device_t::CPU, backend_t::SIMD}, cpu::conv2d_bwd);
  // conv end

  // pool
  Dispatcher::instance().add_func("avg_pool2d_fwd", {device_t::CPU, backend_t::SIMD},
                                  cpu::avg_pool2d_fwd);
  Dispatcher::instance().add_func("avg_pool2d_bwd", {device_t::CPU, backend_t::SIMD},
                                  cpu::avg_pool2d_bwd);

  Dispatcher::instance().add_func("max_pool2d_fwd", {device_t::CPU, backend_t::SIMD},
                                  cpu::max_pool2d_fwd);
  Dispatcher::instance().add_func("max_pool2d_bwd", {device_t::CPU, backend_t::SIMD},
                                  cpu::max_pool2d_bwd);
  // pool end

  // norms
  Dispatcher::instance().add_func("batchnorm2d_fwd", {device_t::CPU, backend_t::SIMD},
                                  cpu::batchnorm2d_fwd);
  Dispatcher::instance().add_func("batchnorm2d_bwd", {device_t::CPU, backend_t::SIMD},
                                  cpu::batchnorm2d_bwd);
  // norms end

  // activations
  Dispatcher::instance().add_func("relu_fwd", {device_t::CPU, backend_t::SIMD}, cpu::relu_fwd);
  Dispatcher::instance().add_func("relu_bwd", {device_t::CPU, backend_t::SIMD}, cpu::relu_bwd);

  Dispatcher::instance().add_func("leaky_relu_fwd", {device_t::CPU, backend_t::SIMD},
                                  cpu::leaky_relu_fwd);
  Dispatcher::instance().add_func("leaky_relu_bwd", {device_t::CPU, backend_t::SIMD},
                                  cpu::leaky_relu_bwd);

  Dispatcher::instance().add_func("sigmoid_fwd", {device_t::CPU, backend_t::SIMD},
                                  cpu::sigmoid_fwd);
  Dispatcher::instance().add_func("sigmoid_bwd", {device_t::CPU, backend_t::SIMD},
                                  cpu::sigmoid_bwd);

  Dispatcher::instance().add_func("silu_fwd", {device_t::CPU, backend_t::SIMD}, cpu::silu_fwd);
  Dispatcher::instance().add_func("silu_bwd", {device_t::CPU, backend_t::SIMD}, cpu::silu_bwd);
  // activations end
}

}  // namespace adept