common.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include <stdint.h>
 
namespace ML {
 
// Dense input uses int64_t until FAISS is updated
typedef int64_t knn_indices_dense_t;
 
typedef int knn_indices_sparse_t;
 
template <typename value_idx, typename value_t>
struct knn_graph {
  knn_graph(value_idx n_rows_, int n_neighbors_)
    : n_rows(n_rows_), n_neighbors(n_neighbors_), knn_indices{nullptr}, knn_dists{nullptr}
  {
  }
 
  knn_graph(value_idx n_rows_, int n_neighbors_, value_idx* knn_indices_, value_t* knn_dists_)
    : n_rows(n_rows_), n_neighbors(n_neighbors_), knn_indices(knn_indices_), knn_dists(knn_dists_)
  {
  }
 
  value_idx* knn_indices;
  value_t* knn_dists;
 
  value_idx n_rows;
  int n_neighbors;
};
 
template <typename T>
struct manifold_inputs_t {
  T* y;
  int n;
  int d;
 
  manifold_inputs_t(T* y_, int n_, int d_) : y(y_), n(n_), d(d_) {}
 
  virtual bool alloc_knn_graph() const = 0;
};
 
template <typename T>
struct manifold_dense_inputs_t : public manifold_inputs_t<T> {
  T* X;
 
  manifold_dense_inputs_t(T* x_, T* y_, int n_, int d_) : manifold_inputs_t<T>(y_, n_, d_), X(x_) {}
 
  bool alloc_knn_graph() const { return true; }
};
 
template <typename value_idx, typename T>
struct manifold_sparse_inputs_t : public manifold_inputs_t<T> {
  value_idx* indptr;
  value_idx* indices;
  T* data;
 
  size_t nnz;
 
  manifold_sparse_inputs_t(
    value_idx* indptr_, value_idx* indices_, T* data_, T* y_, size_t nnz_, int n_, int d_)
    : manifold_inputs_t<T>(y_, n_, d_), indptr(indptr_), indices(indices_), data(data_), nnz(nnz_)
  {
  }
 
  bool alloc_knn_graph() const { return true; }
};
 
template <typename value_idx, typename value_t>
struct manifold_precomputed_knn_inputs_t : public manifold_inputs_t<value_t> {
  manifold_precomputed_knn_inputs_t<value_idx, value_t>(
    value_idx* knn_indices_, value_t* knn_dists_, value_t* y_, int n_, int d_, int n_neighbors_)
    : manifold_inputs_t<value_t>(y_, n_, d_), knn_graph(n_, n_neighbors_, knn_indices_, knn_dists_)
  {
  }
 
  knn_graph<value_idx, value_t> knn_graph;
 
  bool alloc_knn_graph() const { return false; }
};
 
};  // end namespace ML