From b39b3d3eb352940d5dba5560bc07e0ec911be716 Mon Sep 17 00:00:00 2001 From: xytpai Date: Mon, 28 Jul 2025 13:18:49 +0800 Subject: [PATCH] add cat --- src/core/include/tensor.h | 1 + src/core/include/tensor_iterator.h | 8 +++- src/core/include/tensor_shape.h | 11 +++++ src/core/tensor.cpp | 18 +++++++- src/core/tensor_iterator.cpp | 1 + src/core/tensor_shape.cpp | 72 ++++++++++++++++++++++++++++++ src/core/unary_ops.cpp | 2 +- src/register.cpp | 2 + test/test_tensor.py | 14 ++++++ 9 files changed, 126 insertions(+), 3 deletions(-) create mode 100644 src/core/include/tensor_shape.h create mode 100644 src/core/tensor_shape.cpp diff --git a/src/core/include/tensor.h b/src/core/include/tensor.h index 3e6c876..5a17345 100644 --- a/src/core/include/tensor.h +++ b/src/core/include/tensor.h @@ -229,6 +229,7 @@ class Tensor { Tensor select(int64_t dim, int64_t index) const; Tensor narrow(int64_t dim, int64_t start, int64_t length) const; Tensor view(std::vector sizes) const; + bool can_use_32bit_indexing() const; Tensor _half() const; Tensor _float() const; diff --git a/src/core/include/tensor_iterator.h b/src/core/include/tensor_iterator.h index a891ad6..316a721 100644 --- a/src/core/include/tensor_iterator.h +++ b/src/core/include/tensor_iterator.h @@ -42,6 +42,7 @@ class TensorIterator final { bool accumulate_ = false; bool final_output_ = true; bool is_reduction_ = false; + bool check_mem_overlap_ = true; int64_t reduce_dim_ = 0; ScalarType common_dtype_ = ScalarType::Undefined; @@ -82,7 +83,12 @@ class TensorIterator final { SplitUntil32Bit with_32bit_indexing() const; TensorIterator &resize_outputs(bool flag) { - resize_outputs_ = false; + resize_outputs_ = flag; + return *this; + } + + TensorIterator &check_mem_overlap(bool flag) { + check_mem_overlap_ = flag; return *this; } diff --git a/src/core/include/tensor_shape.h b/src/core/include/tensor_shape.h new file mode 100644 index 0000000..6cbdffb --- /dev/null +++ b/src/core/include/tensor_shape.h @@ -0,0 +1,11 @@ +#pragma once + +#include + +#include "tensor_iterator.h" + +namespace gpu { + +Tensor concat(const std::vector tensors, int64_t dim); + +} diff --git a/src/core/tensor.cpp b/src/core/tensor.cpp index b2ef315..763a217 100644 --- a/src/core/tensor.cpp +++ b/src/core/tensor.cpp @@ -267,6 +267,21 @@ Tensor Tensor::view(std::vector sizes) const { return this->as_strided(sizes, {}); } +bool Tensor::can_use_32bit_indexing() const { + int64_t max_value = std::numeric_limits::max(); + if (this->numel() > max_value) { + return false; + } + int64_t max_offset = 1; + for (int d = 0; d < dim_; ++d) { + max_offset += (shape_[d] - 1) * stride_[d] * element_size(dtype_); + } + if (max_offset > max_value) { + return false; + } + return true; +} + Tensor Tensor::_half() const { return gpu::convert(*this, ScalarType::Half); } @@ -360,7 +375,8 @@ std::ostream &operator<<(std::ostream &os, const Tensor &t) { os << "\b], stride=["; for (int i = 0; i < t.dim(); ++i) os << t.stride(i) << ","; - os << "\b], dtype=" << t.dtype(); + os << "\b], storage_offset=" << t.storage_offset(); + os << ", dtype=" << t.dtype(); os << ", numel=" << t.numel() << ", dim=" << t.dim(); os << ", device=" << t.device() << ") {\n"; print_tensor_(os, t); diff --git a/src/core/tensor_iterator.cpp b/src/core/tensor_iterator.cpp index f4625ba..3bcce2f 100644 --- a/src/core/tensor_iterator.cpp +++ b/src/core/tensor_iterator.cpp @@ -90,6 +90,7 @@ void TensorIterator::mark_outputs() { } void TensorIterator::check_mem_overlaps() { + if (!check_mem_overlap_) return; for (int i = 0; i < num_outputs_; ++i) { auto output = tensors_[i]; if (!output->defined()) continue; diff --git a/src/core/tensor_shape.cpp b/src/core/tensor_shape.cpp new file mode 100644 index 0000000..0b645ed --- /dev/null +++ b/src/core/tensor_shape.cpp @@ -0,0 +1,72 @@ +#include "tensor_shape.h" + +namespace gpu { + +// Check to see if the shape of tensors is compatible +// for being concatenated along a given dimension. +inline void check_cat_shape_except_dim( + const Tensor &first, + const Tensor &second, + int64_t dimension, + int64_t index) { + CHECK_FAIL(first.device() == second.device()); + int64_t first_dims = first.dim(); + int64_t second_dims = second.dim(); + CHECK_FAIL( + first_dims == second_dims, + "Tensors must have same number of dimensions: got ", + first_dims, + " and ", + second_dims); + for (int64_t dim = 0; dim < first_dims; dim++) { + if (dim == dimension) { + continue; + } + int64_t first_dim_size = first.shape(dim); + int64_t second_dim_size = second.shape(dim); + CHECK_FAIL( + first_dim_size == second_dim_size, + "Sizes of tensors must match except in dimension ", + dimension, + ". Expected size ", + static_cast(first_dim_size), + " but got size ", + static_cast(second_dim_size), + " for tensor number ", + index, + " in the list."); + } +} + +Tensor concat(const std::vector tensors, int64_t dim) { + dim = maybe_wrap_dim(dim, tensors[0].dim()); + auto out_dtype = tensors[0].dtype(); + + // valid check + bool all_contiguous = true; + bool all_same_dtype = true; + size_t size_at_dim = tensors[0].shape(dim); + + for (int i = 1; i < tensors.size(); ++i) { + const Tensor &t = tensors[i]; + all_same_dtype = all_same_dtype && out_dtype == t.dtype(); + check_cat_shape_except_dim(tensors[0], t, dim, i); + size_at_dim += t.shape(dim); + all_contiguous = all_contiguous && t.is_contiguous(); + } + + auto out_size = tensors[0].sizes(); + out_size[dim] = size_at_dim; + Tensor result = empty(out_size, out_dtype, tensors[0].device()); + + int64_t offset = 0; + for (const Tensor &t : tensors) { + int64_t dim_size = t.shape(dim); + Tensor nt = result.narrow(dim, offset, dim_size); + nt.copy_(t); + offset += dim_size; + } + return result; +} + +} // namespace gpu diff --git a/src/core/unary_ops.cpp b/src/core/unary_ops.cpp index cd4ef6a..4272651 100644 --- a/src/core/unary_ops.cpp +++ b/src/core/unary_ops.cpp @@ -11,7 +11,7 @@ Tensor clone(const Tensor &self) { } Tensor ©_(Tensor &self, const Tensor &other) { - auto iter = TensorIterator().add_output(self).add_input(other).resize_outputs(false).build_for_loops(); + auto iter = TensorIterator().add_output(self).add_input(other).resize_outputs(false).check_mem_overlap(false).build_for_loops(); copy_kernel(iter); return self; } diff --git a/src/register.cpp b/src/register.cpp index 741a263..4a93112 100644 --- a/src/register.cpp +++ b/src/register.cpp @@ -6,6 +6,7 @@ #include "device_info.h" #include "tensor.h" +#include "tensor_shape.h" #include "binary_ops.h" #include "gemm_ops.h" #include "nn_ops.h" @@ -82,6 +83,7 @@ PYBIND11_MODULE(kfunca, m) { m.def("zeros", &zeros); m.def("causal_attention", &gpu::causal_attention); m.def("gemm", &gpu::gemm); + m.def("cat", &gpu::concat); py::class_(m, "tensor") .def("__copy__", [](const Tensor &self) { return Tensor(self); }) .def("__deepcopy__", [](const Tensor &self, py::dict) { return Tensor(self); }) diff --git a/test/test_tensor.py b/test/test_tensor.py index 84c5256..45aae5d 100644 --- a/test/test_tensor.py +++ b/test/test_tensor.py @@ -238,6 +238,20 @@ def test_view(self): arr_t = arr_t.view(5,-1,23).contiguous() + 1 arr_gpu = arr_gpu.view(5,-1,23).contiguous() + 1 assert_allclose(arr_t, arr_gpu) + + def test_cat(self): + arr1 = np.random.uniform(-10000, 10000, size=(5,11,23)).astype(np.float32) + arr2 = np.random.uniform(-10000, 10000, size=(5,13,23)).astype(np.float32) + arr3 = np.random.uniform(-10000, 10000, size=(5,1,23)).astype(np.float32) + arr1_t = torch.from_numpy(arr1) + arr2_t = torch.from_numpy(arr2) + arr3_t = torch.from_numpy(arr3) + arr1_gpu = kfunca.from_numpy(arr1, 0) + arr2_gpu = kfunca.from_numpy(arr2, 0) + arr3_gpu = kfunca.from_numpy(arr3, 0) + arr_t = torch.cat([arr1_t, arr2_t, arr3_t], 1) + arr_gpu = kfunca.cat([arr1_gpu, arr2_gpu, arr3_gpu], 1) + assert_allclose(arr_t, arr_gpu) if __name__ == '__main__':