Skip to content

Added vpu optimized mean int8 #221

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
LeslieXMOS wants to merge 1 commit into
base: develop
Choose a base branch
from
+259 −0
Open

Added vpu optimized mean int8 #221

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
167 changes: 167 additions & 0 deletions lib_nn/src/asm/mean8.S
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,167 @@
// Copyright 2026 XMOS LIMITED.
// This Software is subject to the terms of the XMOS Public Licence: Version 1.
#if defined(__XS3A__)

#include <xs1.h>

// scale_mul is in_scale / out_scale
// only support mean axis is the last axis, which end_dim_size = 1
// mean_dim_size needs to be word aligned

/**
* void mean_int8_asm(
* const int8_t *input,
* int8_t *output,
* const int start_dim_size,
* const int mean_dim_size,
* const int8_t *vpu_buffer, // 64 byte
* const float in_zero_point_sum, // in_zero_point*mean_dim_size
* const float out_zero_point,
* const float scale_mul)
*/

#define FUNCTION_NAME mean_int8_asm

.text
.cc_top FUNCTION_NAME.function,FUNCTION_NAME

#define NSTACKWORDS 32
.issue_mode dual

.align 4

filter:
.word 0x01010101, 0x01010101, 0x01010101, 0x01010101
.word 0x01010101, 0x01010101, 0x01010101, 0x01010101
.word 0x00000000, 0x00000000, 0x00000000, 0x00000000
.word 0x00000000, 0x00000000, 0x00000000, 0x00000000

#define EPV 32

.globl FUNCTION_NAME
.align 16
FUNCTION_NAME:
{ dualentsp NSTACKWORDS ; nop }

@McCrea McCrea Jul 24, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

DUALENTSP_lu6 is usually better unless there's a useful instruction you can packetise - that way it takes a single cycle whichever issue mode the function is entered it.


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]
stw r10, sp[NSTACKWORDS-7]
Comment on lines +46 to +52

@McCrea McCrea Jul 24, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could use some stds to save a bit of time.


#define input r0
#define output r1
#define channel_cnt r2
#define mean_size r3
#define vpu_buffer r4
#define mean_cnt r5

ldw vpu_buffer, sp[NSTACKWORDS+1]

@McCrea McCrea Jul 24, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks like this is mean to find a base address within the current stack frame - should it be an ldaw?

@LeslieXMOS LeslieXMOS Jul 24, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This load the vpu_buffer pointer address passes as the fourth argument of the function.
The vpu buffer is created outside the function, not live within the stack of this function.

@McCrea McCrea Jul 24, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, yes, I see. (Sorry, I forget it's not RISC-V sometimes - where we have 8 argument registers).

ldc r11, 512 // r11: VSETCTRL_TYPE_INT8
{ vsetc r11 ; ldc r8, 16 } // r8: 16

prepare:
{ add mean_cnt, mean_size, 0 ; shl r9, mean_size, 4 } // r9: 16th channel offset in byte
{ add input, input, r9 ; vclrdr } // input: n+16th channel first byte
{ sub input, input, mean_size ; nop } // input: n+15th channel first byte

prepare_filter:
{ ldap r11, filter ; ldc r10, 32 }
{ lsu r9, mean_cnt, r10 ; nop } // r9: mean_cnt < 32, r10: offset in byte
{ bf r9, maccr_prepare ; nop }
{ sub r10, r10, mean_cnt ; nop }
{ add r11, r11, r10 ; add r10, mean_cnt, 0 } // r11: offset filter, r10: actual data goes through filter in next maccr

maccr_prepare:
{ vldc r11[0] ; nop } // load filter to vC

maccr_loop_check:
{ lsu r11, channel_cnt, r8 ; sub r9, r8, channel_cnt } // r11: channel_cnt < 16, r9: tailing loop count
{ bt r11, maccr_tail_prepare ; ldap r11, maccr_loop }
maccr_loop:
{ vlmaccr input[0] ; sub input, input, mean_size } // summation 32 byte in channel n+15
{ vlmaccr input[0] ; sub input, input, mean_size } // summation 32 byte in channel n+14
{ vlmaccr input[0] ; sub input, input, mean_size } // summation 32 byte in channel n+13
{ vlmaccr input[0] ; sub input, input, mean_size } // summation 32 byte in channel n+12
{ vlmaccr input[0] ; sub input, input, mean_size } // summation 32 byte in channel n+11
{ vlmaccr input[0] ; sub input, input, mean_size } // summation 32 byte in channel n+10
{ vlmaccr input[0] ; sub input, input, mean_size } // summation 32 byte in channel n+9
{ vlmaccr input[0] ; sub input, input, mean_size } // summation 32 byte in channel n+8
{ vlmaccr input[0] ; sub input, input, mean_size } // summation 32 byte in channel n+7
{ vlmaccr input[0] ; sub input, input, mean_size } // summation 32 byte in channel n+6
{ vlmaccr input[0] ; sub input, input, mean_size } // summation 32 byte in channel n+5
{ vlmaccr input[0] ; sub input, input, mean_size } // summation 32 byte in channel n+4
{ vlmaccr input[0] ; sub input, input, mean_size } // summation 32 byte in channel n+3
{ vlmaccr input[0] ; sub input, input, mean_size } // summation 32 byte in channel n+2
{ vlmaccr input[0] ; sub input, input, mean_size } // summation 32 byte in channel n+1
{ vlmaccr input[0] ; sub input, input, mean_size } // summation 32 byte in channel n+0
{ shl r9, mean_size, 4 ; add input, input, r10 } // r9: number of byte to loop back, r0: +32 byte offset
{ add input, input, r9 ; ldc r10, 32 } // r0: next input for channel n+15, r10: 32
{ lsu r9, mean_cnt, r10 ; nop } // r9: this is last loop
{ bt r9, output ; sub mean_cnt, mean_cnt, r10 }
{ bf mean_cnt, output ; lsu r9, mean_cnt, r10 } // r9: next is tailing loop
{ bt r9, prepare_filter ; nop }
{ bu maccr_loop_check ; nop }

output:
{ vstd vpu_buffer[0] ; add r9, vpu_buffer, r10 } // vpu_buffer[0]: vD, r9: vpu_buffer+32
{ vstr r9[0] ; ldc r9, 0 } // vpu_buffer[32]: vR, r9: 0
std r3, r0, sp[0]
output_loop:
std r2, r1, sp[1] // store output_ptr and channel_cnt to sp
{ ldc r10, 16 ; nop } // r10: 16 as ld16s access in sub word
{ ld16s r7, vpu_buffer[r9] ; add r9, r9, r10 } // r7: vD[i+0...i+1]
{ ld16s r6, vpu_buffer[r9] ; add r9, r9, 1 } // r6: vR[i+0...i+1]
// prepare full 32-bit accumulator value
{ shl r7, r7, 16 ; shl r6, r6, 16 }
{ shr r6, r6, 16 ; ldc r1, 0 } // r1: exp=0
{ or r0, r7, r6 ; ldc r6, 0 } // r0: 32bit accumulator, r6: 0
bl s32_to_f32 // r0: accumulator value in float
{ ldw r1, sp[NSTACKWORDS+2] ; sub r9, r9, r10 } // r1: in zero point sum, r9: rollback for vD
{ ldw r2, sp[NSTACKWORDS+3] ; eq r7, r9, r10 } // r2: out zero point, r7: last output loop
{ ldw r3, sp[NSTACKWORDS+4] ; nop } // r3: scale
fsub r0, r0, r1 // r0: acc-in_zero_sum
fmacc r0, r2, r0, r3 // r0: r0*scale+out_zero
bl round8 // r0: round(r0) in int8_t
ldd r2, r1, sp[1] // r1: output_ptr, r2: channel_cnt
st8 r0, r1[r6] // store result to output
{ add r1, r1, 1 ; sub r2, r2, 1 } // r2: num of channel out left
{ bf r2, exit ; nop }
{ bf r7, output_loop ; nop }
ldd r3, r0, sp[0]
bu prepare

exit:
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]
ldw r10, sp[NSTACKWORDS-7]
retsp NSTACKWORDS

maccr_tail_prepare:
mul r7, r9, mean_size // r7: input loopback to valid addr offset size (in byte)
{ shl r9, r9, 2 ; nop } // r9: maccr_loop offset (in byte)
{ add r11, r11, r9 ; sub input, input, r7 } // r11: maccr_loop offset start addr
maccr_tail_loop:
{ vlmaccr input[0] ; sub r9, r9, 4 }
{ bt r9, maccr_tail_loop ; nop }
{ bau r11 ; nop }

.cc_bottom FUNCTION_NAME.function
.set FUNCTION_NAME.nstackwords, NSTACKWORDS+s32_to_f32.nstackwords+round8.nstackwords

@McCrea McCrea Jul 24, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should be NSTACKWORDS + (s32_to_f32.nstackwords $M round8.nstackwords) I think (round8 and s32_to_f32 aren't called concurrently.

.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
65 changes: 65 additions & 0 deletions lib_nn/src/asm/round8.S
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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 <xs1.h>

/**
* 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
27 changes: 27 additions & 0 deletions lib_nn/src/c/mean.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,39 @@
#include <math.h>
#include <stdint.h>

/**
* VPU optimized int8 mean calculation
* Only support mean axis be the last axis, which end_dim_size = 1
* mean_dim_size needs to be divided by 4
*/
extern void mean_int8_asm(
const int8_t *input,
int8_t *output,
const int start_dim_size,
const int mean_dim_size,
const int8_t *vpu_buffer, // 64 byte vpu buffer
const float in_zero_point_sum,
const float out_zero_point,
const float scale_mul);

// scale_mul is in_scale / out_scale
void mean_int8(const int8_t *input, int8_t *output, const int start_dim_size,
const int mean_dim_size, const int end_dim_size,
const float in_zero_point, const float out_zero_point,
const float scale_mul) {

#if !defined(NN_USE_REF) && defined(__XS3A__)
if ((end_dim_size == 1) && (mean_dim_size % 4 == 0)) {
int8_t vpu_buffer[64];
float in_zero_point_sum = (float)((int32_t)(in_zero_point*mean_dim_size)); // rounding it to keep the same performance as ref
mean_int8_asm(
input, output, start_dim_size, mean_dim_size,
vpu_buffer,
in_zero_point_sum, out_zero_point, scale_mul);
return;
}
#endif

const int32_t start = -in_zero_point * mean_dim_size;
for (int i = 0; i < start_dim_size; ++i) {
const int i_mul = i * mean_dim_size * end_dim_size;
Expand Down