Fix cmake, add relu CPU impl

Author: danieldk

build2cmake/src/templates/cpu/kernel.cmake

@@ -21,4 +21,4 @@ endforeach()
 {% endif %}
 
 # Add C++ sources to main source list
-list(APPEND SRC {{'"${' + kernel_name + '_CPP_SRC}"'}})
+list(APPEND SRC {{'"${' + kernel_name + '_SRC}"'}})

examples/relu/relu_cpu/relu_cpu.cpp

@@ -0,0 +1,56 @@
+#include <torch/all.h>
+
+#ifdef __SSE__
+#include <xmmintrin.h>
+#endif
+
+#ifdef __ARM_NEON
+#include <arm_neon.h>
+#endif
+
+#ifdef __SSE__
+void relu_forward_sse(float* out, const float* input, size_t size) {
+    size_t i = 0;
+
+    for (; i + 4 <= size; i += 4) {
+        __m128 vec_input = _mm_load_ps(input + i);
+        __m128 vec_zero = _mm_setzero_ps();
+        __m128 vec_output = _mm_max_ps(vec_input, vec_zero);
+        _mm_store_ps(out + i, vec_output);
+    }
+
+    for (; i < size; ++i) {
+        out[i] = input[i] > 0 ? input[i] : 0;
+    }
+}
+#endif
+
+#ifdef __ARM_NEON
+void relu_forward_neon(float* out, const float* input, size_t size) {
+    size_t i = 0;
+
+    for (; i + 4 <= size; i += 4) {
+        float32x4_t vec_input = vld1q_f32(input + i);
+        float32x4_t vec_output = vmaxq_f32(vec_input, vdupq_n_f32(0));
+        vst1q_f32(out + i, vec_output);
+    }
+
+    for (; i < size; ++i) {
+        out[i] = input[i] > 0 ? input[i] : 0;
+    }
+}
+#endif
+
+void relu(torch::Tensor &out, torch::Tensor const &input) {
+    TORCH_CHECK(out.dtype() == torch::kFloat32, "Output tensor must be of dtype float");
+    TORCH_CHECK(input.dtype() == torch::kFloat32, "Input tensor must be of dtype float");
+    TORCH_CHECK(out.numel() == input.numel(), "Input and output tensors must have the same number of elements");
+
+#if defined(__SSE__)
+    relu_forward_sse(out.data_ptr<float>(), input.data_ptr<float>(), input.numel());
+#elif defined(__ARM_NEON)
+    relu_forward_neon(out.data_ptr<float>(), input.data_ptr<float>(), input.numel());
+#else
+    #error "Unsupported architecture; please use a CPU with SSE or ARM NEON support."
+#endif
+}

examples/relu/tests/test_relu.py

@@ -11,7 +11,9 @@ def test_relu():
         device = torch.device("mps")
     elif hasattr(torch, "xpu") and torch.xpu.is_available():
         device = torch.device("xpu")
-    else:
+    elif torch.version.cuda is not None and torch.cuda.is_available():
         device = torch.device("cuda")
+    else:
+        device = torch.device("cpu")
     x = torch.randn(1024, 1024, dtype=torch.float32, device=device)
     torch.testing.assert_allclose(F.relu(x), relu.relu(x))

examples/relu/torch-ext/torch_binding.cpp

@@ -6,7 +6,7 @@
 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("relu(Tensor! out, Tensor input) -> ()");
 #if defined(CPU_KERNEL)
-  ops.impl("relu", torch::kCUDA, &relu);
+  ops.impl("relu", torch::kCPU, &relu);
 #elif defined(CUDA_KERNEL) || defined(ROCM_KERNEL)
   ops.impl("relu", torch::kCUDA, &relu);
 #elif defined(METAL_KERNEL)

overlay.nix

@@ -25,8 +25,8 @@ final: prev: {
           src = final.fetchFromGitHub {
             owner = "huggingface";
             repo = "kernels";
-            rev = "5d21b86a5d611100c10c10b79ffa7965edf567fd";
-            sha256 = "sha256-lKQUVbjhpeXKj1SeZRxgPSsOtBUZ7zQeO6pRoA1h+W8=";
+            rev = "0e18dbf076fc44de5dac4027616e9f3d9e2da45a";
+            sha256 = "sha256-6N1W3jLQIS1yEAdNR2X9CuFdMw4Ia0yzBBVQ4Kujv8U=";
           };
         });
       }