diff --git a/lib_nn/api/nn_layers.h b/lib_nn/api/nn_layers.h index ca3640e8..50550c10 100644 --- a/lib_nn/api/nn_layers.h +++ b/lib_nn/api/nn_layers.h @@ -308,4 +308,23 @@ void mean_int16(const int16_t *input, int16_t *output, const int start_dim_size, const int mean_dim_size, const int end_dim_size, const float scale_mul); +typedef struct { + float lhs_zp; + float rhs_zp; + float in_zp_sum; // channel_size * lhs_zp * rhs_zp + float out_zp; + float scale; // lhs_scale*rhs_scale/out_scale + uint32_t lhs_row_size; + uint32_t channel_size; // lhs col size & rhs row size + uint32_t rhs_col_size; +} nn_mat_mul_real_params_t; +/** + * @brief Execute real matrix multiplication + * @param vpu_buf int8_t vpu_buf[64], need word align + */ +void mat_mul_real_int8( + nn_mat_mul_real_params_t *p, + int8_t *vpu_buf0, int8_t *vpu_buf1, int8_t *vpu_buf2, + int8_t *lhs, int8_t* rhs, int8_t *output); + #endif // LAYERS_H_ diff --git a/lib_nn/src/asm/round8.S b/lib_nn/src/asm/round8.S new file mode 100644 index 00000000..fc3bc5ba --- /dev/null +++ b/lib_nn/src/asm/round8.S @@ -0,0 +1,65 @@ +// Copyright 2026 XMOS LIMITED. +// This Software is subject to the terms of the XMOS Public Licence: Version 1. +#if defined(__XS3A__) + +#include + +/** + * int8_t round8(float r) // round-half-away-from-zero + */ + +#define FUNCTION_NAME round8 + + .text + .cc_top FUNCTION_NAME.function,FUNCTION_NAME + +#define NSTACKWORDS 0 + .issue_mode dual + + .align 4 + + .globl FUNCTION_NAME + .align 16 +FUNCTION_NAME: + { dualentsp NSTACKWORDS ; nop } + + ldc r1, 0x42fe + ldc r2, 0xc300 + { shl r1, r1, 16 ; shl r2, r2, 16 } // r1: 127.0, r2: -128.0 + flt r1, r0, r1 // r1: r < 127.0 + { bf r1, overflow127 ; ldc r3, 22 } // r3: 22 + fgt r2, r0, r2 // r2: r > -128.0 + { bf r2, underflow_neg_128 ; add r2, r0, 0 } // r2: r + + fmant r0, r2 + fsexp r1, r2, r2 // r1: sign, r2: exp, r0: mantissa + + { sub r2, r3, r2 ; nop } // r2: 22-exp + { shr r0, r0, r2 ; ldc r3, 1 } // r0: %.1f(r) + { and r3, r0, r3 ; shr r0, r0, 1 } // r3: rounding bit, r0: integer_part + { add r0, r0, r3 ; bf r1, exit } // r0: rounded integer + { neg r0, r0 ; retsp NSTACKWORDS } + +overflow127: + ldc r0, 127 + retsp NSTACKWORDS + +underflow_neg_128: + ldc r0, 128 + neg r0, r0 // r0: -128 +exit: + retsp NSTACKWORDS + + .cc_bottom FUNCTION_NAME.function + .set FUNCTION_NAME.nstackwords,NSTACKWORDS + .globl FUNCTION_NAME.nstackwords + .set FUNCTION_NAME.maxcores,1 + .globl FUNCTION_NAME.maxcores + .set FUNCTION_NAME.maxtimers,0 + .globl FUNCTION_NAME.maxtimers + .set FUNCTION_NAME.maxchanends,0 + .globl FUNCTION_NAME.maxchanends +.Ltmp1: + .size FUNCTION_NAME, .Ltmp1-FUNCTION_NAME + +#endif diff --git a/lib_nn/src/asm/vect_mat_mul_int8_asm.S b/lib_nn/src/asm/vect_mat_mul_int8_asm.S new file mode 100644 index 00000000..a3f5a5ec --- /dev/null +++ b/lib_nn/src/asm/vect_mat_mul_int8_asm.S @@ -0,0 +1,128 @@ +// Copyright 2026 XMOS LIMITED. +// This Software is subject to the terms of the XMOS Public Licence: Version 1. +#if defined(__XS3A__) + +#include + +// Assume rhs is column major, [col, row] + +/** + * void vect_mat_mul_int8_asm( + * const int8_t *lhs, + * const int8_t *rhs, + * int8_t *vpu_buffer, + * uint32_t channel_size, // lhs size, rhs row size + * uint32_t rhs_col_size); + */ + +#define FUNCTION_NAME vect_mat_mul_int8_asm + + .text + .cc_top FUNCTION_NAME.function,FUNCTION_NAME + +#define NSTACKWORDS 16 + .issue_mode dual + + .align 4 + + .globl FUNCTION_NAME + .align 16 +FUNCTION_NAME: + { dualentsp NSTACKWORDS ; nop } + + stw r4, sp[NSTACKWORDS-1] + stw r5, sp[NSTACKWORDS-2] + stw r6, sp[NSTACKWORDS-3] + stw r7, sp[NSTACKWORDS-4] + stw r8, sp[NSTACKWORDS-5] + stw r9, sp[NSTACKWORDS-6] + +#define lhs r0 +#define rhs r1 +#define vpu_buffer r2 +#define ch_size r3 +#define ch_cnt r6 +#define col r9 + + ldw col, sp[NSTACKWORDS+1] + ldc r11, 512 + { vsetc r11 ; ldc r4, 32 } + { vclrdr ; shl r5, ch_size, 4 } // r5: channel_size*16 + { vstr vpu_buffer[0] ; add r11, vpu_buffer, r4 } // make sure vpu buffer is all 0 + { vstr r11[0] ; nop } // make sure vpu buffer is all 0 + { add ch_cnt, ch_size, 0 ; ldc r11, 16 } + { lsu col, col, r11 ; sub rhs, rhs, ch_size } // rhs point to ch15 + { bf col, maccr_prepare ; ldc col, 0 } + ldw col, sp[NSTACKWORDS+1] + { sub col, r11, col ; nop } // col: tail loop cnt + { shl col, col, 2 ; nop } // col: maccr loop offset (in byte) + +maccr_prepare: + { lsu r11, ch_cnt, r4 ; nop } + { bt r11, lhs_zero_padding ; add r11, lhs, 0 } + { vldc lhs[0] ; add lhs, lhs, r4 } // Load lhs to vC + +maccr_loop_check: + { bt col, tail_prepare ; ldap r11, maccr } + +maccr: + { vlmaccr rhs[0] ; sub rhs, rhs, ch_size } + { vlmaccr rhs[0] ; sub rhs, rhs, ch_size } + { vlmaccr rhs[0] ; sub rhs, rhs, ch_size } + { vlmaccr rhs[0] ; sub rhs, rhs, ch_size } + { vlmaccr rhs[0] ; sub rhs, rhs, ch_size } + { vlmaccr rhs[0] ; sub rhs, rhs, ch_size } + { vlmaccr rhs[0] ; sub rhs, rhs, ch_size } + { vlmaccr rhs[0] ; sub rhs, rhs, ch_size } + { vlmaccr rhs[0] ; sub rhs, rhs, ch_size } + { vlmaccr rhs[0] ; sub rhs, rhs, ch_size } + { vlmaccr rhs[0] ; sub rhs, rhs, ch_size } + { vlmaccr rhs[0] ; sub rhs, rhs, ch_size } + { vlmaccr rhs[0] ; sub rhs, rhs, ch_size } + { vlmaccr rhs[0] ; sub rhs, rhs, ch_size } + { vlmaccr rhs[0] ; sub rhs, rhs, ch_size } + { vlmaccr rhs[0] ; sub rhs, rhs, ch_size } + { add rhs, rhs, r5 ; lsu r11, ch_cnt, r4 } // rollback + { add rhs, rhs, r4 ; bt r11, exit } // processed offset + { sub ch_cnt, ch_cnt, r4 ; ldc r11, 16 } // channel_cnt -= 32 + { bt ch_cnt, maccr_prepare ; nop } + +exit: + { vstd vpu_buffer[0] ; add vpu_buffer, vpu_buffer, r4} + { vstr vpu_buffer[0] ; nop } + ldw r4, sp[NSTACKWORDS-1] + ldw r5, sp[NSTACKWORDS-2] + ldw r6, sp[NSTACKWORDS-3] + ldw r7, sp[NSTACKWORDS-4] + ldw r8, sp[NSTACKWORDS-5] + ldw r9, sp[NSTACKWORDS-6] + retsp NSTACKWORDS + +lhs_zero_padding: + { vstr vpu_buffer[0] ; add vpu_buffer, vpu_buffer, r4 } // store acc vR to vpu buffer + { vldr r11[0] ; mkmsk r11, ch_cnt } // r11: mask of valid lhs byte } + vstrpv vpu_buffer[0], r11 // store valid lhs data to vpu_buffer[63:32] + { vldc vpu_buffer[0] ; sub r11, vpu_buffer, r4 } // r11: vpu_buffer[0] + { vldr r11[0] ; add vpu_buffer, r11, 0 } // restore vpu_buffer to vpu_buffer[0] + { bu maccr_loop_check ; nop } + +tail_prepare: + { add r11, r11, col ; add r7, col, 0 } // r7: tail loop count x4 +maccr_tail: + { vlmaccr lhs[0] ; sub r7, r7, 4 } + { bt r7, maccr_tail ; sub rhs, rhs, ch_size } + { bau r11 ; nop } + + .cc_bottom FUNCTION_NAME.function + .set FUNCTION_NAME.nstackwords, NSTACKWORDS+s32_to_f32.nstackwords+round8.nstackwords + .globl FUNCTION_NAME.nstackwords + .set FUNCTION_NAME.maxcores,1 + .globl FUNCTION_NAME.maxcores + .set FUNCTION_NAME.maxtimers,0 + .globl FUNCTION_NAME.maxtimers + .set FUNCTION_NAME.maxchanends,0 + .globl FUNCTION_NAME.maxchanends +.Ltmp1: + .size FUNCTION_NAME, .Ltmp1-FUNCTION_NAME + +#endif diff --git a/lib_nn/src/c/matmul.c b/lib_nn/src/c/matmul.c new file mode 100644 index 00000000..7391edd2 --- /dev/null +++ b/lib_nn/src/c/matmul.c @@ -0,0 +1,114 @@ + +#include +#include +#include +#include +#include +#include + +#include "nn_op_helper.h" +#include "nn_operator.h" +#include "vpu_sim.h" + +void vect_mat_mul_int8_asm( + const int8_t *lhs, + const int8_t *rhs, + int8_t *vpu_buffer, + uint32_t channel_size, // lhs size, rhs row size + uint32_t rhs_col_size +); + +#if !defined(NN_USE_REF) && defined(__XS3A__) +void mat_mul_real_int8_vpu( + nn_mat_mul_real_params_t *p, + int8_t *vpu_buf0, int8_t *vpu_buf1, int8_t *vpu_buf2, + int8_t *lhs, int8_t* rhs, int8_t *output) { + for (int lhs_row = 0; lhs_row < p->lhs_row_size; ++lhs_row) { + // TODO: optimize it with vpu + int32_t lhs_row_sum = 0; + for (int i = 0; i < p->channel_size; ++i) { + lhs_row_sum += lhs[lhs_row * p->channel_size + i]; + } + for (int rhs_col = 0; rhs_col < p->rhs_col_size; rhs_col+=16) { + int8_t* lhs_temp = &lhs[lhs_row * p->channel_size]; + int8_t* rhs_temp = &rhs[(rhs_col+16)*p->channel_size]; + int process_col = p->rhs_col_size - rhs_col; + process_col = process_col > 16 ? 16 : process_col; + + vect_mat_mul_int8_asm( + lhs_temp, rhs_temp, vpu_buf0, p->channel_size, process_col); + + // TODO: optimize it with vpu + int32_t *buff_temp = vpu_buf1; // cheating here, treating vD:vR as continue space + for (int i = 0; i < process_col; ++i) { + buff_temp[i] = 0; + for (int j = 0; j < p->channel_size; ++j) { + buff_temp[i] += rhs[rhs_col*p->channel_size + i*p->channel_size + j]; + } + } + uint16_t *vD = (uint16_t*)(&vpu_buf0[0]); + uint16_t *vR = (uint16_t*)(&vpu_buf0[32]); + for (int i = 0; i < process_col; ++i) { + uint32_t uacc = (vD[i] << 16) | vR[i]; + int32_t acc = *((int32_t*)(&uacc)); + float accf = + ((float)acc) + -(p->rhs_zp * lhs_row_sum) + -(p->lhs_zp * buff_temp[i]) + +p->in_zp_sum; + accf *= p->scale; + accf += p->out_zp; + if (accf > 127.0f) + accf = 127.0f; + else if (accf < -128.0f) + accf = -128.0f; + output[i] = roundf(accf); + } + output = &output[process_col]; + } + } +} +#endif // NN_USE_REF + +void mat_mul_real_int8_ref( + nn_mat_mul_real_params_t *p, + int8_t *vpu_buf0, int8_t *vpu_buf1, int8_t *vpu_buf2, + int8_t *lhs, int8_t* rhs, int8_t *output) +{ + int out_index = 0; + for (int i = 0; i < p->lhs_row_size; ++i) { + for (int j = 0; j < p->rhs_col_size; ++j) { + double acc = 0.0; + for (int k = 0; k < p->channel_size; ++k) { + int lhs_idx = i*p->channel_size+k; + int rhs_idx = j*p->channel_size+k; + double x = ((double)(lhs[lhs_idx]) - p->lhs_zp); + double y = ((double)(rhs[rhs_idx]) - p->rhs_zp); + acc += x * y; + } + if (p->rhs_col_size == 6 && out_index == 6 && p->channel_size == 4) { + printf("acc before scale: %.10f\n", acc); + } + float quantized_value = (float)acc * p->scale + p->out_zp; + // Clamp the quantized value to int8 range + if (quantized_value > 127.0f) + quantized_value = 127.0f; + else if (quantized_value < -128.0f) + quantized_value = -128.0f; + output[out_index++] = (int8_t)(roundf(quantized_value)); + } + } +} + +// A real mat mul here +void mat_mul_real_int8( + nn_mat_mul_real_params_t *p, + int8_t *vpu_buf0, int8_t *vpu_buf1, int8_t *vpu_buf2, + int8_t *lhs, int8_t* rhs, int8_t *output) { + +#if !defined(NN_USE_REF) && defined(__XS3A__) + mat_mul_real_int8_vpu(p, vpu_buf0, vpu_buf1, vpu_buf2, lhs, rhs, output); +#else + mat_mul_real_int8_ref(p, vpu_buf0, vpu_buf1, vpu_buf2, lhs, rhs, output); +#endif +} \ No newline at end of file