compressed_count_blocks Class — pytorch Architecture

Source Code

aten/src/ATen/native/TensorConversions.cpp lines 2079–2101

template <class index_t>
static index_t compressed_count_blocks(
    const index_t n_compressed, // Tensor size along compressed dimension
    const index_t n_plain, // Tensor size along plain dimension
    const index_t C, // Block size along compressed dimensions
    const index_t P, // Block size along plain dimension
    const index_t Ac[], // Compressed indices
    const index_t Ap[] // Plain indices
) {
  std::vector<index_t> mask(n_plain / P + 1, -1);
  index_t n_blks = 0;
  for (index_t c = 0; c < n_compressed; c++) {
    index_t bc = c / C;
    for (index_t i = Ac[c]; i < Ac[c + 1]; i++) {
      index_t bp = Ap[i] / P;
      if (mask[bp] != bc) {
        mask[bp] = bc;
        n_blks++;
      }
    }
  }
  return n_blks;
}