Feat/op 22 linspace

csrc/aten/ExtendedOps.cpp

@@ -5,9 +5,24 @@
 #include <ATen/ATen.h>
 #include <torch/library.h>
 
+#include <cmath>
+
 namespace tvarant::ops {
+namespace {
 
-// ---- factory (#16–#18) ----
+int64_t arange_length(double start, double end, double step) {
+  TORCH_CHECK(step != 0, "step must be nonzero");
+  if (step > 0) {
+    TORCH_CHECK(end >= start, "upper bound and larger bound inconsistent with step sign");
+  } else {
+    TORCH_CHECK(end <= start, "upper bound and larger bound inconsistent with step sign");
+  }
+  return static_cast<int64_t>(std::ceil((end - start) / step));
+}
+
+}  // namespace
+
+// ---- factory (#16–#18, #21–#22) ----
 
 at::Tensor zeros(
     c10::IntArrayRef size,
@@ -95,6 +110,82 @@ at::Tensor full_like(
   return t;
 }
 
+at::Tensor arange_start_step(
+    const at::Scalar& start,
+    const at::Scalar& end,
+    const at::Scalar& step,
+    std::optional<c10::ScalarType> dtype_opt,
+    std::optional<c10::Layout> layout_opt,
+    std::optional<c10::Device> device_opt,
+    std::optional<bool> pin_memory_opt) {
+  const double s = start.toDouble();
+  const double e = end.toDouble();
+  const double st = step.toDouble();
+  const int64_t n = arange_length(s, e, st);
+  auto out = empty_memory_format(
+      {n}, dtype_opt, layout_opt, device_opt, pin_memory_opt, std::nullopt);
+  if (n == 0) {
+    return out;
+  }
+  if (out.scalar_type() == at::kFloat && out.is_contiguous()) {
+    LaunchParams p;
+    p.kernel = "arange_kernel";
+    p.dst = out.data_ptr();
+    p.scalar = static_cast<float>(s);
+    p.alpha = static_cast<float>(st);
+    p.numel = n;
+    ::tvarant::runtime().launch(p);
+    return out;
+  }
+  auto cpu = at::arange(s, e, st, at::TensorOptions().dtype(out.scalar_type()).device(at::kCPU));
+  return ::tvarant::ops::_copy_from(cpu, out, false);
+}
+
+at::Tensor arange(
+    const at::Scalar& end,
+    std::optional<c10::ScalarType> dtype_opt,
+    std::optional<c10::Layout> layout_opt,
+    std::optional<c10::Device> device_opt,
+    std::optional<bool> pin_memory_opt) {
+  return arange_start_step(0, end, 1, dtype_opt, layout_opt, device_opt, pin_memory_opt);
+}
+
+at::Tensor linspace(
+    const at::Scalar& start,
+    const at::Scalar& end,
+    int64_t steps,
+    std::optional<c10::ScalarType> dtype_opt,
+    std::optional<c10::Layout> layout_opt,
+    std::optional<c10::Device> device_opt,
+    std::optional<bool> pin_memory_opt) {
+  TORCH_CHECK(steps >= 0, "number of steps must be non-negative");
+  const double s = start.toDouble();
+  const double e = end.toDouble();
+  auto out = empty_memory_format(
+      {steps}, dtype_opt, layout_opt, device_opt, pin_memory_opt, std::nullopt);
+  if (steps == 0) {
+    return out;
+  }
+  if (steps == 1) {
+    fill__scalar(out, start);
+    return out;
+  }
+  const double step = (e - s) / static_cast<double>(steps - 1);
+  if (out.scalar_type() == at::kFloat && out.is_contiguous()) {
+    LaunchParams p;
+    p.kernel = "arange_kernel";
+    p.dst = out.data_ptr();
+    p.scalar = static_cast<float>(s);
+    p.alpha = static_cast<float>(step);
+    p.numel = steps;
+    ::tvarant::runtime().launch(p);
+    return out;
+  }
+  auto cpu = at::linspace(
+      s, e, steps, at::TensorOptions().dtype(out.scalar_type()).device(at::kCPU));
+  return ::tvarant::ops::_copy_from(cpu, out, false);
+}
+
 // ---- clone / contiguous (#19–#20) ----
 
 at::Tensor clone(
@@ -343,6 +434,38 @@ at::Tensor wrap_full_like(
       self, fill_value, dtype, layout, device, pin_memory, memory_format);
 }
 
+at::Tensor wrap_arange(
+    const at::Scalar& end,
+    std::optional<c10::ScalarType> dtype,
+    std::optional<c10::Layout> layout,
+    std::optional<c10::Device> device,
+    std::optional<bool> pin_memory) {
+  return tvarant::ops::arange(end, dtype, layout, device, pin_memory);
+}
+
+at::Tensor wrap_arange_start_step(
+    const at::Scalar& start,
+    const at::Scalar& end,
+    const at::Scalar& step,
+    std::optional<c10::ScalarType> dtype,
+    std::optional<c10::Layout> layout,
+    std::optional<c10::Device> device,
+    std::optional<bool> pin_memory) {
+  return tvarant::ops::arange_start_step(
+      start, end, step, dtype, layout, device, pin_memory);
+}
+
+at::Tensor wrap_linspace(
+    const at::Scalar& start,
+    const at::Scalar& end,
+    int64_t steps,
+    std::optional<c10::ScalarType> dtype,
+    std::optional<c10::Layout> layout,
+    std::optional<c10::Device> device,
+    std::optional<bool> pin_memory) {
+  return tvarant::ops::linspace(start, end, steps, dtype, layout, device, pin_memory);
+}
+
 at::Tensor wrap_clone(const at::Tensor& self, std::optional<c10::MemoryFormat> memory_format) {
   return tvarant::ops::clone(self, memory_format);
 }
@@ -366,6 +489,9 @@ TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
   m.impl("ones_like", TORCH_FN(wrap_ones_like));
   m.impl("full", TORCH_FN(wrap_full));
   m.impl("full_like", TORCH_FN(wrap_full_like));
+  m.impl("arange", TORCH_FN(wrap_arange));
+  m.impl("arange.start_step", TORCH_FN(wrap_arange_start_step));
+  m.impl("linspace", TORCH_FN(wrap_linspace));
   m.impl("clone", TORCH_FN(wrap_clone));
   m.impl("contiguous", TORCH_FN(wrap_contiguous));
   m.impl("neg", TORCH_FN(tvarant::ops::neg));

csrc/kernels/host/HostKernels.cpp

@@ -20,6 +20,12 @@ void fill_f32(float* out, float value, int64_t n) {
   }
 }
 
+void arange_f32(float* out, int64_t n, float start, float step) {
+  for (int64_t i = 0; i < n; ++i) {
+    out[i] = start + static_cast<float>(i) * step;
+  }
+}
+
 void copy_f32(float* dst, const float* src, int64_t n) {
   std::memcpy(dst, src, static_cast<size_t>(n) * sizeof(float));
 }

csrc/kernels/host/HostKernels.h

@@ -5,6 +5,7 @@
 namespace tvarant::host {
 
 void fill_f32(float* out, float value, int64_t n);
+void arange_f32(float* out, int64_t n, float start, float step);
 void copy_f32(float* dst, const float* src, int64_t n);
 void add_f32(const float* a, const float* b, float* out, int64_t n, float alpha);
 void mul_f32(const float* a, const float* b, float* out, int64_t n);

csrc/kernels/opencl/elementwise.cl

@@ -1,3 +1,11 @@
+__kernel void arange_kernel(
+    __global float* out, const float start, const float step, const int n) {
+  const int i = get_global_id(0);
+  if (i < n) {
+    out[i] = start + (float)i * step;
+  }
+}
+
 __kernel void silu_kernel(__global const float* in, __global float* out, const int n) {
   const int i = get_global_id(0);
   if (i < n) {

csrc/runtime/CpuSimRuntime.cpp

@@ -111,6 +111,9 @@ class CpuSimRuntime final : public TvarantRuntime {
     const std::string name(p.kernel);
     if (name == "fill_kernel") {
       host::fill_f32(static_cast<float*>(p.dst), p.scalar, p.numel);
+    } else if (name == "arange_kernel") {
+      host::arange_f32(
+          static_cast<float*>(p.dst), p.numel, p.scalar, p.alpha);
     } else if (name == "copy_kernel") {
       host::copy_f32(
           static_cast<float*>(p.dst), static_cast<const float*>(p.src0), p.numel);

csrc/runtime/OpenCLRuntime.cpp

@@ -34,6 +34,10 @@ __kernel void fill_kernel(__global float* out, const float value, const int n) {
   const int i = get_global_id(0);
   if (i < n) out[i] = value;
 }
+__kernel void arange_kernel(__global float* out, const float start, const float step, const int n) {
+  const int i = get_global_id(0);
+  if (i < n) out[i] = start + (float)i * step;
+}
 __kernel void copy_kernel(__global float* dst, __global const float* src, const int n) {
   const int i = get_global_id(0);
   if (i < n) dst[i] = src[i];
@@ -256,6 +260,7 @@ class OpenCLRuntime final : public TvarantRuntime {
       kernels_[name] = k;
     };
     load("fill_kernel");
+    load("arange_kernel");
     load("copy_kernel");
     load("add_kernel");
     load("mul_kernel");
@@ -405,6 +410,11 @@ class OpenCLRuntime final : public TvarantRuntime {
       set_buf(p.dst);
       set_float(p.scalar);
       set_int(n);
+    } else if (name == "arange_kernel") {
+      set_buf(p.dst);
+      set_float(p.scalar);
+      set_float(p.alpha);
+      set_int(n);
     } else if (name == "copy_kernel") {
       set_buf(p.dst);
       set_buf(const_cast<void*>(p.src0));

tests/test_op_arange.py

@@ -0,0 +1,19 @@
+import torch
+import torch_tvarant  # noqa: F401
+
+
+def test_arange_end():
+    a = torch.arange(5, device="tvarant")
+    assert a.device.type == "tvarant"
+    assert torch.equal(a.cpu(), torch.arange(5))
+
+
+def test_arange_start_step():
+    a = torch.arange(2, 10, 2, device="tvarant")
+    assert torch.equal(a.cpu(), torch.arange(2, 10, 2))
+
+
+def test_arange_float():
+    a = torch.arange(0.0, 1.0, 0.25, device="tvarant", dtype=torch.float32)
+    ref = torch.arange(0.0, 1.0, 0.25, dtype=torch.float32)
+    assert torch.allclose(a.cpu(), ref, atol=1e-6)

tests/test_op_linspace.py

@@ -0,0 +1,21 @@
+import torch
+import torch_tvarant  # noqa: F401
+
+
+def test_linspace_basic():
+    a = torch.linspace(0, 1, 5, device="tvarant")
+    assert a.device.type == "tvarant"
+    ref = torch.linspace(0, 1, 5)
+    assert torch.allclose(a.cpu(), ref, atol=1e-6)
+
+
+def test_linspace_single_step():
+    a = torch.linspace(3.0, 7.0, 1, device="tvarant", dtype=torch.float32)
+    ref = torch.linspace(3.0, 7.0, 1, dtype=torch.float32)
+    assert torch.allclose(a.cpu(), ref, atol=1e-6)
+
+
+def test_linspace_empty():
+    a = torch.linspace(0, 1, 0, device="tvarant")
+    assert a.numel() == 0
+    assert a.device.type == "tvarant"