Sync with kernel-builder changes (#1)

- Track dynamic libraries (.so) (257ce3a49add504c1447640ebd7e59d2de950e40)
- Fix function signature (071b5ee9ef6040dfb5a82fea655c60a592218195)
- Make it a flake (93c0f1f9eb0b05fdd14d4c9efa6bf1a2c3269925)
- Build (89d482b71f65ba545b8c24488fc18ac882b2bfc0)


Co-authored-by: Daniël de Kok <danieldk@users.noreply.huggingface.co>

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.so filter=lfs diff=lfs merge=lfs -text

build.toml

@@ -8,7 +8,7 @@ src = [
   "ext-torch/torch_binding.cpp",
   "ext-torch/torch_binding.h",
 ]
-pysrc = ["ext-torch/__init__.py"]
+pyroot = "ext-torch"
 
 [kernel.mel_spectrogram]
 capabilities = ["7.0", "7.2", "7.5", "8.0", "8.6", "8.7", "8.9", "9.0"]

ext-torch/torch_binding.cpp

@@ -5,7 +5,7 @@
 
 
 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
-  ops.def("mel_spectrogram(Tensor out, Tensor samples, Tensor filters, int64_t fft_size, int64_t fft_step) -> Tensor", &mel_spectrogram);
+  ops.def("mel_spectrogram(Tensor! out, Tensor samples, Tensor filters, int fft_size, int fft_step) -> ()", &mel_spectrogram);
   ops.impl("mel_spectrogram", torch::kCUDA, &mel_spectrogram);
 }
 

ext-torch/torch_binding.h

@@ -3,8 +3,8 @@
 #include <torch/torch.h>
 
 void mel_spectrogram(
-    torch::Tensor out,
-    torch::Tensor samples,
-    torch::Tensor filters,
+    torch::Tensor &out,
+    torch::Tensor const &samples,
+    torch::Tensor const &filters,
     int64_t fft_size,
-    int64_t fft_step);
+    int64_t fft_step);

flake.nix

@@ -0,0 +1,14 @@
+{
+  description = "Flake for mel_spectogram";
+
+  inputs = {
+    kernel-builder.url = "git+ssh://git@github.com/huggingface/kernel-builder";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genFlakeOutputs ./.;
+}

mel_spectrogram/mel_spectrogram.cu

@@ -100,11 +100,11 @@ __global__ void normalize_mel_spec_kernel(
 }
 
 void mel_spectrogram(
-    torch::Tensor& out,        // [n_mel, n_frames]
-    torch::Tensor& samples,    // [n_samples]
-    torch::Tensor& filters,    // [n_mel, n_fft]
-    const int fft_size,
-    const int fft_step) {
+    torch::Tensor &out,        // [n_mel, n_frames]
+    torch::Tensor const &samples,    // [n_samples]
+    torch::Tensor const &filters,    // [n_mel, n_fft]
+    const int64_t fft_size,
+    const int64_t fft_step) {
     
     TORCH_CHECK(samples.is_cuda(), "Input samples must be a CUDA tensor");
     TORCH_CHECK(filters.is_cuda(), "Mel filters must be a CUDA tensor");
@@ -160,4 +160,4 @@ void mel_spectrogram(
         n_frames);
     
     cufftDestroy(plan);
-}
+}