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

src/backends/cpu/activations.cpp
152 строки6 KB
kolkir
Activation functions implementaions
12 фев 2025, 00:20
12 фев 2025, 00:20a5158cd
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE "../src/backends/cpu/activations.cpp"
#include <hwy/foreach_target.h>

#include <hwy/highway.h>
#include <hwy/per_target.h>
#include "activations-inl.hpp"

#if HWY_ONCE

#include <adept/backends/cpu/activations.hpp>
#include <adept/threading.hpp>
#include <adept/types_dispatch.hpp>

namespace adept::cpu {
Tensor relu_fwd(const Tensor& input) {
  auto result = Tensor::empty(input.properties());
  DISPATCH_TYPE(input.dtype(), [&]() {
    const auto* in_base_ptr = input.const_data_ptr<scalar_t>();
    auto* out_base_ptr = result.mutable_data_ptr<scalar_t>();
    parallel_for<scalar_t>(0, input.shape().numel(), [&](auto begin, auto end) {
      auto* in_data_ptr = in_base_ptr + begin;
      auto* out_data_ptr = out_base_ptr + begin;
      auto size = end - begin;
      HWY_EXPORT_AND_DYNAMIC_DISPATCH_T(simd_relu_fwd<scalar_t>)
      (out_data_ptr, in_data_ptr, size, input.is_shape_aligned());
    });
  });
  return result;
}

Tensor relu_bwd(const Tensor& input, const Tensor& out_grad) {
  auto result = Tensor::empty(input.properties());
  DISPATCH_TYPE(input.dtype(), [&]() {
    const auto* in_base_ptr = input.const_data_ptr<scalar_t>();
    const auto* grad_base_ptr = out_grad.const_data_ptr<scalar_t>();
    auto* out_base_ptr = result.mutable_data_ptr<scalar_t>();
    parallel_for<scalar_t>(0, input.shape().numel(), [&](auto begin, auto end) {
      const auto* in_data_ptr = in_base_ptr + begin;
      const auto* grad_data_ptr = grad_base_ptr + begin;
      auto* out_data_ptr = out_base_ptr + begin;
      auto size = end - begin;
      HWY_EXPORT_AND_DYNAMIC_DISPATCH_T(simd_relu_bwd<scalar_t>)
      (out_data_ptr, in_data_ptr, grad_data_ptr, size, input.is_shape_aligned());
    });
  });
  return result;
}

Tensor leaky_relu_fwd(const Tensor& input, float32_t negative_slope) {
  auto result = Tensor::empty(input.properties());
  DISPATCH_TYPE(input.dtype(), [&]() {
    const auto* in_base_ptr = input.const_data_ptr<scalar_t>();
    auto* out_base_ptr = result.mutable_data_ptr<scalar_t>();
    parallel_for<scalar_t>(0, input.shape().numel(), [&](auto begin, auto end) {
      auto* in_data_ptr = in_base_ptr + begin;
      auto* out_data_ptr = out_base_ptr + begin;
      auto size = end - begin;
      HWY_EXPORT_AND_DYNAMIC_DISPATCH_T(simd_leaky_relu_fwd<scalar_t>)
      (out_data_ptr, in_data_ptr, negative_slope, size, input.is_shape_aligned());
    });
  });
  return result;
}

Tensor leaky_relu_bwd(const Tensor& input, const Tensor& out_grad, float32_t negative_slope) {
  auto result = Tensor::empty(input.properties());
  DISPATCH_TYPE(input.dtype(), [&]() {
    const auto* in_base_ptr = input.const_data_ptr<scalar_t>();
    const auto* grad_base_ptr = out_grad.const_data_ptr<scalar_t>();
    auto* out_base_ptr = result.mutable_data_ptr<scalar_t>();
    parallel_for<scalar_t>(0, input.shape().numel(), [&](auto begin, auto end) {
      const auto* in_data_ptr = in_base_ptr + begin;
      const auto* grad_data_ptr = grad_base_ptr + begin;
      auto* out_data_ptr = out_base_ptr + begin;
      auto size = end - begin;
      HWY_EXPORT_AND_DYNAMIC_DISPATCH_T(simd_leaky_relu_bwd<scalar_t>)
      (out_data_ptr, in_data_ptr, grad_data_ptr, negative_slope, size, input.is_shape_aligned());
    });
  });
  return result;
}

Tensor sigmoid_fwd(const Tensor& input) {
  auto result = Tensor::empty(input.properties());
  DISPATCH_TYPE(input.dtype(), [&]() {
    const auto* in_base_ptr = input.const_data_ptr<scalar_t>();
    auto* out_base_ptr = result.mutable_data_ptr<scalar_t>();
    parallel_for<scalar_t>(0, input.shape().numel(), [&](auto begin, auto end) {
      auto* in_data_ptr = in_base_ptr + begin;
      auto* out_data_ptr = out_base_ptr + begin;
      auto size = end - begin;
      HWY_EXPORT_AND_DYNAMIC_DISPATCH_T(simd_sigmoid_fwd<scalar_t>)
      (out_data_ptr, in_data_ptr, size, input.is_shape_aligned());
    });
  });
  return result;
}

Tensor sigmoid_bwd(const Tensor& op_result, const Tensor& out_grad) {
  auto result = Tensor::empty(op_result.properties());
  DISPATCH_TYPE(op_result.dtype(), [&]() {
    const auto* r_base_ptr = op_result.const_data_ptr<scalar_t>();
    const auto* grad_base_ptr = out_grad.const_data_ptr<scalar_t>();
    auto* out_base_ptr = result.mutable_data_ptr<scalar_t>();
    parallel_for<scalar_t>(0, op_result.shape().numel(), [&](auto begin, auto end) {
      const auto* r_data_ptr = r_base_ptr + begin;
      const auto* grad_data_ptr = grad_base_ptr + begin;
      auto* out_data_ptr = out_base_ptr + begin;
      auto size = end - begin;
      HWY_EXPORT_AND_DYNAMIC_DISPATCH_T(simd_sigmoid_bwd<scalar_t>)
      (out_data_ptr, r_data_ptr, grad_data_ptr, size, op_result.is_shape_aligned());
    });
  });
  return result;
}

Tensor silu_fwd(const Tensor& input) {
  auto result = Tensor::empty(input.properties());
  DISPATCH_TYPE(input.dtype(), [&]() {
    const auto* in_base_ptr = input.const_data_ptr<scalar_t>();
    auto* out_base_ptr = result.mutable_data_ptr<scalar_t>();
    parallel_for<scalar_t>(0, input.shape().numel(), [&](auto begin, auto end) {
      auto* in_data_ptr = in_base_ptr + begin;
      auto* out_data_ptr = out_base_ptr + begin;
      auto size = end - begin;
      HWY_EXPORT_AND_DYNAMIC_DISPATCH_T(simd_silu_fwd<scalar_t>)
      (out_data_ptr, in_data_ptr, size, input.is_shape_aligned());
    });
  });
  return result;
}

Tensor silu_bwd(const Tensor& input, const Tensor& out_grad) {
  auto result = Tensor::empty(input.properties());
  DISPATCH_TYPE(input.dtype(), [&]() {
    const auto* r_base_ptr = input.const_data_ptr<scalar_t>();
    const auto* grad_base_ptr = out_grad.const_data_ptr<scalar_t>();
    auto* out_base_ptr = result.mutable_data_ptr<scalar_t>();
    parallel_for<scalar_t>(0, input.shape().numel(), [&](auto begin, auto end) {
      const auto* r_data_ptr = r_base_ptr + begin;
      const auto* grad_data_ptr = grad_base_ptr + begin;
      auto* out_data_ptr = out_base_ptr + begin;
      auto size = end - begin;
      HWY_EXPORT_AND_DYNAMIC_DISPATCH_T(simd_silu_bwd<scalar_t>)
      (out_data_ptr, r_data_ptr, grad_data_ptr, size, input.is_shape_aligned());
    });
  });
  return result;
}
}  // namespace adept::cpu
#endif  // HWY_ONCE