Add concat

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<int64_t> sizes) const;
+    bool can_use_32bit_indexing() const;
 
     Tensor _half() const;
     Tensor _float() const;

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;
     }
 

src/core/include/tensor_shape.h

@@ -0,0 +1,11 @@
+#pragma once
+
+#include <vector>
+
+#include "tensor_iterator.h"
+
+namespace gpu {
+
+Tensor concat(const std::vector<Tensor> tensors, int64_t dim);
+
+}

src/core/tensor.cpp

@@ -267,6 +267,21 @@ Tensor Tensor::view(std::vector<int64_t> sizes) const {
     return this->as_strided(sizes, {});
 }
 
+bool Tensor::can_use_32bit_indexing() const {
+    int64_t max_value = std::numeric_limits<int32_t>::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);

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;

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<long long>(first_dim_size),
+            " but got size ",
+            static_cast<long long>(second_dim_size),
+            " for tensor number ",
+            index,
+            " in the list.");
+    }
+}
+
+Tensor concat(const std::vector<Tensor> 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

src/core/unary_ops.cpp

@@ -11,7 +11,7 @@ Tensor clone(const Tensor &self) {
 }
 
 Tensor &copy_(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;
 }

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_<Tensor>(m, "tensor")
         .def("__copy__", [](const Tensor &self) { return Tensor(self); })
         .def("__deepcopy__", [](const Tensor &self, py::dict) { return Tensor(self); })

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__':