diff --git a/src/native/cpu/ops/silu/silu.h b/src/native/cpu/ops/silu/silu.h new file mode 100644 index 000000000..d3317656c --- /dev/null +++ b/src/native/cpu/ops/silu/silu.h @@ -0,0 +1,59 @@ +#ifndef INFINI_OPS_CPU_SILU_SILU_H_ +#define INFINI_OPS_CPU_SILU_SILU_H_ + +#include + +#include "base/silu.h" +#include "common/generic_utils.h" +#include "native/cpu/caster_.h" + +namespace infini::ops { + +template <> +class Operator : public Silu, + Caster { + public: + using Silu::Silu; + + void operator()(const Tensor input, Tensor out) const override { + DispatchFunc( + out_type_, + [&](auto tag) { + using T = typename decltype(tag)::type; + Compute(input, out); + }, + "Operator::operator()"); + } + + private: + template + void Compute(const Tensor input, Tensor out) const { + using ComputeType = std::conditional_t || + IsFP16, + float, T>; + + const auto* input_ptr = static_cast(input.data()); + auto* out_ptr = static_cast(out.data()); + + auto get_idx = [&](Tensor::Size i, bool is_contig, const auto* shape, + const auto* strides) { + return is_contig ? i : utils::IndexToOffset(i, ndim_, shape, strides); + }; + +#pragma omp parallel for + for (Tensor::Size i = 0; i < output_size_; ++i) { + const auto input_idx = get_idx( + i, is_input_contiguous_, input_shape_.data(), input_strides_.data()); + const auto out_idx = get_idx(i, is_out_contiguous_, out_shape_.data(), + out_strides_.data()); + const ComputeType input_value = Cast(input_ptr[input_idx]); + const ComputeType sigmoid = static_cast( + 1.0 / (1.0 + std::exp(-static_cast(input_value)))); + out_ptr[out_idx] = Cast(input_value * sigmoid); + } + } +}; + +} // namespace infini::ops + +#endif