Add torch library bindings (guard main, add torch wrapper functions)

minkowski_spectrum/minkowski_spectrum.cu

@@ -1,3 +1,7 @@
+#ifdef TORCH_EXTENSION_NAME
+#include <torch/torch.h>
+#endif
+
 /*
  * Multifractal Singularity Spectrum of the Minkowski Question Mark Function
  *
@@ -166,6 +170,8 @@ __global__ void compute_tau(int num_q, double q_min, double q_step,
 
 /* ---- Host ---- */
 
+#ifndef TORCH_EXTENSION_NAME
+
 int main(int argc, char **argv) {
     int A_max = argc > 1 ? atoi(argv[1]) : 50;
     int N     = argc > 2 ? atoi(argv[2]) : 40;
@@ -318,3 +324,47 @@ int main(int argc, char **argv) {
     free(h_tau); free(h_q); free(h_alpha); free(h_f);
     return 0;
 }
+#endif
+
+
+#ifdef TORCH_EXTENSION_NAME
+std::vector<torch::Tensor> compute_spectrum(int64_t A_max, int64_t N) {
+    TORCH_CHECK(A_max >= 1 && A_max <= 100, "A_max must be 1-100");
+    TORCH_CHECK(N >= 1 && N <= 48, "N must be 1-48");
+
+    // Q_MIN, Q_MAX, Q_STEP, Q_COUNT are already defined as macros above
+    double *d_tau;
+    cudaMalloc(&d_tau, Q_COUNT * sizeof(double));
+
+    int threads = 32;
+    int blocks = (Q_COUNT + threads - 1) / threads;
+    compute_tau<<<blocks, threads>>>(Q_COUNT, Q_MIN, Q_STEP, (int)A_max, (int)N, d_tau);
+    cudaDeviceSynchronize();
+
+    auto tau = torch::zeros({Q_COUNT}, torch::dtype(torch::kFloat64));
+    cudaMemcpy(tau.data_ptr<double>(), d_tau, Q_COUNT * sizeof(double), cudaMemcpyDeviceToHost);
+    cudaFree(d_tau);
+
+    auto q = torch::zeros({Q_COUNT}, torch::dtype(torch::kFloat64));
+    auto alpha = torch::zeros({Q_COUNT}, torch::dtype(torch::kFloat64));
+    auto f_alpha = torch::zeros({Q_COUNT}, torch::dtype(torch::kFloat64));
+
+    double *pq = q.data_ptr<double>();
+    double *pt = tau.data_ptr<double>();
+    double *pa = alpha.data_ptr<double>();
+    double *pf = f_alpha.data_ptr<double>();
+
+    for (int i = 0; i < Q_COUNT; i++) pq[i] = Q_MIN + i * Q_STEP;
+
+    pa[0] = -(pt[1] - pt[0]) / Q_STEP;
+    for (int i = 1; i < Q_COUNT - 1; i++)
+        pa[i] = -(pt[i+1] - pt[i-1]) / (2.0 * Q_STEP);
+    pa[Q_COUNT-1] = -(pt[Q_COUNT-1] - pt[Q_COUNT-2]) / Q_STEP;
+
+    for (int i = 0; i < Q_COUNT; i++)
+        pf[i] = pq[i] * pa[i] + pt[i];
+
+    return {q, tau, alpha, f_alpha};
+}
+#endif
+

torch-ext/torch_binding.cpp

@@ -3,4 +3,7 @@
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.doc() = "Minkowski Question-Mark Singularity Spectrum CUDA kernel";
+  m.def("compute_spectrum", &compute_spectrum,
+        "Compute multifractal spectrum tau(q), alpha(q), f(alpha)",
+        py::arg("A_max"), py::arg("N"));
 }

torch-ext/torch_binding.h

@@ -1,3 +1,4 @@
 #pragma once
 #include <torch/torch.h>
-// See minkowski_spectrum/minkowski_spectrum.cu for kernel API
+
+std::vector<torch::Tensor> compute_spectrum(int64_t A_max, int64_t N);