slice is the opposite of cat

Author: txus

src/tensor/cpu/ops.cpp

@@ -607,12 +607,16 @@ 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>&);
 
+// tril
 template Tensor<bfloat16, CPU> tril(const TensorView<bfloat16, CPU>&, bool);
 template Tensor<int, CPU> tril(const TensorView<int, CPU>&, bool);
+
+// slice
 template Tensor<bfloat16, CPU> slice(const TensorView<bfloat16, CPU>&, int, size_t, size_t);
 template Tensor<float, CPU> slice(const TensorView<float, CPU>&, int, size_t, size_t);
 template Tensor<float, CPU> slice(const TensorView<const float, CPU>&, int, size_t, size_t);
-template Tensor<int, CPU> slice(const TensorView<int, CPU>&, int, size_t, size_t);
+
+// matmul
 template Tensor<bfloat16, CPU> matmul(const TensorView<bfloat16, CPU>&,
                                       const TensorView<bfloat16, CPU>&);
 template Tensor<bfloat16, CPU> matmul(const TensorView<bfloat16, CPU>&,

src/tensor/cuda/CMakeLists.txt

@@ -1,6 +1,6 @@
 file(GLOB HEADER_LIST CONFIGURE_DEPENDS "${PROJECT_SOURCE_DIR}/include/tensor/*.hpp")
 
-add_library(tensor_cuda STATIC storage.cu loader.cu ops.cu kernels/fill.cu kernels/argmax.cu kernels/arange.cu kernels/sum.cu kernels/max.cu kernels/masked_fill.cu kernels/cat.cu kernels/map.cu kernels/zip.cu kernels/tril.cu)
+add_library(tensor_cuda STATIC storage.cu loader.cu ops.cu kernels/fill.cu kernels/argmax.cu kernels/arange.cu kernels/sum.cu kernels/max.cu kernels/masked_fill.cu kernels/cat.cu kernels/map.cu kernels/zip.cu kernels/tril.cu kernels/slice.cu)
 
 set_target_properties(tensor_cuda PROPERTIES
   CUDA_SEPARABLE_COMPILATION ON

src/tensor/cuda/kernels/slice.cu

@@ -0,0 +1,76 @@
+#include "slice.cuh"
+#include "utils.cuh"
+#include <cstddef>
+#include <tensor/device_type.hpp>
+#include <cuda_runtime.h>
+
+namespace tensor::kernels {
+
+using namespace dtype;
+
+template <typename T>
+__global__ void slice_kernel(Cuda<T>* out, const Cuda<T>* input, size_t start_offset, size_t chunk_size, size_t source_stride) {
+  size_t operation_idx = blockIdx.x;
+
+  auto in_base = (operation_idx * source_stride) + start_offset;
+  auto out_base = operation_idx * chunk_size;
+
+  for (size_t element = threadIdx.x; element < chunk_size; element += blockDim.x) {
+    out[out_base + element] = input[in_base + element];
+  }
+}
+
+template <typename T>
+Tensor<T, CUDA> slice(const TensorView<T, CUDA>& view, int dim, size_t start, size_t end) {
+  assert(view.is_contiguous() && "tensor should be contiguous");
+
+  auto shape = view.shape;
+
+  if (dim < 0) {
+    dim = static_cast<int>(shape.size()) + dim;
+  }
+
+  Shape new_shape{shape};
+  new_shape[dim] = end - start;
+
+  // product of all dimensions after dim
+  size_t inner_stride = 1;
+  for (size_t idx = dim + 1; idx < shape.size(); ++idx) {
+    inner_stride *= shape[idx];
+  }
+
+  // product of all dimensions before dim
+  size_t outer_iterations = 1;
+  for (size_t idx = 0; idx < static_cast<size_t>(dim); ++idx) {
+    outer_iterations *= shape[idx];
+  }
+
+  size_t source_stride = shape[dim] * inner_stride;
+  size_t chunk_size = (end - start) * inner_stride;
+  size_t start_offset = start * inner_stride;
+
+  size_t n_elements = outer_iterations * chunk_size;
+  TensorStorage<T, CUDA> storage(n_elements);
+  Tensor<T, CUDA> out{new_shape, std::move(storage)};
+
+  // fast path: if slicing on first dimension, just use cudaMemcpy
+  if (dim == 0) {
+    size_t bytes = n_elements * sizeof(T);
+    CUDA_CHECK(cudaMemcpy(out.data(), view.data + start_offset, bytes, cudaMemcpyDeviceToDevice)); // NOLINT
+    return out;
+  }
+
+  size_t block_size = cuda::get_block_size(chunk_size);
+
+  auto* out_d = reinterpret_cast<Cuda<T>*>(out.data()); // NOLINT
+  auto* in_d = reinterpret_cast<const Cuda<T>*>(view.data); // NOLINT
+
+  slice_kernel<T><<<outer_iterations, block_size>>>(out_d, in_d, start_offset, chunk_size, source_stride);
+
+  return out;
+}
+
+template Tensor<bfloat16, CUDA> slice(const TensorView<bfloat16, CUDA>& view, int dim, size_t start, size_t end);
+template Tensor<float, CUDA> slice(const TensorView<float, CUDA>& view, int dim, size_t start, size_t end);
+
+} // namespace tensor::kernels

src/tensor/cuda/kernels/slice.cuh

@@ -0,0 +1,10 @@
+#pragma once
+
+#include <tensor/tensor.hpp>
+
+namespace tensor::kernels {
+
+template<typename T>
+Tensor<T, CUDA> slice(const TensorView<T, CUDA>& view, int dim, size_t start, size_t end);
+
+} // namespace tensor::kernels

src/tensor/cuda/kernels/tril.cu

@@ -51,8 +51,6 @@ template <typename T> Tensor<T, CUDA> tril(const TensorView<T, CUDA>& tensor, bo
   auto* in_d = reinterpret_cast<Cuda<T>*>(tensor.data); // NOLINT
   Cuda<T> diagonal_d = to_device_type(diagonal, CUDA{});
 
-  fmt::println("Grid: {} by {}", grid_size, block_size);
-
   tril_kernel<Cuda<T>><<<grid_size, block_size>>>(device_data, in_d, diagonal_d, cols, rows);
 
   return out;

src/tensor/cuda/ops.cu

@@ -14,6 +14,7 @@
 #include "kernels/map.cuh"
 #include "kernels/zip.cuh"
 #include "kernels/tril.cuh"
+#include "kernels/slice.cuh"
 #include "kernels/utils.cuh"
 
 namespace tensor {
@@ -135,4 +136,14 @@ Tensor<int, CUDA> tril(const TensorView<int, CUDA>& tensor, bool diagonal) {
   return kernels::tril(tensor, diagonal);
 }
 
+template <>
+Tensor<bfloat16, CUDA> slice(const TensorView<bfloat16, CUDA>& view, int dim, size_t start, size_t end) {
+  return kernels::slice(view, dim, start, end);
+}
+
+template <>
+Tensor<float, CUDA> slice(const TensorView<float, CUDA>& view, int dim, size_t start, size_t end) {
+  return kernels::slice(view, dim, start, end);
+}
+
 } // namespace tensor

tests/tensor/cuda/test_ops.cpp

@@ -309,3 +309,75 @@ TEST(TensorCUDATest, TrilBf16) {
   exp = {1, 1, 0, 0, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1};
   tensor_is_close<bfloat16>(diag.span(), std::span(exp));
 }
+
+TEST(TensorCUDATest, SliceBf16FirstDim) {
+  SKIP_IF_NO_GPU();
+  // Tensor shape {4, 3}: 4 rows, 3 cols
+  // Data: row0=[1,2,3], row1=[4,5,6], row2=[7,8,9], row3=[10,11,12]
+  Tensor<bfloat16, CPU> tensor({4, 3});
+  for (int i = 0; i < 12; ++i) {
+    tensor.set_(i, bfloat16(i + 1));
+  }
+
+  auto gpu_tensor = tensor.cuda();
+
+  // Slice rows 1 to 3 (exclusive), so rows 1 and 2
+  Tensor<bfloat16, CUDA> result = slice(gpu_tensor.view(), 0, 1, 3);
+
+  auto result_cpu = result.cpu();
+
+  Shape expected_shape = {2, 3};
+  EXPECT_EQ(result_cpu.shape(), expected_shape);
+
+  // Expected: row1=[4,5,6], row2=[7,8,9]
+  std::vector<bfloat16> exp = {4, 5, 6, 7, 8, 9};
+  tensor_is_close<bfloat16>(result_cpu.span(), std::span(exp));
+}
+
+TEST(TensorCUDATest, SliceBf16LastDim) {
+  SKIP_IF_NO_GPU();
+  // Tensor shape {2, 6}
+  // Data: row0=[1,2,3,4,5,6], row1=[7,8,9,10,11,12]
+  Tensor<bfloat16, CPU> tensor({2, 6});
+  for (int i = 0; i < 12; ++i) {
+    tensor.set_(i, bfloat16(i + 1));
+  }
+
+  auto gpu_tensor = tensor.cuda();
+
+  // Slice cols 2 to 5 (exclusive), so cols 2, 3, 4
+  Tensor<bfloat16, CUDA> result = slice(gpu_tensor.view(), 1, 2, 5);
+
+  auto result_cpu = result.cpu();
+
+  Shape expected_shape = {2, 3};
+  EXPECT_EQ(result_cpu.shape(), expected_shape);
+
+  // Expected: row0=[3,4,5], row1=[9,10,11]
+  std::vector<bfloat16> exp = {3, 4, 5, 9, 10, 11};
+  tensor_is_close<bfloat16>(result_cpu.span(), std::span(exp));
+}
+
+TEST(TensorCUDATest, SliceBf16MiddleDim) {
+  SKIP_IF_NO_GPU();
+  // Tensor shape {2, 4, 3}: 2 batches, 4 rows, 3 cols
+  Tensor<bfloat16, CPU> tensor({2, 4, 3});
+  for (int i = 0; i < 24; ++i) {
+    tensor.set_(i, bfloat16(i + 1));
+  }
+
+  auto gpu_tensor = tensor.cuda();
+
+  // Slice dim 1 (rows) from 1 to 3, keeping 2 rows
+  Tensor<bfloat16, CUDA> result = slice(gpu_tensor.view(), 1, 1, 3);
+
+  auto result_cpu = result.cpu();
+
+  Shape expected_shape = {2, 2, 3};
+  EXPECT_EQ(result_cpu.shape(), expected_shape);
+
+  // Batch 0: rows 1-2 = [4,5,6, 7,8,9]
+  // Batch 1: rows 1-2 = [16,17,18, 19,20,21]
+  std::vector<bfloat16> exp = {4, 5, 6, 7, 8, 9, 16, 17, 18, 19, 20, 21};
+  tensor_is_close<bfloat16>(result_cpu.span(), std::span(exp));
+}