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

hausdorff_spectrum/hausdorff_spectrum.cu

@@ -1,3 +1,7 @@
+#ifdef TORCH_EXTENSION_NAME
+#include <torch/torch.h>
+#endif
+
 /*
  * Hausdorff Dimension Spectrum of Continued Fraction Cantor Sets
  *
@@ -200,6 +204,8 @@ void format_subset(uint32_t mask, int max_d, char *buf, int buflen) {
  * Host: main
  * ============================================================ */
 
+#ifndef TORCH_EXTENSION_NAME
+
 int main(int argc, char **argv) {
     int max_d = argc > 1 ? atoi(argv[1]) : 10;
     int N     = argc > 2 ? atoi(argv[2]) : 40;
@@ -384,3 +390,29 @@ int main(int argc, char **argv) {
 
     return 0;
 }
+#endif
+
+
+#ifdef TORCH_EXTENSION_NAME
+torch::Tensor hausdorff_all(int64_t max_d, int64_t N) {
+    TORCH_CHECK(max_d >= 1 && max_d <= 24, "max_d must be 1-24");
+    TORCH_CHECK(N >= 1 && N <= 48, "N must be 1-48");
+
+    int total = (1 << max_d) - 1;
+    auto results = torch::zeros({total}, torch::dtype(torch::kFloat64));
+    double *d_results;
+    cudaMalloc(&d_results, 1024 * sizeof(double));
+
+    for (int offset = 0; offset < total; offset += 1024) {
+        int batch = std::min(1024, total - offset);
+        uint32_t start_mask = offset + 1;
+        batch_hausdorff<<<batch, 1>>>(start_mask, batch, (int)max_d, (int)N, d_results);
+        cudaDeviceSynchronize();
+        cudaMemcpy(results.data_ptr<double>() + offset, d_results,
+                   batch * sizeof(double), cudaMemcpyDeviceToHost);
+    }
+    cudaFree(d_results);
+    return results;
+}
+#endif
+

torch-ext/torch_binding.cpp

@@ -3,4 +3,7 @@
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.doc() = "Hausdorff Dimension Spectrum CUDA kernel";
+  m.def("hausdorff_all", &hausdorff_all,
+        "Compute Hausdorff dimensions for all subsets of {1,...,max_d}",
+        py::arg("max_d"), py::arg("N"));
 }

torch-ext/torch_binding.h

@@ -1,3 +1,4 @@
 #pragma once
 #include <torch/torch.h>
-// See hausdorff_spectrum/hausdorff_spectrum.cu for kernel API
+
+torch::Tensor hausdorff_all(int64_t max_d, int64_t N);