CS348Project/llama.cpp
5
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2025-10-27 08:21:30 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
342c728d031d50673feded797520a44127d73379
llama.cpp/ggml/src/ggml-cpu/spacemit/ime.cpp
muggle-stack 342c728d03 ggml : fix SpaceMit IME array out-of-bounds in task assignment (#16629) 
Fix incorrect task-to-batch index calculation in the quantization phase.

The bug caused out-of-bounds access to qnbitgemm_args array when
compute_idx exceeded per_gemm_block_count_m, leading to invalid
pointer dereferences and SIGBUS errors.

Correctly map tasks to batches by dividing compute_idx by
per_gemm_block_count_m instead of block_size_m.

Example:
  batch_feature=1, gemm_m=30, block_size_m=4
  per_gemm_block_count_m = 8, task_count = 8

  Old: gemm_idx = 4/4 = 1 (out of bounds  New: gemm_idx = 4/8 = 0 (correct)

Tested on SpaceMit K1 RISC-V64 with qwen2.5:0.5b model.

Co-authored-by: muggle <mingjun.rong@spacemit.com>
2025-10-17 13:01:23 +03:00

1026 lines
42 KiB
C++
Raw Blame History

#define GGML_COMMON_IMPL_CPP
#define GGML_COMMON_DECL_CPP
#include "ime.h"
#include "ggml-backend-impl.h"
#include "ggml-common.h"
#include "ggml-cpu.h"
#include "ime_kernels.h"
#include "traits.h"
#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdio> // for GGML_ASSERT
#include <stdexcept>
#include <thread>
// clang-format off
#if defined(__riscv)
#if !defined(__riscv_v) || !defined(__riscv_v_intrinsic)
#error "riscv v extension or v_intrinsic not enabled"
#else
#include <riscv_vector.h>
#endif
#if !defined(__riscv_zfh)
#error "riscv zfh extension not enabled"
#endif
#if defined(RISCV64_SPACEMIT_IME1)
#else
#error "RISCV64_SPACEMIT_IME1 not defined"
#endif
#else
#error "riscv not enabled in this build"
#endif
#if defined(__GNUC__)
#pragma GCC diagnostic ignored "-Woverlength-strings"
#pragma GCC diagnostic ignored "-Wcast-qual"
#pragma GCC diagnostic ignored "-Wunused-parameter"
#endif
#if defined(RISCV64_SPACEMIT_IME1)
#define QGEMM_STRIDEN_THREAD_ALIGN 16
#else
#define QGEMM_STRIDEN_THREAD_ALIGN 32
#endif
// clang-format on
struct qnbitgemm_spacemit_ime_args {
const float * a_ptr = nullptr;
size_t lda = 0;
const std::byte * packed_quant_b_data = nullptr;
const float * quant_b_scale = nullptr;
const void * quant_b_zp = nullptr;
const float * quant_b_blksum = nullptr;
const float * bias = nullptr;
float * c_ptr = nullptr;
size_t ldc = 0;
};
constexpr size_t div_round_up(size_t up, size_t down) {
return (up + down - 1) / down;
}
constexpr size_t q8_blk_size(size_t blk_len) {
const size_t blk_size = sizeof(float) + blk_len * sizeof(int8_t);
// Currently, the strictest alignment requirement of a block is for a float.
// Ensure contiguous blocks are suitably aligned.
assert(blk_size % alignof(float) == 0);
return blk_size;
}
namespace ggml::cpu::riscv64_spacemit {
const int num_ai_cores = std::thread::hardware_concurrency() / 2;
} // namespace ggml::cpu::riscv64_spacemit
static void sqnbitgemm_spacemit_ime_i8i4(const size_t blk_len,
const size_t gemm_k,
const qnbitgemm_spacemit_ime_args * gemm_args,
void * const per_gemm_ws,
const size_t m_start,
const size_t m_count,
const size_t n_start,
const size_t n_count) {
constexpr size_t scale_stride = sizeof(uint16_t);
constexpr size_t blk_bitwidth = 4;
const size_t k_blks = div_round_up(gemm_k, blk_len);
const size_t lda = k_blks * q8_blk_size(blk_len);
const size_t ldc = gemm_args->ldc;
const size_t ldb = k_blks * (blk_len * blk_bitwidth / 8);
const std::byte * quant_a_ptr = static_cast<const std::byte *>(per_gemm_ws) + m_start * lda;
const size_t zero_point_stride = gemm_args->quant_b_zp != nullptr ? sizeof(uint8_t) : 0;
const size_t packed_b_stride = ldb + k_blks * (scale_stride + zero_point_stride);
const std::byte * packed_quant_b_data = gemm_args->packed_quant_b_data + n_start * packed_b_stride;
float * c_ptr = gemm_args->c_ptr + m_start * ldc + n_start;
size_t count_n = 0;
const size_t compute_block_count_n = m_count == 1 ? n_count : 16;
for (size_t n = 0; n < n_count; n += count_n) {
count_n = std::min(n_count - n, compute_block_count_n);
const std::byte * a_row = quant_a_ptr;
const std::byte * b_col = packed_quant_b_data + n * packed_b_stride;
const std::byte * b_col_zp = (zero_point_stride != 0) ? b_col : nullptr;
float * c_blk = c_ptr + n;
int32_t rows_remaining = m_count;
while (rows_remaining > 0) {
const auto rows_handled = sqnbitgemm_spacemit_ime::ime1::gemm_kernel_i8i4(
blk_len, a_row, b_col, nullptr, b_col_zp, c_blk, rows_remaining, count_n, gemm_k, k_blks, ldc, nullptr,
scale_stride);
c_blk += rows_handled * ldc;
a_row += rows_handled * lda;
rows_remaining -= rows_handled;
}
}
}
template <int K> constexpr int QK_0() {
if constexpr (K == 4) {
return QK4_0;
}
if constexpr (K == 8) {
return QK8_0;
}
return -1;
}
template <int K, int N> struct block {
ggml_half d[N]; // deltas for N qK_0 blocks
uint8_t qs[(QK_0<K>() * N * K) / 8]; // quants for N qK_0 blocks
};
template <int K, int N> struct block_with_zp {
ggml_half d[N]; // deltas for N qK_1 blocks
uint8_t zp[N]; // zero points for N qK_1 blocks
uint8_t qs[(QK_0<K>() * N * K) / 8]; // quants for N qK_1 blocks
};
// control size
static_assert(sizeof(block<4, 16>) == 16 * sizeof(ggml_half) + QK4_0 * 8, "wrong block<4,16> size/padding");
static_assert(sizeof(block_with_zp<4, 16>) == 16 * sizeof(ggml_half) + QK4_0 * 8 + 16 * sizeof(uint8_t),
"wrong block_with_zp<4,16> size/padding");
static_assert(sizeof(block<8, 16>) == 16 * sizeof(ggml_half) + QK4_0 * 16, "wrong block<8,16> size/padding");
using block_q4_0x16 = block<4, 16>;
using block_q4_1x16 = block_with_zp<4, 16>;
using block_q8_0x16 = block<8, 16>;
static block_q4_0x16 make_block_q4_0x16(block_q4_0 * in, unsigned int blck_size_interleave) {
block_q4_0x16 out;
GGML_ASSERT(QK4_0 / blck_size_interleave == 2);
for (int i = 0; i < 16; i++) {
out.d[i] = in[i].d;
}
for (int i = 0; i < 16; i++) {
// [0, 15], in.d & 0x0F
for (int j = 0; j < QK4_0 / 4; j++) {
//src [b0 b16] ......... [b8 b24] ......... [b15 b31]
//dst [b0 b8] ......... [b7 b15]
out.qs[i * QK4_0 / 4 + j] = (in[i].qs[j] & 0x0F) | ((in[i].qs[j + QK4_0 / 4] & 0x0F) << 4);
}
}
for (int i = 0; i < 16; i++) {
// [16, 31], in.d & 0xF0
for (int j = 0; j < QK4_0 / 4; j++) {
//src [b0 b16] ......... [b8 b24] ......... [b15 b31]
//dst [b16 b24] ......... [b23 b31]
out.qs[4 * QK4_0 + i * QK4_0 / 4 + j] = ((in[i].qs[j] & 0xF0) >> 4) | (in[i].qs[j + QK4_0 / 4] & 0xF0);
}
}
return out;
}
static block_q4_1x16 make_block_q4_1x16(block_q4_1 * in, unsigned int blck_size_interleave) {
block_q4_1x16 out;
GGML_ASSERT(QK4_1 / blck_size_interleave == 2);
for (int i = 0; i < 16; i++) {
float d = GGML_FP16_TO_FP32(in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.d);
float m = GGML_FP16_TO_FP32(in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.m);
float mid = -std::nearbyintf(m / d);
mid = std::min(15.0f, std::max(0.0f, mid));
out.d[i] = GGML_FP32_TO_FP16(d);
out.zp[i] = static_cast<uint8_t>(mid);
}
for (int i = 0; i < 16; i++) {
// [0, 15], in.d & 0x0F
for (int j = 0; j < QK4_1 / 4; j++) {
//src [b0 b16] ......... [b8 b24] ......... [b15 b31]
//dst [b0 b8] ......... [b7 b15]
out.qs[i * QK4_1 / 4 + j] = (in[i].qs[j] & 0x0F) | ((in[i].qs[j + QK4_1 / 4] & 0x0F) << 4);
}
}
for (int i = 0; i < 16; i++) {
// [16, 31], in.d & 0xF0
for (int j = 0; j < QK4_1 / 4; j++) {
//src [b0 b16] ......... [b8 b24] ......... [b15 b31]
//dst [b16 b24] ......... [b23 b31]
out.qs[4 * QK4_1 + i * QK4_1 / 4 + j] = ((in[i].qs[j] & 0xF0) >> 4) | (in[i].qs[j + QK4_1 / 4] & 0xF0);
}
}
return out;
}
static int repack_q4_0_to_q4_0_16_bl(struct ggml_tensor * t,
int interleave_block,
const void * GGML_RESTRICT data,
size_t data_size) {
GGML_ASSERT(t->type == GGML_TYPE_Q4_0);
GGML_ASSERT(interleave_block == 16);
constexpr int nrows_interleaved = 16;
block_q4_0x16 * dst = (block_q4_0x16 *) t->data;
const block_q4_0 * src = (const block_q4_0 *) data;
block_q4_0 dst_tmp[16];
int nrow = ggml_nrows(t);
int nblocks = t->ne[0] / QK4_0;
GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q4_0));
if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % QK4_0 != 0) {
return -1;
}
for (int b = 0; b < nrow; b += nrows_interleaved) {
for (int64_t x = 0; x < nblocks; x++) {
for (int i = 0; i < nrows_interleaved; i++) {
dst_tmp[i] = src[x + i * nblocks];
}
*dst++ = make_block_q4_0x16(dst_tmp, interleave_block);
}
src += nrows_interleaved * nblocks;
}
return 0;
GGML_UNUSED(data_size);
}
static int repack_q4_1_to_q4_1_16_bl(struct ggml_tensor * t,
int interleave_block,
const void * GGML_RESTRICT data,
size_t data_size) {
GGML_ASSERT(t->type == GGML_TYPE_Q4_1);
GGML_ASSERT(interleave_block == 16);
constexpr int nrows_interleaved = 16;
block_q4_1x16 * dst = (block_q4_1x16 *) t->data;
const block_q4_1 * src = (const block_q4_1 *) data;
block_q4_1 dst_tmp[16];
int nrow = ggml_nrows(t);
int nblocks = t->ne[0] / QK4_1;
GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q4_1));
if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % QK4_1 != 0) {
return -1;
}
for (int b = 0; b < nrow; b += nrows_interleaved) {
for (int64_t x = 0; x < nblocks; x++) {
for (int i = 0; i < nrows_interleaved; i++) {
dst_tmp[i] = src[x + i * nblocks];
}
*dst++ = make_block_q4_1x16(dst_tmp, interleave_block);
}
src += nrows_interleaved * nblocks;
}
return 0;
GGML_UNUSED(data_size);
}
static inline void get_scale_min_k4(int j,
const uint8_t * GGML_RESTRICT q,
uint8_t * GGML_RESTRICT d,
uint8_t * GGML_RESTRICT m) {
if (j < 4) {
*d = q[j] & 63;
*m = q[j + 4] & 63;
} else {
*d = (q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4);
*m = (q[j + 4] >> 4) | ((q[j - 0] >> 6) << 4);
}
}
static int repack_q4_k_to_q4_1_16_bl(struct ggml_tensor * t,
int interleave_block,
const void * GGML_RESTRICT data,
size_t data_size) {
GGML_ASSERT(t->type == GGML_TYPE_Q4_K);
GGML_ASSERT(interleave_block == 16);
GGML_ASSERT(QK_K / QK4_1 == 8);
constexpr int nrows_interleaved = 16;
block_q4_1x16 * dst = (block_q4_1x16 *) t->data;
const block_q4_K * src = (const block_q4_K *) data;
block_q4_1 dst_tmp[16];
int nrow = ggml_nrows(t);
int nblocks = t->ne[0] / QK_K;
if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % QK_K != 0) {
return -1;
}
for (int b = 0; b < nrow; b += nrows_interleaved) {
for (int64_t x = 0; x < nblocks; x++) {
for (int j = 0; j < 8; j++) {
for (int i = 0; i < nrows_interleaved; i++) {
uint8_t sc, m;
const float d = GGML_FP16_TO_FP32(src[x + i * nblocks].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.d);
const float min =
GGML_FP16_TO_FP32(src[x + i * nblocks].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.dmin);
get_scale_min_k4(j, src[x + i * nblocks].scales, &sc, &m);
const float d1 = d * sc;
const float m1 = min * m;
dst_tmp[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.d = GGML_FP32_TO_FP16(d1);
dst_tmp[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.m = GGML_FP32_TO_FP16(-m1);
// src -> [b0, b32] [b1, b33] ... [b31, b63]
// dst -> [b0, b16] [b1, b17] ... [b15, b31] [b32, b48] [b33, b49] ... [b47, b63]
const uint8_t * q = src[x + i * nblocks].qs + (j / 2) * QK4_1;
if (j % 2 == 0) {
for (int ii = 0; ii < 16; ii++) {
dst_tmp[i].qs[ii] = (q[ii] & 0x0F) | ((q[ii + 16] & 0x0F) << 4);
}
} else {
for (int ii = 0; ii < 16; ii++) {
dst_tmp[i].qs[ii] = ((q[ii] & 0xF0) >> 4) | (q[ii + 16] & 0xF0);
}
}
}
*dst++ = make_block_q4_1x16(dst_tmp, interleave_block);
}
}
src += nrows_interleaved * nblocks;
}
return 0;
GGML_UNUSED(data_size);
}
namespace ggml::cpu::riscv64_spacemit {
template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS>
int repack(struct ggml_tensor *, const void *, size_t);
template <> int repack<block_q4_0, 8, 16>(struct ggml_tensor * t, const void * data, size_t data_size) {
return repack_q4_0_to_q4_0_16_bl(t, 16, data, data_size);
}
template <> int repack<block_q4_1, 8, 16>(struct ggml_tensor * t, const void * data, size_t data_size) {
return repack_q4_1_to_q4_1_16_bl(t, 16, data, data_size);
}
template <> int repack<block_q4_K, 8, 16>(struct ggml_tensor * t, const void * data, size_t data_size) {
return repack_q4_k_to_q4_1_16_bl(t, 16, data, data_size);
}
class tensor_traits_base : public ggml::cpu::tensor_traits {
public:
virtual int repack(struct ggml_tensor * t, const void * data, size_t data_size) = 0;
};
template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS> class tensor_traits : public tensor_traits_base {
bool work_size(int /* n_threads */, const struct ggml_tensor * op, size_t & size) override {
switch (op->op) {
case GGML_OP_MUL_MAT:
size = ggml_row_size(GGML_TYPE_Q8_0, ggml_nelements(op->src[1])) * 4;
size = ((size + QK4_0 - 1) / QK4_0) * (QK4_0 * sizeof(float) + sizeof(float));
return true;
default:
// GGML_ABORT("fatal error");
break;
}
return false;
}
bool compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) override {
switch (op->op) {
case GGML_OP_MUL_MAT:
if (op->src[0]->type == GGML_TYPE_Q4_0 || //
op->src[0]->type == GGML_TYPE_Q4_1 || //
op->src[0]->type == GGML_TYPE_Q4_K) {
forward_mul_mat_q4(params, op);
return true;
}
default:
// GGML_ABORT("fatal error");
break;
}
return false;
}
void forward_mul_mat_q4(ggml_compute_params * params, ggml_tensor * op) {
const ggml_tensor * src0 = op->src[0];
const ggml_tensor * src1 = op->src[1];
ggml_tensor * dst = op;
GGML_TENSOR_BINARY_OP_LOCALS
int ith = params->ith;
int nth = params->nth;
[[maybe_unused]] const enum ggml_type type = src0->type;
void * w_data = (void *) src0->data;
const float * feature = (const float *) src1->data;
float * output = (float *) dst->data;
const size_t batch_feature = ne12 * ne13;
[[maybe_unused]] const size_t batch_weight = ne02 * ne03;
const size_t gemm_m = ne11;
const size_t gemm_k = ne10;
const size_t gemm_n = ne01;
GGML_ASSERT(batch_weight == 1);
const size_t block_count_k = div_round_up(gemm_k, QK4_0);
const size_t per_gemm_workspace_size = gemm_m * block_count_k * q8_blk_size(QK4_0);
const size_t per_gemm_workspace_stride =
div_round_up(per_gemm_workspace_size, alignof(uint64_t)) * alignof(uint64_t);
const size_t gemm_workspace_size = batch_feature * per_gemm_workspace_stride;
const size_t desired_wsize = gemm_workspace_size + alignof(uint64_t) - 1;
if (ith == 0 && params->wsize < desired_wsize) {
throw std::runtime_error("wsize less than desired_wsize");
}
std::vector<qnbitgemm_spacemit_ime_args> qnbitgemm_args(batch_feature);
for (size_t i = 0; i < batch_feature; i++) {
qnbitgemm_args[i].a_ptr = feature + gemm_m * gemm_k * i;
qnbitgemm_args[i].lda = gemm_k;
qnbitgemm_args[i].packed_quant_b_data = (const std::byte *) w_data;
qnbitgemm_args[i].quant_b_scale = nullptr;
if constexpr (std::is_same_v<BLOC_TYPE, block_q4_0>) {
qnbitgemm_args[i].quant_b_zp = nullptr;
} else {
qnbitgemm_args[i].quant_b_zp = w_data;
}
qnbitgemm_args[i].bias = nullptr;
qnbitgemm_args[i].c_ptr = output + gemm_m * gemm_n * i;
qnbitgemm_args[i].ldc = gemm_n;
}
const uintptr_t ws_ptr = reinterpret_cast<uintptr_t>(params->wdata);
void * ws = reinterpret_cast<void *>((ws_ptr + alignof(uint64_t) - 1) & (~(alignof(uint64_t) - 1)));
const size_t quant_a_stride = block_count_k * q8_blk_size(QK4_0);
{
constexpr size_t block_size_m = 4;
size_t per_gemm_block_count_m = div_round_up(gemm_m, block_size_m);
int32_t task_count = batch_feature * per_gemm_block_count_m;
int32_t task_per_thread = (task_count + nth - 1) / nth;
int32_t start = ith * task_per_thread;
int32_t end = std::min((ith + 1) * task_per_thread, task_count);
for (int32_t compute_idx = start; compute_idx < end; compute_idx++) {
int32_t gemm_idx = compute_idx / per_gemm_block_count_m;
int32_t block_idx_in_gemm = compute_idx % per_gemm_block_count_m;
int32_t m_idx = block_idx_in_gemm * block_size_m;
const qnbitgemm_spacemit_ime_args & data = qnbitgemm_args[gemm_idx];
int32_t rows_tobe_handled = (gemm_m - m_idx) > block_size_m ? block_size_m : (gemm_m - m_idx);
if (rows_tobe_handled == block_size_m) {
const float * a_row_ptr = data.a_ptr + m_idx * data.lda;
std::byte * quant_a_row_ptr =
static_cast<std::byte *>(ws) + gemm_idx * per_gemm_workspace_stride + m_idx * quant_a_stride;
sqnbitgemm_spacemit_ime::ime1::quantize_a_4row_i8(QK4_0, a_row_ptr, gemm_k, quant_a_row_ptr);
} else {
while (rows_tobe_handled) {
const float * a_row_ptr = data.a_ptr + m_idx * data.lda;
std::byte * quant_a_row_ptr = static_cast<std::byte *>(ws) +
gemm_idx * per_gemm_workspace_stride + m_idx * quant_a_stride;
sqnbitgemm_spacemit_ime::ime1::quantize_a_row_i8(QK4_0, a_row_ptr, gemm_k, quant_a_row_ptr);
rows_tobe_handled -= 1;
m_idx += 1;
}
}
}
}
ggml_barrier(params->threadpool);
if (ith >= ggml::cpu::riscv64_spacemit::num_ai_cores) {
return;
}
nth = std::min(nth, int{ ggml::cpu::riscv64_spacemit::num_ai_cores });
size_t threads_per_gemm = nth / batch_feature;
constexpr size_t gemm_m_stride = 128;
size_t nc = gemm_n;
const size_t gemm_m_blocked = div_round_up(gemm_m, gemm_m_stride);
const size_t max_nc = div_round_up(gemm_n * gemm_m_blocked, threads_per_gemm);
if (max_nc < nc) {
nc = std::min(nc, div_round_up(max_nc, QGEMM_STRIDEN_THREAD_ALIGN) * QGEMM_STRIDEN_THREAD_ALIGN);
}
const size_t gemm_n_stride = nc;
const size_t thread_count_m = div_round_up(gemm_m, gemm_m_stride);
const size_t thread_count_n = div_round_up(gemm_n, gemm_n_stride);
threads_per_gemm = thread_count_m * thread_count_n;
{
int task_count = batch_feature * threads_per_gemm;
int task_per_thread = (task_count + nth - 1) / nth;
int start = ith * task_per_thread;
int end = std::min((ith + 1) * task_per_thread, task_count);
for (int compute_idx = start; compute_idx < end; compute_idx++) {
const auto gemm_i = compute_idx / threads_per_gemm;
const auto blk_i = compute_idx % threads_per_gemm;
const auto * data = &qnbitgemm_args[gemm_i];
const auto tid_n = blk_i / thread_count_m;
const auto tid_m = blk_i % thread_count_m;
const size_t m_start = tid_m * gemm_m_stride;
const size_t m_count = std::min(gemm_m - m_start, (size_t) gemm_m_stride);
const size_t n_start = tid_n * gemm_n_stride;
const size_t n_count = std::min(gemm_n - n_start, (size_t) gemm_n_stride);
void * per_gemm_ws = reinterpret_cast<std::byte *>(ws) + gemm_i * per_gemm_workspace_stride;
sqnbitgemm_spacemit_ime_i8i4(QK4_0, gemm_k, data, per_gemm_ws, m_start, m_count, n_start, n_count);
}
}
}
int repack(struct ggml_tensor * t, const void * data, size_t data_size) override {
GGML_LOG_DEBUG("%s: repack tensor %s with %s_%dx%d\n", __func__, t->name, ggml_type_name(t->type),
(int) NB_COLS, (int) INTER_SIZE);
return ggml::cpu::riscv64_spacemit::repack<BLOC_TYPE, INTER_SIZE, NB_COLS>(t, data, data_size);
}
};
class tensor_traits_common : public tensor_traits_base {
bool work_size(int /* n_threads */, const struct ggml_tensor * op, size_t & size) override {
switch (op->op) {
case GGML_OP_NORM:
case GGML_OP_RMS_NORM:
size = 0;
return true;
default:
// GGML_ABORT("fatal error");
break;
}
return false;
}
bool compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) override {
switch (op->op) {
case GGML_OP_NORM:
forward_norm_f32(params, op);
return true;
case GGML_OP_RMS_NORM:
forward_rms_norm_f32(params, op);
return true;
default:
// GGML_ABORT("fatal error");
break;
}
return false;
}
void forward_norm_f32(ggml_compute_params * params, ggml_tensor * op) {
const ggml_tensor * src0 = op->src[0];
ggml_tensor * dst = op;
GGML_ASSERT(ggml_are_same_shape(src0, dst));
GGML_ASSERT(src0->nb[0] == sizeof(float));
const int ith = params->ith;
const int nth = params->nth;
GGML_TENSOR_UNARY_OP_LOCALS
float epsilon;
memcpy(&epsilon, dst->op_params, sizeof(float));
GGML_ASSERT(epsilon > 0.0f);
auto * input = (float *) src0->data;
auto * output = (float *) dst->data;
const auto hidden_size = ne00;
const auto task_count = ne01 * ne02 * ne03;
const auto task_per_thread = (task_count + nth - 1) / nth;
const auto task_begin = ith * task_per_thread;
const auto task_end = std::min((ith + 1) * task_per_thread, task_count);
for (auto task_idx = task_begin; task_idx < task_end; task_idx++) {
auto offset = task_idx * hidden_size;
auto * p_input = const_cast<float *>(input + offset);
auto * p_output = output + offset;
auto * p_temp_output = p_output;
auto * p_gamma_data = (const float *) nullptr;
auto * p_beta_data = (const float *) nullptr;
size_t gvl = __riscv_vsetvlmax_e32m4();
vfloat32m4_t sum = __riscv_vfmv_v_f_f32m4(0.f, gvl);
vfloat32m4_t sum_sq = __riscv_vfmv_v_f_f32m4(0.f, gvl);
int64_t length = hidden_size;
while (length > 0) {
gvl = __riscv_vsetvl_e32m4(length);
// load data
vfloat32m4_t src_data = __riscv_vle32_v_f32m4(p_input, gvl);
sum = __riscv_vfadd_vv_f32m4(sum, src_data, gvl);
sum_sq = __riscv_vfmacc_vv_f32m4(sum_sq, src_data, src_data, gvl);
__riscv_vse32_v_f32m4(p_temp_output, src_data, gvl);
p_input += gvl;
p_temp_output += gvl;
length -= gvl;
}
gvl = __riscv_vsetvlmax_e32m1();
float mean = 0.f;
vfloat32m1_t zero_v = __riscv_vfmv_v_f_f32m1(0.f, gvl);
vfloat32m1_t mean_v =
__riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m4_f32m1(sum, 0), __riscv_vget_v_f32m4_f32m1(sum, 1), gvl);
mean_v = __riscv_vfadd_vv_f32m1(mean_v, __riscv_vget_v_f32m4_f32m1(sum, 2), gvl);
mean_v = __riscv_vfadd_vv_f32m1(mean_v, __riscv_vget_v_f32m4_f32m1(sum, 3), gvl);
mean_v = __riscv_vfredusum_vs_f32m1_f32m1(mean_v, zero_v, gvl);
mean = __riscv_vfmv_f_s_f32m1_f32(mean_v);
mean /= hidden_size;
vfloat32m1_t mean_square_v = __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m4_f32m1(sum_sq, 0),
__riscv_vget_v_f32m4_f32m1(sum_sq, 1), gvl);
mean_square_v = __riscv_vfadd_vv_f32m1(mean_square_v, __riscv_vget_v_f32m4_f32m1(sum_sq, 2), gvl);
mean_square_v = __riscv_vfadd_vv_f32m1(mean_square_v, __riscv_vget_v_f32m4_f32m1(sum_sq, 3), gvl);
mean_square_v = __riscv_vfredusum_vs_f32m1_f32m1(mean_square_v, zero_v, gvl);
float mean_square = __riscv_vfmv_f_s_f32m1_f32(mean_square_v);
mean_square /= hidden_size;
mean_square = sqrt(mean_square - mean * mean + epsilon);
mean_square = 1.0f / mean_square;
length = hidden_size;
p_temp_output = p_output;
if (p_gamma_data == nullptr && p_beta_data == nullptr) {
while (length > 0) {
gvl = __riscv_vsetvl_e32m4(length);
vfloat32m4_t src_data = __riscv_vle32_v_f32m4(p_temp_output, gvl);
src_data = __riscv_vfsub_vf_f32m4(src_data, mean, gvl);
src_data = __riscv_vfmul_vf_f32m4(src_data, mean_square, gvl);
__riscv_vse32_v_f32m4(p_output, src_data, gvl);
p_temp_output += gvl;
p_output += gvl;
length -= gvl;
}
} else if (p_beta_data == nullptr) {
while (length > 0) {
gvl = __riscv_vsetvl_e32m4(length);
vfloat32m4_t src_data = __riscv_vle32_v_f32m4(p_temp_output, gvl);
vfloat32m4_t gamma_data_v = __riscv_vle32_v_f32m4(p_gamma_data, gvl);
src_data = __riscv_vfsub_vf_f32m4(src_data, mean, gvl);
src_data = __riscv_vfmul_vf_f32m4(src_data, mean_square, gvl);
src_data = __riscv_vfmul_vv_f32m4(src_data, gamma_data_v, gvl);
__riscv_vse32_v_f32m4(p_output, src_data, gvl);
p_temp_output += gvl;
p_output += gvl;
p_gamma_data += gvl;
length -= gvl;
}
} else if (p_gamma_data != nullptr) {
while (length > 0) {
gvl = __riscv_vsetvl_e32m4(length);
vfloat32m4_t src_data = __riscv_vle32_v_f32m4(p_temp_output, gvl);
vfloat32m4_t gamma_data_v = __riscv_vle32_v_f32m4(p_gamma_data, gvl);
src_data = __riscv_vfsub_vf_f32m4(src_data, mean, gvl);
src_data = __riscv_vfmul_vf_f32m4(src_data, mean_square, gvl);
src_data = __riscv_vfmul_vv_f32m4(src_data, gamma_data_v, gvl);
vfloat32m4_t beta_data_v = __riscv_vle32_v_f32m4(p_beta_data, gvl);
src_data = __riscv_vfadd_vv_f32m4(src_data, beta_data_v, gvl);
p_beta_data += gvl;
__riscv_vse32_v_f32m4(p_output, src_data, gvl);
p_temp_output += gvl;
p_output += gvl;
p_gamma_data += gvl;
length -= gvl;
}
}
}
}
void forward_rms_norm_f32(ggml_compute_params * params, ggml_tensor * op) {
const ggml_tensor * src0 = op->src[0];
ggml_tensor * dst = op;
GGML_ASSERT(ggml_are_same_shape(src0, dst));
GGML_ASSERT(src0->nb[0] == sizeof(float));
const int ith = params->ith;
const int nth = params->nth;
GGML_TENSOR_UNARY_OP_LOCALS
float epsilon;
memcpy(&epsilon, dst->op_params, sizeof(float));
GGML_ASSERT(epsilon > 0.0f);
auto * input = (float *) src0->data;
auto * output = (float *) dst->data;
const auto hidden_size = ne00;
const auto task_count = ne01 * ne02 * ne03;
const auto task_per_thread = (task_count + nth - 1) / nth;
const auto task_begin = ith * task_per_thread;
const auto task_end = std::min((ith + 1) * task_per_thread, task_count);
for (auto task_idx = task_begin; task_idx < task_end; task_idx++) {
auto offset = task_idx * hidden_size;
auto * p_input = const_cast<float *>(input + offset);
auto * p_output = output + offset;
auto * p_temp_output = p_output;
auto * p_gamma_data = (const float *) nullptr;
auto * p_beta_data = (const float *) nullptr;
size_t gvl = __riscv_vsetvlmax_e32m4();
// vfloat32m4_t sum = __riscv_vfmv_v_f_f32m4(0.f, gvl);
vfloat32m4_t sum_sq = __riscv_vfmv_v_f_f32m4(0.f, gvl);
int64_t length = hidden_size;
while (length > 0) {
gvl = __riscv_vsetvl_e32m4(length);
// load data
vfloat32m4_t src_data = __riscv_vle32_v_f32m4(p_input, gvl);
sum_sq = __riscv_vfmacc_vv_f32m4(sum_sq, src_data, src_data, gvl);
__riscv_vse32_v_f32m4(p_temp_output, src_data, gvl);
p_input += gvl;
p_temp_output += gvl;
length -= gvl;
}
gvl = __riscv_vsetvlmax_e32m1();
// float mean = 0.f;
vfloat32m1_t zero_v = __riscv_vfmv_v_f_f32m1(0.f, gvl);
vfloat32m1_t mean_square_v = __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m4_f32m1(sum_sq, 0),
__riscv_vget_v_f32m4_f32m1(sum_sq, 1), gvl);
mean_square_v = __riscv_vfadd_vv_f32m1(mean_square_v, __riscv_vget_v_f32m4_f32m1(sum_sq, 2), gvl);
mean_square_v = __riscv_vfadd_vv_f32m1(mean_square_v, __riscv_vget_v_f32m4_f32m1(sum_sq, 3), gvl);
mean_square_v = __riscv_vfredusum_vs_f32m1_f32m1(mean_square_v, zero_v, gvl);
float mean_square = __riscv_vfmv_f_s_f32m1_f32(mean_square_v);
mean_square /= hidden_size;
mean_square = sqrt(mean_square + epsilon);
mean_square = 1.0f / mean_square;
length = hidden_size;
p_temp_output = p_output;
if (p_gamma_data == nullptr && p_beta_data == nullptr) {
while (length > 0) {
gvl = __riscv_vsetvl_e32m4(length);
vfloat32m4_t src_data = __riscv_vle32_v_f32m4(p_temp_output, gvl);
src_data = __riscv_vfmul_vf_f32m4(src_data, mean_square, gvl);
__riscv_vse32_v_f32m4(p_output, src_data, gvl);
p_temp_output += gvl;
p_output += gvl;
length -= gvl;
}
} else if (p_beta_data == nullptr) {
while (length > 0) {
gvl = __riscv_vsetvl_e32m4(length);
vfloat32m4_t src_data = __riscv_vle32_v_f32m4(p_temp_output, gvl);
vfloat32m4_t gamma_data_v = __riscv_vle32_v_f32m4(p_gamma_data, gvl);
src_data = __riscv_vfmul_vf_f32m4(src_data, mean_square, gvl);
src_data = __riscv_vfmul_vv_f32m4(src_data, gamma_data_v, gvl);
__riscv_vse32_v_f32m4(p_output, src_data, gvl);
p_temp_output += gvl;
p_output += gvl;
p_gamma_data += gvl;
length -= gvl;
}
} else if (p_gamma_data != nullptr) {
while (length > 0) {
gvl = __riscv_vsetvl_e32m4(length);
vfloat32m4_t src_data = __riscv_vle32_v_f32m4(p_temp_output, gvl);
vfloat32m4_t gamma_data_v = __riscv_vle32_v_f32m4(p_gamma_data, gvl);
src_data = __riscv_vfmul_vf_f32m4(src_data, mean_square, gvl);
src_data = __riscv_vfmul_vv_f32m4(src_data, gamma_data_v, gvl);
vfloat32m4_t beta_data_v = __riscv_vle32_v_f32m4(p_beta_data, gvl);
src_data = __riscv_vfadd_vv_f32m4(src_data, beta_data_v, gvl);
p_beta_data += gvl;
__riscv_vse32_v_f32m4(p_output, src_data, gvl);
p_temp_output += gvl;
p_output += gvl;
p_gamma_data += gvl;
length -= gvl;
}
}
}
}
int repack(struct ggml_tensor * t, const void * data, size_t data_size) override {
memcpy(t->data, data, data_size);
return 0;
}
};
static const tensor_traits<block_q4_0, 8, 16> q4_0_16x8_q8_0;
static const tensor_traits<block_q4_1, 8, 16> q4_1_16x8_q8_0;
static const tensor_traits<block_q4_K, 8, 16> q4_k_16x8_q8_0;
static const tensor_traits_common rvv_impl;
} // namespace ggml::cpu::riscv64_spacemit
static const ggml::cpu::tensor_traits * ggml_riscv64_spacemit_get_optimal_repack_type(const struct ggml_tensor * cur) {
if (cur->type == GGML_TYPE_Q4_0) {
if (cur->ne[1] % 16 == 0) {
return &ggml::cpu::riscv64_spacemit::q4_0_16x8_q8_0;
}
} else if (cur->type == GGML_TYPE_Q4_1) {
if (cur->ne[1] % 16 == 0) {
return &ggml::cpu::riscv64_spacemit::q4_1_16x8_q8_0;
}
} else if (cur->type == GGML_TYPE_Q4_K) {
if (cur->ne[1] % 16 == 0) {
return &ggml::cpu::riscv64_spacemit::q4_k_16x8_q8_0;
}
} else if (cur->type == GGML_TYPE_F32) {
return &ggml::cpu::riscv64_spacemit::rvv_impl;
}
return nullptr;
}
static enum ggml_status ggml_backend_riscv64_spacemit_buffer_init_tensor(ggml_backend_buffer_t buffer,
struct ggml_tensor * tensor) {
tensor->extra =
(void *) const_cast<ggml::cpu::tensor_traits *>(ggml_riscv64_spacemit_get_optimal_repack_type(tensor));
GGML_UNUSED(buffer);
return GGML_STATUS_SUCCESS;
}
static void ggml_backend_riscv64_spacemit_buffer_set_tensor(ggml_backend_buffer_t buffer,
struct ggml_tensor * tensor,
const void * data,
size_t offset,
size_t size) {
GGML_ASSERT(offset == 0);
GGML_ASSERT(size == ggml_nbytes(tensor));
auto tensor_traits = (ggml::cpu::riscv64_spacemit::tensor_traits_base *) tensor->extra;
if (tensor_traits) {
auto OK = tensor_traits->repack(tensor, data, size);
GGML_ASSERT(OK == 0);
}
GGML_UNUSED(buffer);
}
static const char * ggml_backend_cpu_riscv64_spacemit_buffer_type_get_name(ggml_backend_buffer_type_t buft) {
return "CPU_RISCV64_SPACEMIT";
GGML_UNUSED(buft);
}
static ggml_backend_buffer_t ggml_backend_cpu_riscv64_spacemit_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft,
size_t size) {
ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(ggml_backend_cpu_buffer_type(), size);
if (buffer == nullptr) {
return nullptr;
}
buffer->buft = buft;
buffer->iface.init_tensor = ggml_backend_riscv64_spacemit_buffer_init_tensor;
buffer->iface.set_tensor = ggml_backend_riscv64_spacemit_buffer_set_tensor;
buffer->iface.get_tensor = nullptr;
buffer->iface.cpy_tensor = nullptr;
return buffer;
}
static size_t ggml_backend_cpu_riscv64_spacemit_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) {
return 64;
GGML_UNUSED(buft);
}
static size_t ggml_backend_cpu_riscv64_spacemit_nbytes(ggml_backend_buffer_type_t buft,
const struct ggml_tensor * tensor) {
for (int i = 0; i < GGML_MAX_DIMS; ++i) {
if (tensor->ne[i] <= 0) {
return 0;
}
}
size_t nbytes;
const size_t blck_size = ggml_blck_size(tensor->type);
if (blck_size == 1) {
nbytes = ggml_type_size(tensor->type);
for (int i = 0; i < GGML_MAX_DIMS; ++i) {
nbytes += (tensor->ne[i] - 1) * tensor->nb[i];
}
} else {
nbytes = tensor->ne[0] * tensor->nb[0] / blck_size;
if (tensor->type == GGML_TYPE_Q4_K) {
GGML_ASSERT(nbytes % sizeof(block_q4_K) == 0);
nbytes = (nbytes / sizeof(block_q4_K)) * sizeof(block_q4_1) * 8;
for (int i = 1; i < GGML_MAX_DIMS; ++i) {
nbytes += (tensor->ne[i] - 1) * (tensor->nb[i] / sizeof(block_q4_K)) * sizeof(block_q4_1) * 8;
}
} else {
for (int i = 1; i < GGML_MAX_DIMS; ++i) {
nbytes += (tensor->ne[i] - 1) * tensor->nb[i];
}
}
}
GGML_UNUSED(buft);
return nbytes;
}
namespace ggml::cpu::riscv64_spacemit {
class extra_buffer_type : ggml::cpu::extra_buffer_type {
bool supports_op(ggml_backend_dev_t, const struct ggml_tensor * op) override {
switch (op->op) {
case GGML_OP_MUL_MAT:
if (op->src[0]->buffer && (ggml_n_dims(op->src[0]) == 2) &&
op->src[0]->buffer->buft == ggml_backend_cpu_riscv64_spacemit_buffer_type() &&
ggml_riscv64_spacemit_get_optimal_repack_type(op->src[0])) {
if (op->src[1]->buffer && !ggml_backend_buft_is_host(op->src[1]->buffer->buft)) {
return false;
}
if (op->src[1]->type == GGML_TYPE_F32) {
return true;
}
}
break;
case GGML_OP_NORM:
case GGML_OP_RMS_NORM:
if (op->src[0]->type == GGML_TYPE_F32) {
return true;
}
break;
default:
// GGML_ABORT("fatal error");
break;
}
return false;
}
ggml::cpu::tensor_traits * get_tensor_traits(const struct ggml_tensor * op) override {
switch (op->op) {
case GGML_OP_MUL_MAT:
if (op->src[0]->buffer && op->src[0]->buffer->buft == ggml_backend_cpu_riscv64_spacemit_buffer_type()) {
return (ggml::cpu::tensor_traits *) op->src[0]->extra;
}
break;
case GGML_OP_NORM:
case GGML_OP_RMS_NORM:
return (ggml::cpu::tensor_traits *) (&ggml::cpu::riscv64_spacemit::rvv_impl);
default:
// GGML_ABORT("fatal error");
break;
}
return nullptr;
}
};
} // namespace ggml::cpu::riscv64_spacemit
ggml_backend_buffer_type_t ggml_backend_cpu_riscv64_spacemit_buffer_type(void) {
static struct ggml_backend_buffer_type ggml_backend_cpu_buffer_type_riscv64_spacemit = {
/* .iface = */
{
/* .get_name = */ ggml_backend_cpu_riscv64_spacemit_buffer_type_get_name,
/* .alloc_buffer = */ ggml_backend_cpu_riscv64_spacemit_buffer_type_alloc_buffer,
/* .get_alignment = */ ggml_backend_cpu_riscv64_spacemit_buffer_type_get_alignment,
/* .get_max_size = */ nullptr,
/* .get_alloc_size = */ ggml_backend_cpu_riscv64_spacemit_nbytes,
/* .is_host = */ nullptr,
},
/* .device = */
ggml_backend_reg_dev_get(ggml_backend_cpu_reg(), 0),
/* .context = */
new ggml::cpu::riscv64_spacemit::extra_buffer_type(),
};
return &ggml_backend_cpu_buffer_type_riscv64_spacemit;
}
Reference in New Issue View Git Blame Copy Permalink