1#include "zn_layouts.h"
 2
 3zn32_pmat_layout::zn32_pmat_layout(uint64_t nrows, uint64_t ncols)
 4    : nrows(nrows),  //
 5      ncols(ncols),  //
 6      data((ZN32_VMP_PMAT*)malloc((nrows * ncols + 7) * sizeof(int32_t))) {}
 7
 8zn32_pmat_layout::~zn32_pmat_layout() { free(data); }
 9
10int32_t* zn32_pmat_layout::get_addr(uint64_t row, uint64_t col) const {
11  REQUIRE_DRAMATICALLY(row < nrows, "row overflow" << row << " / " << nrows);
12  REQUIRE_DRAMATICALLY(col < ncols, "col overflow" << col << " / " << ncols);
13  const uint64_t nblk = ncols >> 5;
14  const uint64_t rem_ncols = ncols & 31;
15  uint64_t blk = col >> 5;
16  uint64_t col_rem = col & 31;
17  if (blk < nblk) {
18    // column is part of a full block
19    return (int32_t*)data + blk * nrows * 32 + row * 32 + col_rem;
20  } else {
21    // column is part of the last block
22    return (int32_t*)data + blk * nrows * 32 + row * rem_ncols + col_rem;
23  }
24}
25int32_t zn32_pmat_layout::get(uint64_t row, uint64_t col) const { return *get_addr(row, col); }
26int32_t zn32_pmat_layout::get_zext(uint64_t row, uint64_t col) const {
27  if (row >= nrows || col >= ncols) return 0;
28  return *get_addr(row, col);
29}
30void zn32_pmat_layout::set(uint64_t row, uint64_t col, int32_t value) { *get_addr(row, col) = value; }
31void zn32_pmat_layout::fill_random() {
32  int32_t* d = (int32_t*)data;
33  for (uint64_t i = 0; i < nrows * ncols; ++i) d[i] = uniform_i64_bits(32);
34}
35thash zn32_pmat_layout::content_hash() const { return test_hash(data, nrows * ncols * sizeof(int32_t)); }
36
37template <typename T>
38std::vector<int32_t> vmp_product(const T* vec, uint64_t vec_size, uint64_t out_size, const zn32_pmat_layout& mat) {
39  uint64_t rows = std::min(vec_size, mat.nrows);
40  uint64_t cols = std::min(out_size, mat.ncols);
41  std::vector<int32_t> res(out_size, 0);
42  for (uint64_t j = 0; j < cols; ++j) {
43    for (uint64_t i = 0; i < rows; ++i) {
44      res[j] += vec[i] * mat.get(i, j);
45    }
46  }
47  return res;
48}
49
50template std::vector<int32_t> vmp_product(const int8_t* vec, uint64_t vec_size, uint64_t out_size,
51                                          const zn32_pmat_layout& mat);
52template std::vector<int32_t> vmp_product(const int16_t* vec, uint64_t vec_size, uint64_t out_size,
53                                          const zn32_pmat_layout& mat);
54template std::vector<int32_t> vmp_product(const int32_t* vec, uint64_t vec_size, uint64_t out_size,
55                                          const zn32_pmat_layout& mat);