diff --git a/src/core/include/index_ops.h b/src/core/include/index_ops.h new file mode 100644 index 0000000..b00525b --- /dev/null +++ b/src/core/include/index_ops.h @@ -0,0 +1,11 @@ +#pragma once + +#include + +#include "tensor_iterator.h" + +namespace gpu { + +Tensor &index_put_(Tensor &self, const std::vector &indices, const Tensor &values); + +} // namespace gpu diff --git a/src/core/include/tensor.h b/src/core/include/tensor.h index d12b9b8..8c1282a 100644 --- a/src/core/include/tensor.h +++ b/src/core/include/tensor.h @@ -30,6 +30,17 @@ Tensor empty_like_reduced(const Tensor &self, int dim, ScalarType dtype); Tensor zeros(std::vector shape, ScalarType dtype, int device = 0); std::ostream &operator<<(std::ostream &os, const Tensor &t); +template +std::ostream &operator<<(std::ostream &os, const std::vector &vec) { + os << "["; + for (auto i = 0; i < vec.size(); ++i) { + os << vec[i]; + if (i + 1 != vec.size()) os << ", "; + } + os << "]"; + return os; +} + template struct d_array { T val[vec_size] = {0}; @@ -251,4 +262,5 @@ class Tensor { std::tuple topk(int64_t k, int64_t dim, bool largest) const; std::tuple mean_var(int64_t reduce_dim, bool take_sqrt) const; std::tuple norm_stat(int64_t dim) const; + Tensor &index_put_(const std::vector &indices, const Tensor &values); }; diff --git a/src/core/include/tensor_iterator.h b/src/core/include/tensor_iterator.h index 316a721..234a477 100644 --- a/src/core/include/tensor_iterator.h +++ b/src/core/include/tensor_iterator.h @@ -129,6 +129,11 @@ class TensorIterator final { return shape_; } + std::vector shape_vec() const { + std::vector vec(shape_, shape_ + ndim_); + return vec; + } + int64_t stride_bytes(int arg, int dim) const { return stride_bytes_[arg][dim]; } diff --git a/src/core/index_ops.cpp b/src/core/index_ops.cpp new file mode 100644 index 0000000..002a40e --- /dev/null +++ b/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 &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 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 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 diff --git a/src/core/tensor.cpp b/src/core/tensor.cpp index 7e1f791..c9c7cbe 100644 --- a/src/core/tensor.cpp +++ b/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 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::mean_var(int64_t reduce_dim, bool take_sqrt) std::tuple Tensor::norm_stat(int64_t dim) const { return gpu::norm_stat(*this, dim); } + +Tensor &Tensor::index_put_(const std::vector &indices, const Tensor &values) { + return gpu::index_put_(*this, indices, values); +} diff --git a/src/core/tensor_iterator.cpp b/src/core/tensor_iterator.cpp index 3bcce2f..da6204b 100644 --- a/src/core/tensor_iterator.cpp +++ b/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()); } } } diff --git a/src/device/include/index_ops_kernel.h b/src/device/include/index_ops_kernel.h new file mode 100644 index 0000000..a6f4911 --- /dev/null +++ b/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 index_size, const std::vector index_stride); diff --git a/src/device/index_ops_kernel.cu b/src/device/index_ops_kernel.cu new file mode 100644 index 0000000..0b69219 --- /dev/null +++ b/src/device/index_ops_kernel.cu @@ -0,0 +1,15 @@ +#include "tensor_index.h" + +template +struct IndexOffsetFunctor { + DEVICE void operator()(char *out_data, const char *in_data, int64_t offset) const { + *reinterpret_cast(out_data + offset) = *reinterpret_cast(in_data); + } +}; + +void index_put_kernel(TensorIterator &iter, const std::vector index_size, const std::vector index_stride) { + DISPATCH_BASIC_TYPES(iter.dtype(), "index_put_kernel", [&]() { + auto offset_fn = IndexOffsetFunctor(); + gpu_index_kernel(iter, index_size, index_stride, offset_fn); + }); +} diff --git a/src/device/utils/tensor_index.h b/src/device/utils/tensor_index.h new file mode 100644 index 0000000..1acf2d3 --- /dev/null +++ b/src/device/utils/tensor_index.h @@ -0,0 +1,143 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#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 +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 +static void launch_index_kernel(const int64_t N, const func_t f) { + CHECK_FAIL(N >= 0 && N <= std::numeric_limits::max()); + if (N == 0) { + return; + } + auto kernel = IndexElementwiseKernel(N, f); + [[maybe_unused]] int nb = (N + nt * vt - 1) / (nt * vt); + Launcher::GetInstance()->submit(0, {nb}, {nt}, kernel); +} + +template +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(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 sizes, + memory::array strides, + memory::array 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 sizes_; + memory::array strides_; + memory::array index_ptrs_; + offset_fn_t offset_fn_; + func_t f_; +}; + +template +void gpu_index_kernel( + TensorIterator &iter, + const std::vector index_size, + const std::vector index_stride, + const func_t f) { + const auto num_indices = index_size.size(); + CHECK_FAIL(num_indices == index_stride.size()); + CHECK_FAIL(static_cast(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(sub_iter, index_size, index_stride, f); + } + return; + } + + auto sizes = memory::array{}; + auto strides = memory::array{}; + auto index_ptrs = memory::array{}; + 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(iter.data_ptr(0)); + char *in_ptr = static_cast(iter.data_ptr(1)); + + auto offset_calc = make_offset_calculator<3>(iter); + auto offset_fn = IndexingOffsetFunctor( + out_ptr, in_ptr, num_indices, sizes, strides, index_ptrs, offset_calc, f); + launch_index_kernel(iter.numel(), offset_fn); +} diff --git a/src/register.cpp b/src/register.cpp index 9b2869b..93aa0d4 100644 --- a/src/register.cpp +++ b/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); diff --git a/test/test_tensor.py b/test/test_tensor.py index 873de90..c111d41 100644 --- a/test/test_tensor.py +++ b/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__':