Added: a real matmul function and vector matrix mul assmebly for xs3

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_

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 <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

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 <xs1.h>
+
+// 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

lib_nn/src/c/matmul.c

@@ -0,0 +1,114 @@
+
+#include <assert.h>
+#include <math.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#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
+}