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

ramanujan/ramanujan_v2.cu

@@ -1,3 +1,7 @@
+#ifdef TORCH_EXTENSION_NAME
+#include <torch/torch.h>
+#endif
+
 /*
  * Ramanujan Machine v2: ASYMMETRIC-DEGREE polynomial CF search
  *
@@ -332,6 +336,8 @@ __global__ void search_kernel(
 
 /* ── Main ──────────────────────────────────────────────────── */
 
+#ifndef TORCH_EXTENSION_NAME
+
 int main(int argc, char **argv) {
     if (argc < 5) {
         printf("Usage: %s <deg_a> <deg_b> <range_a> <range_b> [cf_depth] [gpu_id]\n", argv[0]);
@@ -534,3 +540,66 @@ int main(int argc, char **argv) {
     cudaFree(d_hit_count); cudaFree(d_unmatched_count);
     return 0;
 }
+#endif
+
+
+#ifdef TORCH_EXTENSION_NAME
+std::vector<torch::Tensor> search(int64_t deg_a, int64_t deg_b,
+                                   int64_t range_a, int64_t range_b,
+                                   int64_t cf_depth) {
+    int max_hits = 500000, max_unmatched = 1000000;
+    Hit *d_hits, *d_unmatched;
+    int *d_hit_count, *d_unmatched_count;
+    cudaMalloc(&d_hits, max_hits * sizeof(Hit));
+    cudaMalloc(&d_unmatched, max_unmatched * sizeof(Hit));
+    cudaMalloc(&d_hit_count, sizeof(int));
+    cudaMalloc(&d_unmatched_count, sizeof(int));
+    cudaMemset(d_hit_count, 0, sizeof(int));
+    cudaMemset(d_unmatched_count, 0, sizeof(int));
+
+    long long width_a = 2 * range_a + 1;
+    long long width_b = 2 * range_b + 1;
+    long long total_a = 1, total_b = 1;
+    for (int i = 0; i <= deg_a; i++) total_a *= width_a;
+    for (int i = 0; i <= deg_b; i++) total_b *= width_b;
+    long long total = total_a * total_b;
+
+    long long chunk = 1000000;
+    for (long long start = 0; start < total; start += chunk) {
+        long long count = std::min(chunk, total - start);
+        int blocks = ((int)count + 255) / 256;
+        search_kernel<<<blocks, 256>>>(start, count,
+            (int)deg_a, (int)deg_b, (int)range_a, (int)range_b, (int)cf_depth,
+            d_hits, d_hit_count, max_hits,
+            d_unmatched, d_unmatched_count, max_unmatched);
+        cudaDeviceSynchronize();
+    }
+
+    int h_hc, h_uc;
+    cudaMemcpy(&h_hc, d_hit_count, sizeof(int), cudaMemcpyDeviceToHost);
+    cudaMemcpy(&h_uc, d_unmatched_count, sizeof(int), cudaMemcpyDeviceToHost);
+    h_hc = std::min(h_hc, max_hits);
+    h_uc = std::min(h_uc, max_unmatched);
+
+    auto hit_vals = torch::zeros({h_hc}, torch::kFloat64);
+    auto hit_match = torch::zeros({h_hc}, torch::kInt32);
+    if (h_hc > 0) {
+        std::vector<Hit> hh(h_hc);
+        cudaMemcpy(hh.data(), d_hits, h_hc * sizeof(Hit), cudaMemcpyDeviceToHost);
+        for (int i = 0; i < h_hc; i++) {
+            hit_vals.data_ptr<double>()[i] = hh[i].value;
+            hit_match.data_ptr<int>()[i] = hh[i].match_const;
+        }
+    }
+    auto um_vals = torch::zeros({h_uc}, torch::kFloat64);
+    if (h_uc > 0) {
+        std::vector<Hit> uh(h_uc);
+        cudaMemcpy(uh.data(), d_unmatched, h_uc * sizeof(Hit), cudaMemcpyDeviceToHost);
+        for (int i = 0; i < h_uc; i++) um_vals.data_ptr<double>()[i] = uh[i].value;
+    }
+
+    cudaFree(d_hits); cudaFree(d_unmatched);
+    cudaFree(d_hit_count); cudaFree(d_unmatched_count);
+    return {hit_vals, hit_match, um_vals};
+}
+#endif

torch-ext/torch_binding.cpp

@@ -3,4 +3,6 @@
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.doc() = "Ramanujan Machine v2 (Asymmetric-Degree CF Search) CUDA kernel";
+  m.def("search", &search, py::arg("deg_a"), py::arg("deg_b"),
+        py::arg("range_a"), py::arg("range_b"), py::arg("cf_depth") = 300);
 }

torch-ext/torch_binding.h

@@ -1,3 +1,6 @@
 #pragma once
 #include <torch/torch.h>
-// See ramanujan/ramanujan_v2.cu for kernel API
+
+std::vector<torch::Tensor> search(int64_t deg_a, int64_t deg_b,
+                                   int64_t range_a, int64_t range_b,
+                                   int64_t cf_depth);