| 94 | |
| 95 | template <MetricType metric> |
| 96 | bool |
| 97 | ProductQuantizer<metric>::EncodeOneImpl(const DataType* data, uint8_t* codes) { |
| 98 | const DataType* cur = data; |
| 99 | Vector<float> tmp(this->allocator_); |
| 100 | if constexpr (metric == MetricType::METRIC_TYPE_COSINE) { |
| 101 | tmp.resize(this->dim_); |
| 102 | Normalize(data, tmp.data(), this->dim_); |
| 103 | cur = tmp.data(); |
| 104 | } |
| 105 | for (int i = 0; i < pq_dim_; ++i) { |
| 106 | // TODO(LHT): use blas |
| 107 | float nearest_dis = std::numeric_limits<float>::max(); |
| 108 | uint8_t nearest_id = 0; |
| 109 | const float* query = cur + i * subspace_dim_; |
| 110 | const float* base = this->codebooks_.data() + i * subspace_dim_ * CENTROIDS_PER_SUBSPACE; |
| 111 | for (int j = 0; j < CENTROIDS_PER_SUBSPACE; ++j) { |
| 112 | float dist = FP32ComputeL2Sqr(query, base + j * subspace_dim_, subspace_dim_); |
| 113 | if (dist < nearest_dis) { |
| 114 | nearest_dis = dist; |
| 115 | nearest_id = static_cast<uint8_t>(j); |
| 116 | } |
| 117 | } |
| 118 | codes[i] = nearest_id; |
| 119 | } |
| 120 | return true; |
| 121 | } |
| 122 | |
| 123 | template <MetricType metric> |
| 124 | bool |
no test coverage detected