File size: 1,650 Bytes
57c2394 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | #pragma once
#include "packed_matmul_cpu.h"
#include <cstdint>
namespace orbitquant::cpu {
enum class ActivationIsa : std::uint8_t {
Portable,
Neon,
Avx2,
Avx512,
};
struct ActivationArgs {
void *out;
void const *x;
// Exactly one of the permutation pointers is set (int64 checkpoints and the
// int32 hot-path constants are both accepted).
std::int64_t const *permutation;
std::int32_t const *permutation_i32;
std::int8_t const *signs;
float const *centroids;
float const *boundaries;
ScalarKind scalar_kind;
ActivationIsa isa;
std::int64_t rows;
std::int64_t dim;
std::int64_t boundary_count;
std::int64_t block_size;
float eps;
float inv_sqrt_block;
};
void activation_fwht_msvc_avx2(float *values, std::int64_t block_size);
float activation_squared_norm_msvc_avx2(
void const *data,
ScalarKind scalar_kind,
std::int64_t offset,
std::int64_t dim);
void activation_quantize_lookup_msvc_avx2(
ActivationArgs const &args,
float const *scratch,
std::int64_t output_offset,
float norm);
struct AdalnArgs {
void *out;
void const *x;
std::uint8_t const *packed_weight;
float const *scales;
float const *bias;
bool has_bias;
std::int64_t rows;
std::int64_t out_features;
std::int64_t in_features;
std::int64_t group_size;
std::int64_t num_groups;
std::int64_t padded_in_features;
};
using AdalnRangeFn = void (*)(
AdalnArgs const &args,
std::int64_t out_start,
std::int64_t out_end);
void packed_adaln_msvc_avx2_range(
AdalnArgs const &args,
std::int64_t out_start,
std::int64_t out_end);
} // namespace orbitquant::cpu
|