Arange is now like the other kernels

Author: txus

src/tensor/cpu/ops.cpp

@@ -598,13 +598,15 @@ template Tensor<bfloat16, CPU> cat(const TensorView<bfloat16, CPU>&,
                                    const TensorView<bfloat16, CPU>&, int);
 template Tensor<float, CPU> cat(const TensorView<float, CPU>&, const TensorView<float, CPU>&, int);
 
+// elementwise maps
 template Tensor<float, CPU> cos(const TensorView<float, CPU>& tensor);
 template Tensor<float, CPU> sin(const TensorView<float, CPU>& tensor);
 template Tensor<float, CPU> exp(const TensorView<float, CPU>& tensor);
 
 template Tensor<bfloat16, CPU> pow(bfloat16, const TensorView<bfloat16, CPU>&);
 template Tensor<float, CPU> pow(const TensorView<float, CPU>&, float);
 template Tensor<float, CPU> pow(float, const TensorView<float, CPU>&);
+
 template Tensor<bfloat16, CPU> tril(const TensorView<bfloat16, CPU>&, bool);
 template Tensor<int, CPU> tril(const TensorView<int, CPU>&, bool);
 template Tensor<bfloat16, CPU> slice(const TensorView<bfloat16, CPU>&, int, size_t, size_t);

src/tensor/cuda/kernels/arange.cu

@@ -1,5 +1,5 @@
 #include "arange.cuh"
-#include <cstddef>
+#include "utils.cuh"
 
 namespace tensor::kernels {
 
@@ -14,8 +14,31 @@ __global__ void arange_kernel(DeviceT* out, DeviceT start, DeviceT end, DeviceT
   }
 }
 
-template __global__ void arange_kernel<Cuda<float>>(Cuda<float>*, Cuda<float>, Cuda<float>, Cuda<float>, size_t);
-template __global__ void arange_kernel<Cuda<int>>(Cuda<int>*, Cuda<int>, Cuda<int>, Cuda<int>, size_t);
-template __global__ void arange_kernel<Cuda<bfloat16>>(Cuda<bfloat16>*, Cuda<bfloat16>, Cuda<bfloat16>, Cuda<bfloat16>, size_t);
+template <typename T> Tensor<T, CUDA> arange(T start, T end, T step) {
+  auto n_elements = static_cast<size_t>((end - start) / step);
+
+  TensorStorage<T, CUDA> storage(n_elements);
+
+  Shape shape{n_elements};
+
+  Tensor<T, CUDA> out{shape, std::move(storage)};
+
+  size_t block_size = cuda::get_block_size(n_elements);
+  size_t grid_size = cuda::get_grid_size(n_elements, block_size);
+
+  // Convert to device-native types for kernel call
+  auto* device_data = reinterpret_cast<Cuda<T>*>(out.data()); // NOLINT
+  Cuda<T> device_start = to_device_type(start, CUDA{});
+  Cuda<T> device_end = to_device_type(end, CUDA{});
+  Cuda<T> device_step = to_device_type(step, CUDA{});
+
+  arange_kernel<Cuda<T>><<<grid_size, block_size>>>(device_data, device_start, device_end, device_step, n_elements);
+
+  return out;
+}
+
+template Tensor<bfloat16, CUDA> arange(bfloat16 start, bfloat16 end, bfloat16 step);
+template Tensor<float, CUDA> arange(float start, float end, float step);
+template Tensor<int, CUDA> arange(int start, int end, int step);
 
 } // namespace tensor::kernels

src/tensor/cuda/kernels/arange.cuh

@@ -1,18 +1,12 @@
 #pragma once
 
-#include <cuda_runtime.h>
+#include <tensor/tensor.hpp>
 #include <tensor/device_type.hpp>
-#include <cstddef>
 
 namespace tensor::kernels {
 
 using namespace dtype;
 
-template<typename DeviceT>
-__global__ void arange_kernel(DeviceT* out, DeviceT start, DeviceT end, DeviceT step, size_t n);
-
-extern template __global__ void arange_kernel<Cuda<float>>(Cuda<float>*, Cuda<float>, Cuda<float>, Cuda<float>, size_t);
-extern template __global__ void arange_kernel<Cuda<int>>(Cuda<int>*, Cuda<int>, Cuda<int>, Cuda<int>, size_t);
-extern template __global__ void arange_kernel<Cuda<bfloat16>>(Cuda<bfloat16>*, Cuda<bfloat16>, Cuda<bfloat16>, Cuda<bfloat16>, size_t);
+template <typename T> Tensor<T, CUDA> arange(T start, T end, T step);
 
 } // namespace tensor::kernels

src/tensor/cuda/ops.cu

@@ -20,32 +20,15 @@ namespace tensor {
 using namespace dtype;
 using namespace device;
 
-template <typename T, typename D> Tensor<T, D> arange(T start, T end, T step) {
-  auto n_elements = static_cast<size_t>((end - start) / step);
-
-  TensorStorage<T, D> storage(n_elements);
-
-  Shape shape{n_elements};
-
-  Tensor<T, D> out{shape, std::move(storage)};
-
-  size_t block_size = cuda::get_block_size(n_elements);
-  size_t grid_size = cuda::get_grid_size(n_elements, block_size);
-
-  // Convert to device-native types for kernel call
-  auto* device_data = reinterpret_cast<Cuda<T>*>(out.data()); // NOLINT
-  Cuda<T> device_start = to_device_type(start, D{});
-  Cuda<T> device_end = to_device_type(end, D{});
-  Cuda<T> device_step = to_device_type(step, D{});
-
-  kernels::arange_kernel<<<grid_size, block_size>>>(device_data, device_start, device_end, device_step, n_elements);
-
-  return out;
+template <> Tensor<bfloat16, CUDA> arange(bfloat16 start, bfloat16 end, bfloat16 step) {
+  return kernels::arange(start, end, step);
+}
+template <> Tensor<float, CUDA> arange(float start, float end, float step) {
+  return kernels::arange(start, end, step);
+}
+template <> Tensor<int, CUDA> arange(int start, int end, int step) {
+  return kernels::arange(start, end, step);
 }
-
-template Tensor<int, CUDA> arange(int start, int end, int step);
-template Tensor<float, CUDA> arange(float start, float end, float step);
-template Tensor<bfloat16, CUDA> arange(bfloat16 start, bfloat16 end, bfloat16 step);
 
 template <typename T, typename D>
 void replace_from_(Tensor<T, D>& out, const TensorView<T, D>& input) {