chromium/third_party/xnnpack/src/src/xnnpack/packq.h

// Copyright 2024 Google LLC
//
// This source code is licensed under the BSD-style license found in the
// LICENSE file in the root directory of this source tree.

#ifndef __XNNPACK_SRC_XNNPACK_PACKQ_H
#define __XNNPACK_SRC_XNNPACK_PACKQ_H

#include <assert.h>
#include <stddef.h>
#include <stdint.h>

#include "xnnpack/common.h"
#include "xnnpack/config-types.h"
#include "xnnpack/config.h"
#include "xnnpack/math.h"

#ifdef __cplusplus
extern "C" {
#endif

// These functions have been adapted from KleidiAI's
// `kai_run_lhs_quant_pack_qai8dxp_f32` as a reference scalar implementation.

inline static size_t k_roundedup(size_t k, size_t kr, size_t sr) { … }

inline static size_t lhs_packed_stride(size_t k, size_t mr_packed, size_t kr,
                                       size_t sr) { … }

XNN_INLINE static size_t xnn_x8_packq_f32qp8_packed_offset(
    size_t m_idx, size_t k, size_t mr_packed, size_t kr, size_t sr) { … }

XNN_INLINE static size_t xnn_x8_packq_f32qp8_packed_size(size_t m, size_t k,
                                                         size_t mr_packed,
                                                         size_t kr, size_t sr) { … }

XNN_INLINE static size_t xnn_x8_packq_f32qp8_gemm_packed_size(size_t m,
                                                              size_t k) { … }

XNN_INLINE static int8_t xnn_x8_packq_f32qp8_get_quantized(
    size_t m_idx, size_t k_idx, const int8_t* lhs_packed, size_t k,
    size_t mr_packed, size_t kr, size_t sr) { … }

XNN_INLINE static float xnn_x8_packq_f32qp8_get_dequantized(
    size_t m_idx, size_t k_idx, const int8_t* lhs_packed, size_t k,
    size_t mr_packed, size_t kr, size_t sr) { … }

#define DECLARE_X8_PACKQ_UKERNEL_FUNCTION(fn_name) …

DECLARE_X8_PACKQ_UKERNEL_FUNCTION(xnn_x8_packq_f32qp8_ukernel__scalar_u1)

#if XNN_ENABLE_KLEIDIAI
DECLARE_X8_PACKQ_UKERNEL_FUNCTION(xnn_x8_packq_f32qp8_ukernel__aarch64_neon_u2)
#endif  // XNN_ENABLE_KLEIDIAI

#ifdef __cplusplus
}  // extern "C"
#endif

#endif  // __XNNPACK_SRC_XNNPACK_PACKQ_H