KMeans_CUDA_PyBindings.cu

// Python Bindings for KMeans_CUDA

#include "KMeans_CUDA.h"

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>

namespace py = pybind11;

// Model / Data
KMeans_CUDA* model;
float* h_data;

// Copies memory to GPU
// FIXME -- Find a better way to pass and store the data.
void KMeans_CUDA_constructor(py::array_t<double> data, int N, int D, int K) {
    h_data = new float[N];

    auto buffer = data.request();
    double* ptr = static_cast<double*>(buffer.ptr);
    for (ssize_t i = 0; i < buffer.size; i++) {
        h_data[i] = ptr[i];
    }

    model = new KMeans_CUDA(h_data, N, D, K);
}

// Run one epoch
void one_epoch() {
    model->one_epoch();
}

// Runs `epochs` number of iterations
void fit(int epochs) {
    model->fit(epochs);
}

// Compute error
float compute_error() {
    return model->compute_error();
}

// Predictions for the current data with current clusters.
py::array_t<int> predictions() {
    vector<int> predictions = model->predictions();
    return py::array_t<int>(
        {predictions.size()},
        {sizeof(int)},
        predictions.data()
    );

}

// Deletes CPU memory. GPU memory is deleted in the KMeans_CUDA class.
// FIXME -- is there a way to call the destructor without having the user do so in Python?
void release_memory() {
    delete model;
    delete[] h_data;
}

// Binding functions
PYBIND11_MODULE(KMeans_CUDA, m) {
    m.def("KMeans_CUDA_constructor", &KMeans_CUDA_constructor);
    m.def("one_epoch", &one_epoch);
    m.def("fit", &fit);
    m.def("compute_error", &compute_error);
    m.def("predictions", &predictions);
    m.def("release_memory", &release_memory);
}