Add index_put_

src/core/include/index_ops.h

@@ -0,0 +1,11 @@
+#pragma once
+
+#include <vector>
+
+#include "tensor_iterator.h"
+
+namespace gpu {
+
+Tensor &index_put_(Tensor &self, const std::vector<Tensor> &indices, const Tensor &values);
+
+} // namespace gpu

src/core/include/tensor.h

@@ -30,6 +30,17 @@ Tensor empty_like_reduced(const Tensor &self, int dim, ScalarType dtype);
 Tensor zeros(std::vector<int64_t> shape, ScalarType dtype, int device = 0);
 std::ostream &operator<<(std::ostream &os, const Tensor &t);
 
+template <typename T>
+std::ostream &operator<<(std::ostream &os, const std::vector<T> &vec) {
+    os << "[";
+    for (auto i = 0; i < vec.size(); ++i) {
+        os << vec[i];
+        if (i + 1 != vec.size()) os << ", ";
+    }
+    os << "]";
+    return os;
+}
+
 template <typename T, int vec_size>
 struct d_array {
     T val[vec_size] = {0};
@@ -251,4 +262,5 @@ class Tensor {
     std::tuple<Tensor, Tensor> topk(int64_t k, int64_t dim, bool largest) const;
     std::tuple<Tensor, Tensor> mean_var(int64_t reduce_dim, bool take_sqrt) const;
     std::tuple<Tensor, Tensor> norm_stat(int64_t dim) const;
+    Tensor &index_put_(const std::vector<Tensor> &indices, const Tensor &values);
 };

src/core/include/tensor_iterator.h

@@ -129,6 +129,11 @@ class TensorIterator final {
         return shape_;
     }
 
+    std::vector<int64_t> shape_vec() const {
+        std::vector<int64_t> vec(shape_, shape_ + ndim_);
+        return vec;
+    }
+
     int64_t stride_bytes(int arg, int dim) const {
         return stride_bytes_[arg][dim];
     }

src/core/index_ops.cpp

@@ -0,0 +1,40 @@
+#include "tensor_iterator.h"
+#include "index_ops_kernel.h"
+
+namespace gpu {
+
+Tensor &index_put_(Tensor &self, const std::vector<Tensor> &indices, const Tensor &values) {
+    CHECK_FAIL(indices.size() == self.dim(),
+               "Number of indices must match the number of dimensions in the tensor.");
+    CHECK_FAIL(self.defined() && values.defined(),
+               "Both self and values tensors must be defined.");
+    CHECK_FAIL(self.dtype() == values.dtype(),
+               "Data types of self and values tensors must match.");
+
+    int64_t element_size_bytes = element_size(self.dtype());
+    std::vector<int64_t> indexed_sizes, indexed_strides;
+    for (int dim = 0; dim < indices.size(); ++dim) {
+        indexed_sizes.push_back(self.shape(dim));
+        indexed_strides.push_back(self.stride(dim) * element_size_bytes);
+    }
+    auto replacement_shape = indices[0].sizes();
+    std::vector<int64_t> replacement_strides(replacement_shape.size(), 0);
+    auto self_ = self.as_strided(replacement_shape, replacement_strides);
+
+    if (std::find(indexed_sizes.begin(), indexed_sizes.end(), 0) != indexed_sizes.end()) {
+        CHECK_FAIL(
+            false, "index is out of bounds for dimension with size 0");
+    }
+
+    auto iter = TensorIterator().add_output(self_).check_mem_overlap(false).resize_outputs(false);
+    iter = iter.add_input(values);
+    for (auto &index : indices) {
+        CHECK_FAIL(index.defined() && index.dtype() == ScalarType::Long, "Indices must be of type Long.");
+        iter = iter.add_input(index);
+    }
+    iter = iter.build();
+    index_put_kernel(iter, indexed_sizes, indexed_strides);
+    return self;
+}
+
+} // namespace gpu

src/core/tensor.cpp

@@ -15,6 +15,7 @@
 #include "reduce_ops.h"
 #include "sort_ops.h"
 #include "norm_ops.h"
+#include "index_ops.h"
 #include "accumulate_type.h"
 
 using namespace utils::memory;
@@ -149,6 +150,11 @@ Tensor Tensor::as_strided(std::vector<int64_t> sizes,
     out.numel_ = numel;
     out.is_contiguous_ = false;
     out.storage_offset_ = storage_offset;
+    // TODO: remove it
+    for (int i = ndim; i < MAX_TENSOR_DIMS; i++) {
+        out.shape_[i] = 0;
+        out.stride_[i] = 0;
+    }
     return out;
 }
 
@@ -452,3 +458,7 @@ std::tuple<Tensor, Tensor> Tensor::mean_var(int64_t reduce_dim, bool take_sqrt)
 std::tuple<Tensor, Tensor> Tensor::norm_stat(int64_t dim) const {
     return gpu::norm_stat(*this, dim);
 }
+
+Tensor &Tensor::index_put_(const std::vector<Tensor> &indices, const Tensor &values) {
+    return gpu::index_put_(*this, indices, values);
+}

src/core/tensor_iterator.cpp

@@ -140,7 +140,8 @@ void TensorIterator::mark_resize_outputs() {
                 continue;
             }
             // for reduction, output size does not match shape_, as output is reduced size, and shape_ is size of the input
-            CHECK_FAIL(is_reduction_, "output with shape doesn't match the broadcast shape ");
+            CHECK_FAIL(is_reduction_, "output with shape ", output->shape(), " doesn't match the broadcast shape ",
+                       this->shape_vec());
         }
     }
 }

src/device/include/index_ops_kernel.h

@@ -0,0 +1,5 @@
+#pragma once
+
+#include "tensor_iterator.h"
+
+void index_put_kernel(TensorIterator &iter, const std::vector<int64_t> index_size, const std::vector<int64_t> index_stride);

src/device/index_ops_kernel.cu

@@ -0,0 +1,15 @@
+#include "tensor_index.h"
+
+template <typename scalar_t>
+struct IndexOffsetFunctor {
+    DEVICE void operator()(char *out_data, const char *in_data, int64_t offset) const {
+        *reinterpret_cast<scalar_t *>(out_data + offset) = *reinterpret_cast<const scalar_t *>(in_data);
+    }
+};
+
+void index_put_kernel(TensorIterator &iter, const std::vector<int64_t> index_size, const std::vector<int64_t> index_stride) {
+    DISPATCH_BASIC_TYPES(iter.dtype(), "index_put_kernel", [&]() {
+        auto offset_fn = IndexOffsetFunctor<scalar_t>();
+        gpu_index_kernel(iter, index_size, index_stride, offset_fn);
+    });
+}

src/device/utils/tensor_index.h

@@ -0,0 +1,143 @@
+#pragma once
+
+#include <limits>
+#include <tuple>
+#include <utility>
+#include <iostream>
+#include <vector>
+#include <array>
+
+#include "function_traits.h"
+#include "tensor_iterator.h"
+#include "scalar_type.h"
+#include "tensor_offset_calculator.h"
+#include "tensor_memory_access.h"
+#include "exception.h"
+#include "array.h"
+#include "launcher.h"
+
+static constexpr int launch_bound2 = 4;
+static constexpr int launch_size_nd = 128;
+
+template <int nt, int vt, typename func_t>
+struct IndexElementwiseKernel {
+    DEVICE void operator()(ITEM &item) const {
+        const auto tid = item.thread_idx_x();
+        const auto nv = nt * vt;
+        auto idx = nv * item.block_idx_x() + tid;
+#pragma unroll
+        for (int i = 0; i < vt; i++) {
+            if (idx < N_) {
+                f_(idx);
+                idx += nt;
+            }
+        }
+    }
+    IndexElementwiseKernel(int64_t N, func_t f) :
+        N_(N), f_(f) {
+    }
+
+private:
+    int64_t N_;
+    func_t f_;
+};
+
+template <int nt, int vt, typename func_t>
+static void launch_index_kernel(const int64_t N, const func_t f) {
+    CHECK_FAIL(N >= 0 && N <= std::numeric_limits<int32_t>::max());
+    if (N == 0) {
+        return;
+    }
+    auto kernel = IndexElementwiseKernel<nt, vt, func_t>(N, f);
+    [[maybe_unused]] int nb = (N + nt * vt - 1) / (nt * vt);
+    Launcher::GetInstance()->submit(0, {nb}, {nt}, kernel);
+}
+
+template <typename offset_fn_t, typename func_t>
+struct IndexingOffsetFunctor {
+    DEVICE void operator()(int idx) const {
+        const auto offsets = offset_fn_.get(idx);
+        auto out_data = out_ptr_ + offsets[0];
+        auto in_data = in_ptr_ + offsets[1];
+
+        int64_t offset = 0;
+#pragma unroll
+        for (int i = 0; i < num_indices_; i++) {
+            int64_t index = *reinterpret_cast<int64_t *>(index_ptrs_[i] + offsets[2]);
+            // assert(-sizes_[i] <= index && index < sizes_[i] && "index out of bounds");
+            if (index < 0) {
+                index += sizes_[i];
+            }
+            offset += index * strides_[i];
+        }
+
+        f_(out_data, in_data, offset);
+    }
+    IndexingOffsetFunctor(
+        char *out_ptr,
+        const char *in_ptr,
+        int num_indices,
+        memory::array<int64_t, MAX_TENSOR_DIMS> sizes,
+        memory::array<int64_t, MAX_TENSOR_DIMS> strides,
+        memory::array<char *, MAX_TENSOR_DIMS> index_ptrs,
+        offset_fn_t offset_fn,
+        func_t f) :
+        out_ptr_(out_ptr),
+        in_ptr_(in_ptr),
+        num_indices_(num_indices),
+        sizes_(sizes),
+        strides_(strides),
+        index_ptrs_(index_ptrs),
+        offset_fn_(offset_fn),
+        f_(f) {
+    }
+
+private:
+    char *out_ptr_;
+    const char *in_ptr_;
+    int num_indices_;
+    memory::array<int64_t, MAX_TENSOR_DIMS> sizes_;
+    memory::array<int64_t, MAX_TENSOR_DIMS> strides_;
+    memory::array<char *, MAX_TENSOR_DIMS> index_ptrs_;
+    offset_fn_t offset_fn_;
+    func_t f_;
+};
+
+template <typename func_t>
+void gpu_index_kernel(
+    TensorIterator &iter,
+    const std::vector<int64_t> index_size,
+    const std::vector<int64_t> index_stride,
+    const func_t f) {
+    const auto num_indices = index_size.size();
+    CHECK_FAIL(num_indices == index_stride.size());
+    CHECK_FAIL(static_cast<int64_t>(num_indices) == iter.ntensors() - 2);
+
+    if (iter.numel() == 0) {
+        return;
+    }
+
+    if (!iter.can_use_32bit_indexing()) {
+        for (auto &sub_iter : iter.with_32bit_indexing()) {
+            gpu_index_kernel<func_t>(sub_iter, index_size, index_stride, f);
+        }
+        return;
+    }
+
+    auto sizes = memory::array<int64_t, MAX_TENSOR_DIMS>{};
+    auto strides = memory::array<int64_t, MAX_TENSOR_DIMS>{};
+    auto index_ptrs = memory::array<char *, MAX_TENSOR_DIMS>{};
+    for (unsigned i = 0; i < num_indices; i++) {
+        sizes[i] = index_size[i];
+        strides[i] = index_stride[i];
+        index_ptrs[i] = (char *)iter.data_ptr(i + 2);
+    }
+
+    char *out_ptr = static_cast<char *>(iter.data_ptr(0));
+    char *in_ptr = static_cast<char *>(iter.data_ptr(1));
+
+    auto offset_calc = make_offset_calculator<3>(iter);
+    auto offset_fn = IndexingOffsetFunctor<decltype(offset_calc), func_t>(
+        out_ptr, in_ptr, num_indices, sizes, strides, index_ptrs, offset_calc, f);
+    launch_index_kernel<launch_size_nd, launch_bound2>(iter.numel(), offset_fn);
+}

src/register.cpp

@@ -207,6 +207,7 @@ PYBIND11_MODULE(kfunca, m) {
         .def("mean", &Tensor::mean)
         .def("mean_var", &Tensor::mean_var)
         .def("norm_stat", &Tensor::norm_stat)
+        .def("index_put_", &Tensor::index_put_)
         .def("half", &Tensor::_half)
         .def("bfloat16", &Tensor::_bfloat16)
         .def("float", &Tensor::_float);

test/test_tensor.py

@@ -268,6 +268,19 @@ def test_split(self):
         assert_allclose(arr_t1, arr_gpu1)
         assert_allclose(arr_t2, arr_gpu2)
         assert_allclose(arr_t3, arr_gpu3)
+
+    def test_index_put(self):
+        arr = np.random.uniform(-10000, 10000, size=(13, 15)).astype(np.float32)
+        arr_gpu = kfunca.from_numpy(arr, 0)
+        indices = [kfunca.from_numpy(np.array([0, 5, 1, 2]).astype('q'), 0), 
+                   kfunca.from_numpy(np.array([0, 11, 1, 0]).astype('q'), 0)]
+        values = kfunca.from_numpy(np.random.uniform(-10000, 10000, size=(4)).astype(np.float32), 0)
+        arr_gpu.index_put_(indices, values)
+        arr_gpu_pt = torch.from_numpy(arr)
+        indices_t = [torch.from_numpy(indices[0].numpy()), torch.from_numpy(indices[1].numpy())]
+        values_pt = torch.from_numpy(values.numpy())
+        arr_gpu_pt.index_put_(indices_t, values_pt)
+        assert_allclose(arr_gpu, arr_gpu_pt)
 
 
 if __name__ == '__main__':