Function _matmul

src/infer.cpp:121–157

#endif

static void _matmul(
  float* xout, float* x, float* w, int n, int d, 
  const int* block_size, float* scale,
  void* unused_aqb
) {
  // W (d,n) @ x (n,) -> xout (d,)
  (void)unused_aqb;
  static float one = 1.0f;
  int dummy_block_size[2] = {d, n};
  if (scale == nullptr) {
    scale = &one;
    block_size = dummy_block_size;
  }
  int scale_num_cols = (n + block_size[1] - 1) / block_size[1];
  for (int scale_i = 0; scale_i < cdiv(d, block_size[0]); scale_i++) {
    int ii;
#pragma omp parallel for private(ii)
    for (ii = 0; ii < block_size[0]; ii++) {
      int i = scale_i * block_size[0] + ii;
      if (i >= d) {
        continue;
      }
      float val = 0.0f;
      for (int scale_j = 0; scale_j < cdiv(n, block_size[1]); scale_j++) {
        float scale_val = scale[scale_i * scale_num_cols + scale_j];
        for (int jj = 0; jj < block_size[1]; jj++) {
          int j = scale_j * block_size[1] + jj;
          if (j >= n) {
            break;
          }
          val += (w[i * n + j] * x[j]) * scale_val;
        }
      }
      xout[i] = val;
    }
  }
}

// matmul supporting float16 weights via the F16C extension, which allows
// conversion into float32 values before calculations.