diff --git a/csrc/aten/ExtendedOps.cpp b/csrc/aten/ExtendedOps.cpp index 046f06d..6dc9b61 100644 --- a/csrc/aten/ExtendedOps.cpp +++ b/csrc/aten/ExtendedOps.cpp @@ -5,9 +5,24 @@ #include #include +#include + 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(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 dtype_opt, + std::optional layout_opt, + std::optional device_opt, + std::optional 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(s); + p.alpha = static_cast(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 dtype_opt, + std::optional layout_opt, + std::optional device_opt, + std::optional 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 dtype_opt, + std::optional layout_opt, + std::optional device_opt, + std::optional 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(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(s); + p.alpha = static_cast(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 dtype, + std::optional layout, + std::optional device, + std::optional 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 dtype, + std::optional layout, + std::optional device, + std::optional 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 dtype, + std::optional layout, + std::optional device, + std::optional pin_memory) { + return tvarant::ops::linspace(start, end, steps, dtype, layout, device, pin_memory); +} + at::Tensor wrap_clone(const at::Tensor& self, std::optional 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)); diff --git a/csrc/kernels/host/HostKernels.cpp b/csrc/kernels/host/HostKernels.cpp index f617852..6186e05 100644 --- a/csrc/kernels/host/HostKernels.cpp +++ b/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(i) * step; + } +} + void copy_f32(float* dst, const float* src, int64_t n) { std::memcpy(dst, src, static_cast(n) * sizeof(float)); } diff --git a/csrc/kernels/host/HostKernels.h b/csrc/kernels/host/HostKernels.h index 2cffb4c..f4a3cef 100644 --- a/csrc/kernels/host/HostKernels.h +++ b/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); diff --git a/csrc/kernels/opencl/elementwise.cl b/csrc/kernels/opencl/elementwise.cl index 4fa48a2..a44d495 100644 --- a/csrc/kernels/opencl/elementwise.cl +++ b/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) { diff --git a/csrc/runtime/CpuSimRuntime.cpp b/csrc/runtime/CpuSimRuntime.cpp index 9cdf0cb..42a63c3 100644 --- a/csrc/runtime/CpuSimRuntime.cpp +++ b/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(p.dst), p.scalar, p.numel); + } else if (name == "arange_kernel") { + host::arange_f32( + static_cast(p.dst), p.numel, p.scalar, p.alpha); } else if (name == "copy_kernel") { host::copy_f32( static_cast(p.dst), static_cast(p.src0), p.numel); diff --git a/csrc/runtime/OpenCLRuntime.cpp b/csrc/runtime/OpenCLRuntime.cpp index 247ca27..7b3d4e7 100644 --- a/csrc/runtime/OpenCLRuntime.cpp +++ b/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(p.src0)); diff --git a/tests/test_op_arange.py b/tests/test_op_arange.py new file mode 100644 index 0000000..d897de7 --- /dev/null +++ b/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) diff --git a/tests/test_op_linspace.py b/tests/test_op_linspace.py new file mode 100644 index 0000000..9d16bf4 --- /dev/null +++ b/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"