Method process_linear_weights

src/models/model.cc:389–425

This method runs some precomputations on linear weights when possible.

    // This method runs some precomputations on linear weights when possible.
    void Model::process_linear_weights() {
      if (_device != Device::CPU)
        return;  // There is currently no processing for non CPU device.

      const bool pack_weights = cpu::pack_gemm_weights(_effective_compute_type);
      const bool transpose = true;
      const float alpha = 1;

      const auto variable_index = _variable_index;
      for (const auto& pair : variable_index) {
        const std::string& name = pair.first;
        if (!is_linear_weight(name))
          continue;

        const StorageView& weight = *pair.second;
        const DataType dtype = weight.dtype();
        const dim_t k = weight.dim(1);
        const dim_t n = weight.dim(0);

        // If the target Gemm implementation prefers the u8s8s32 format, we can shift
        // the input of linear layers to the u8 domain and add a compensation term.
        // This term only depends on the linear weight, so we can compute it once and
        // store it as a model variable.
        if (dtype == DataType::INT8 && cpu::prefer_u8s8s32_gemm()) {
          StorageView compensation = ops::Gemm::compensate_u8_input(weight, transpose, k, n, alpha);
          register_variable(name + "_compensation", std::move(compensation));
        }

        // If requested, linear weights can be packed for the Gemm call.
        // Only 2D weights are supported; Conv weights are 3D and use a different code path.
        if (pack_weights && is_packable(name) && weight.rank() == 2) {
          StorageView packed_weight = ops::Gemm::pack_b_input(weight, transpose, k, n, alpha);
          register_variable(name + "_packed", std::move(packed_weight));
          remove_variable(name);  // The original weight is no longer needed.
        }
      }
    }

    static DataType get_dtype_from_item_size(uint8_t item_size) {
      // This is the old (and flawed) logic of resolving the dtype of saved variables.