diff --git a/CMakeLists.txt b/CMakeLists.txt index 73d6c6b..f08d2cf 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -26,6 +26,7 @@ include_directories(${ROOTDIR}/src/core) include_directories(${ROOTDIR}/src/core/include) include_directories(${ROOTDIR}/src/core/utils) include_directories(${ROOTDIR}/src/core/utils/memory) +# set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g") file(GLOB CORE_SRCS "src/core/*.cpp") add_library(kfunca_core SHARED ${CORE_SRCS}) diff --git a/src/core/include/tensor.h b/src/core/include/tensor.h index 6e2a1cd..40a098b 100644 --- a/src/core/include/tensor.h +++ b/src/core/include/tensor.h @@ -1,26 +1,11 @@ #pragma once -#include -#include -#include -#include #include -#include "exception.h" -#include "data_ptr.h" -#include "intrusive_ptr.h" -#include "scalar_type.h" -#include "device_allocator.h" -#include "memory_overlap.h" +#include "tensor_impl.h" #define MAX_TENSOR_DIMS 12 -using namespace utils::memory; - -inline int maybe_wrap_dim(int d, int ndim) { - return d < 0 ? (ndim + d) % ndim : d; -} - class Tensor; Tensor empty(std::vector shape, ScalarType dtype, int device = 0); Tensor empty(int64_t *shape, int ndim, ScalarType dtype, int device, bool inverse = false); @@ -30,91 +15,9 @@ 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 -struct d_array { - T val[vec_size] = {0}; - d_array() { - } - d_array(double value) { - *reinterpret_cast(&val[0]) = value; - } - operator double() const { - return *reinterpret_cast(const_cast(&val[0])); - } - T &operator[](int i) { - return val[i]; - } - T const &operator[](int i) const { - return val[i]; - } - bool equals(d_array &other) { -#pragma unroll - for (int i = 0; i < vec_size; ++i) { - if (val[i] != other[i]) { - return false; - } - } - return true; - } - bool equals(T (&other)[vec_size]) { -#pragma unroll - for (int i = 0; i < vec_size; ++i) { - if (val[i] != other[i]) { - return false; - } - } - return true; - } -}; - -typedef d_array dim_t; -using any_t = d_array; -std::ostream &operator<<(std::ostream &os, const dim_t &d); - -class TensorStorage : public intrusive_ptr_target { -protected: - size_t size_; - int device_; - DataPtr ptr_; - -public: - TensorStorage() : - size_(0), device_(-1), ptr_() { - } - TensorStorage(size_t size, int device) : - size_(size), device_(device) { - ptr_ = DeviceAllocator::GetInstance()->allocate(size, device); - } - size_t size() const { - return size_; - } - int device() const { - return device_; - } - void *data_ptr() const { - return ptr_.get(); - } - template - T *data_ptr() const { - return reinterpret_cast(ptr_.get()); - } - bool defined() const { - return static_cast(ptr_); - } -}; - class Tensor { - int dim_; - dim_t shape_; - dim_t stride_; - ScalarType dtype_; - int64_t numel_; - intrusive_ptr storage_; - int64_t storage_offset_ = 0; - bool is_contiguous_ = true; - - void new_storage_(int device); - +private: + intrusive_ptr impl_; friend Tensor empty(std::vector shape, ScalarType dtype, int device); friend Tensor empty(int64_t *shape, int ndim, ScalarType dtype, int device, bool inverse); friend Tensor empty_like(const Tensor &self); @@ -122,91 +25,73 @@ class Tensor { friend Tensor empty_like_reduced(const Tensor &self, int dim, ScalarType dtype); friend Tensor zeros(std::vector shape, ScalarType dtype, int device); - Tensor(std::vector &shape, ScalarType dtype); - Tensor(std::vector &shape, std::vector &strides, ScalarType dtype); - Tensor(int64_t *shape, int ndim, ScalarType dtype, bool inverse); - public: - Tensor() : - storage_() { - } - Tensor(const Tensor &other) : - dim_(other.dim_), shape_(other.shape_), stride_(other.stride_), - dtype_(other.dtype_), numel_(other.numel_), - storage_(other.storage_), storage_offset_(other.storage_offset_), - is_contiguous_(other.is_contiguous_) { - } - Tensor &operator=(const Tensor &other) { - dim_ = other.dim_; - shape_ = other.shape_; - stride_ = other.stride_; - dtype_ = other.dtype_; - numel_ = other.numel_; - storage_ = other.storage_; - storage_offset_ = other.storage_offset_; - is_contiguous_ = other.is_contiguous_; - return *this; - } + Tensor() = default; + Tensor(const Tensor &other) = default; + Tensor &operator=(const Tensor &other) = default; Tensor(Tensor &&other) = default; Tensor &operator=(Tensor &&other) = default; + TensorImpl *impl() const { + return impl_.get(); + } int dim() const { - return dim_; + return impl_.get()->dim(); } int64_t shape(int d) const { - d = maybe_wrap_dim(d, dim_); - return shape_[d]; + return impl_.get()->shape(d); } dim_t &shape() { - return shape_; + return impl_.get()->shape(); } std::vector sizes() const { - std::vector vec(shape_.val, shape_.val + dim_); - return vec; + return impl_.get()->sizes(); } std::vector strides() const { - std::vector vec(stride_.val, stride_.val + dim_); - return vec; + return impl_.get()->strides(); } int64_t stride(int d) const { - return stride_[d]; + return impl_.get()->stride(d); } dim_t &stride() { - return stride_; + return impl_.get()->stride(); } ScalarType dtype() const { - return dtype_; + return impl_.get()->dtype(); } int64_t numel() const { - return numel_; + return impl_.get()->numel(); } void *data_ptr() const { - return (char *)storage_.get()->data_ptr() + storage_offset_ * element_size(dtype_); + return impl_.get()->data_ptr(); } template T *data_ptr() const { - return reinterpret_cast(storage_.get()->data_ptr()) + storage_offset_; + return impl_.get()->data_ptr(); } size_t storage_bytes() const { - return storage_.get()->size(); + return impl_.get()->storage_bytes(); } size_t storage_ref_count() const { - return storage_.ref_count(); + return impl_.get()->storage_ref_count(); + } + size_t impl_ref_count() const { + return impl_.ref_count(); } int64_t storage_offset() const { - return storage_offset_; + return impl_.get()->storage_offset(); } intrusive_ptr storage() const { - return storage_; + return impl_.get()->storage(); } bool defined() const { - return storage_.get() != nullptr; + return impl_.get() && impl_.get()->defined(); } int device() const { - return storage_.get()->device(); + return impl_.get()->device(); } int64_t element_size_in_bytes() const { - return element_size(dtype_); + return impl_.get()->element_size_in_bytes(); } std::string to_string() const { std::ostringstream oss; @@ -214,7 +99,7 @@ class Tensor { return oss.str(); } bool is_contiguous() const { - return is_contiguous_; + return impl_.get()->is_contiguous(); } void copy_from_cpu_ptr(void *ptr); diff --git a/src/core/include/tensor_impl.h b/src/core/include/tensor_impl.h new file mode 100644 index 0000000..f0bc229 --- /dev/null +++ b/src/core/include/tensor_impl.h @@ -0,0 +1,195 @@ +#pragma once + +#include +#include + +#include "data_ptr.h" +#include "exception.h" +#include "intrusive_ptr.h" +#include "scalar_type.h" +#include "device_allocator.h" +#include "memory_overlap.h" + +#define MAX_TENSOR_DIMS 12 + +using namespace utils::memory; + +inline int maybe_wrap_dim(int d, int ndim) { + return d < 0 ? (ndim + d) % ndim : d; +} + +template +struct d_array { + T val[vec_size] = {0}; + d_array() { + } + d_array(double value) { + *reinterpret_cast(&val[0]) = value; + } + operator double() const { + return *reinterpret_cast(const_cast(&val[0])); + } + T &operator[](int i) { + return val[i]; + } + T const &operator[](int i) const { + return val[i]; + } + bool equals(d_array &other) { +#pragma unroll + for (int i = 0; i < vec_size; ++i) { + if (val[i] != other[i]) { + return false; + } + } + return true; + } + bool equals(T (&other)[vec_size]) { +#pragma unroll + for (int i = 0; i < vec_size; ++i) { + if (val[i] != other[i]) { + return false; + } + } + return true; + } +}; + +typedef d_array dim_t; +using any_t = d_array; +std::ostream &operator<<(std::ostream &os, const dim_t &d); + +class TensorStorage : public intrusive_ptr_target { +protected: + size_t size_; + int device_; + DataPtr ptr_; + +public: + TensorStorage() : + size_(0), device_(-1), ptr_() { + } + TensorStorage(size_t size, int device) : + size_(size), device_(device) { + ptr_ = DeviceAllocator::GetInstance()->allocate(size, device); + } + size_t size() const { + return size_; + } + int device() const { + return device_; + } + void *data_ptr() const { + return ptr_.get(); + } + template + T *data_ptr() const { + return reinterpret_cast(ptr_.get()); + } + bool defined() const { + return static_cast(ptr_); + } +}; + +class TensorImpl : public intrusive_ptr_target { +private: + int dim_; + dim_t shape_; + dim_t stride_; + ScalarType dtype_; + int64_t numel_; + intrusive_ptr storage_; + int64_t storage_offset_ = 0; + bool is_contiguous_ = true; + +public: + TensorImpl(std::vector &shape, ScalarType dtype); + TensorImpl(std::vector &shape, std::vector &strides, ScalarType dtype); + TensorImpl(int64_t *shape, int ndim, ScalarType dtype, bool inverse); + TensorImpl() : + storage_() { + } + TensorImpl(const TensorImpl &other) : + dim_(other.dim_), shape_(other.shape_), stride_(other.stride_), + dtype_(other.dtype_), numel_(other.numel_), + storage_(other.storage_), storage_offset_(other.storage_offset_), + is_contiguous_(other.is_contiguous_) { + } + TensorImpl &operator=(const TensorImpl &other) { + dim_ = other.dim_; + shape_ = other.shape_; + stride_ = other.stride_; + dtype_ = other.dtype_; + numel_ = other.numel_; + storage_ = other.storage_; + storage_offset_ = other.storage_offset_; + is_contiguous_ = other.is_contiguous_; + return *this; + } + TensorImpl(TensorImpl &&other) = default; + TensorImpl &operator=(TensorImpl &&other) = default; + + int dim() const { + return dim_; + } + int64_t shape(int d) const { + d = maybe_wrap_dim(d, dim_); + return shape_[d]; + } + dim_t &shape() { + return shape_; + } + std::vector sizes() const { + std::vector vec(shape_.val, shape_.val + dim_); + return vec; + } + std::vector strides() const { + std::vector vec(stride_.val, stride_.val + dim_); + return vec; + } + int64_t stride(int d) const { + return stride_[d]; + } + dim_t &stride() { + return stride_; + } + ScalarType dtype() const { + return dtype_; + } + int64_t numel() const { + return numel_; + } + void *data_ptr() const { + return (char *)storage_.get()->data_ptr() + storage_offset_ * element_size(dtype_); + } + template + T *data_ptr() const { + return reinterpret_cast(storage_.get()->data_ptr()) + storage_offset_; + } + size_t storage_bytes() const { + return storage_.get()->size(); + } + size_t storage_ref_count() const { + return storage_.ref_count(); + } + int64_t storage_offset() const { + return storage_offset_; + } + intrusive_ptr storage() const { + return storage_; + } + bool defined() const { + return storage_.get() != nullptr; + } + int device() const { + return storage_.get()->device(); + } + int64_t element_size_in_bytes() const { + return element_size(dtype_); + } + bool is_contiguous() const { + return is_contiguous_; + } + void new_storage_(int device); + void as_strided_(std::vector sizes, std::vector strides, int64_t storage_offset = 0); +}; diff --git a/src/core/tensor.cpp b/src/core/tensor.cpp index c9c7cbe..aaad7d8 100644 --- a/src/core/tensor.cpp +++ b/src/core/tensor.cpp @@ -1,13 +1,7 @@ -#include #include -#include #include "tensor.h" #include "tensor_shape.h" -#include "exception.h" -#include "data_ptr.h" -#include "intrusive_ptr.h" -#include "scalar_type.h" #include "memory_engine.h" #include "binary_ops.h" #include "unary_ops.h" @@ -18,65 +12,58 @@ #include "index_ops.h" #include "accumulate_type.h" -using namespace utils::memory; +Tensor empty(std::vector shape, ScalarType dtype, int device) { + Tensor output; + auto impl = new TensorImpl(shape, dtype); + impl->new_storage_(device); + output.impl_.unsafe_set_ptr(impl); + return output; +} -void Tensor::new_storage_(int device) { - auto [min_offset, max_offset] = compute_offset_range(shape_, stride_, dim_); - size_t offset_range = max_offset - min_offset + 1; - size_t bytes = offset_range * element_size(dtype_); - auto ptr = new TensorStorage(bytes, device); - storage_.unsafe_set_ptr(ptr); +Tensor empty(int64_t *shape, int ndim, ScalarType dtype, int device, bool inverse) { + Tensor output; + auto impl = new TensorImpl(shape, ndim, dtype, inverse); + impl->new_storage_(device); + output.impl_.unsafe_set_ptr(impl); + return output; } -std::ostream &operator<<(std::ostream &os, const dim_t &d) { - os << "dim_t:"; - for (int i = 0; i < MAX_TENSOR_DIMS; i++) - os << d[i] << ", "; - os << "\n"; - return os; +Tensor empty_like(const Tensor &self) { + auto sizes = self.sizes(); + Tensor output; + auto impl = new TensorImpl(sizes, self.dtype()); + impl->new_storage_(self.device()); + output.impl_.unsafe_set_ptr(impl); + return output; } -Tensor::Tensor(std::vector &shape, ScalarType dtype) { - CHECK_FAIL(shape.size() <= MAX_TENSOR_DIMS); - dtype_ = dtype; - dim_ = shape.size(); - numel_ = 1; - for (int i = dim_ - 1; i >= 0; i--) { - stride_[i] = numel_; - numel_ *= shape[i]; - shape_[i] = shape[i]; - } +Tensor empty_strided(std::vector shape, std::vector strides, ScalarType dtype, int device) { + Tensor output; + auto impl = new TensorImpl(shape, strides, dtype); + impl->new_storage_(device); + output.impl_.unsafe_set_ptr(impl); + return output; } -Tensor::Tensor(std::vector &shape, std::vector &strides, ScalarType dtype) { - CHECK_FAIL(shape.size() <= MAX_TENSOR_DIMS); - CHECK_FAIL(strides.size() <= MAX_TENSOR_DIMS); - dtype_ = dtype; - dim_ = shape.size(); - CHECK_FAIL(dim_ == strides.size()); - numel_ = 1; - for (int i = dim_ - 1; i >= 0; i--) { - stride_[i] = strides[i]; - numel_ *= shape[i]; - shape_[i] = shape[i]; +Tensor empty_like_reduced(const Tensor &self, int dim, ScalarType dtype) { + auto sizes = self.sizes(); + if (dim >= 0) { + sizes[dim] = 1; } + Tensor output; + auto impl = new TensorImpl(sizes, dtype); + impl->new_storage_(self.device()); + output.impl_.unsafe_set_ptr(impl); + return output; } -Tensor::Tensor(int64_t *shape, int ndim, ScalarType dtype, bool inverse) { - CHECK_FAIL(ndim <= MAX_TENSOR_DIMS); - dtype_ = dtype; - dim_ = ndim; - numel_ = 1; - int is; - for (int i = dim_ - 1; i >= 0; i--) { - stride_[i] = numel_; - if (!inverse) - is = i; - else - is = dim_ - 1 - i; - numel_ *= shape[is]; - shape_[i] = shape[is]; - } +Tensor zeros(std::vector shape, ScalarType dtype, int device) { + Tensor output; + auto impl = new TensorImpl(shape, dtype); + impl->new_storage_(device); + output.impl_.unsafe_set_ptr(impl); + dmemset_zeros(output.data_ptr(), output.storage_bytes()); + return output; } void Tensor::copy_from_cpu_ptr(void *ptr) { @@ -104,62 +91,31 @@ Tensor &Tensor::fill_(const any_t &value) { } int64_t Tensor::offset(const std::vector &indices) const { - CHECK_FAIL(indices.size() == dim_); + CHECK_FAIL(indices.size() == dim()); int64_t flat_index = 0; for (size_t i = 0; i < indices.size(); ++i) { - flat_index += indices[i] * stride_[i]; + flat_index += indices[i] * stride(i); } return flat_index; } Tensor Tensor::contiguous() const { - if (is_contiguous_) + if (is_contiguous()) return *this; return gpu::clone(*this); } Tensor Tensor::as_strided(std::vector sizes, std::vector strides, int64_t storage_offset) const { - auto ndim = sizes.size(); - bool has_strides = strides.size() > 0; - if (has_strides) { - CHECK_FAIL(ndim == strides.size()); - } else { - strides.reserve(ndim); - int64_t cumprod = 1; - for (int i = ndim - 1; i >= 0; i--) { - strides[i] = cumprod; - cumprod *= sizes[i]; - } - } - // in-bounds check - auto [min_offset, max_offset] = compute_offset_range(sizes, strides, ndim); - min_offset += storage_offset; - max_offset += storage_offset; - CHECK_FAIL(min_offset >= 0); - CHECK_FAIL(max_offset * element_size(dtype_) < this->storage_bytes()); - // create tensor view - Tensor out(*this); - out.dim_ = ndim; - int64_t numel = 1; - for (int i = 0; i < ndim; i++) { - out.shape_[i] = sizes[i]; - numel *= sizes[i]; - out.stride_[i] = strides[i]; - } - 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; + Tensor output; + auto impl = new TensorImpl(*(impl_.get())); + impl->as_strided_(sizes, strides, storage_offset); + output.impl_.unsafe_set_ptr(impl); + return output; } Tensor Tensor::permute(const std::vector dims) const { - const auto ndim = dim_; + const auto ndim = dim(); CHECK_FAIL(ndim == dims.size()); auto new_sizes = std::vector(ndim); auto new_strides = std::vector(ndim); @@ -168,8 +124,8 @@ Tensor Tensor::permute(const std::vector dims) const { int d = maybe_wrap_dim(dims[i], ndim); CHECK_FAIL(!seen_dims[d], "permute(): duplicate dims are not allowed."); seen_dims[d] = true; - new_sizes[i] = this->shape_[d]; - new_strides[i] = this->stride_[d]; + new_sizes[i] = this->shape(d); + new_strides[i] = this->stride(d); } return as_strided(new_sizes, new_strides); } @@ -252,7 +208,7 @@ Tensor Tensor::narrow(int64_t dim, int64_t start, int64_t length) const { } Tensor Tensor::view(std::vector sizes) const { - CHECK_FAIL(this->is_contiguous_); + CHECK_FAIL(this->is_contiguous()); int64_t cumprod = 1; bool has_neg_dim = false; int64_t neg_dim = -1; @@ -267,10 +223,10 @@ Tensor Tensor::view(std::vector sizes) const { } } if (has_neg_dim) { - sizes[neg_dim] = this->numel_ / cumprod; + sizes[neg_dim] = this->numel() / cumprod; cumprod *= sizes[neg_dim]; } - CHECK_FAIL(cumprod == this->numel_); + CHECK_FAIL(cumprod == this->numel()); return this->as_strided(sizes, {}); } @@ -280,8 +236,8 @@ bool Tensor::can_use_32bit_indexing() const { return false; } int64_t max_offset = 1; - for (int d = 0; d < dim_; ++d) { - max_offset += (shape_[d] - 1) * stride_[d] * element_size(dtype_); + for (int d = 0; d < this->dim(); ++d) { + max_offset += (shape(d) - 1) * stride(d) * element_size(dtype()); } if (max_offset > max_value) { return false; @@ -305,48 +261,6 @@ Tensor Tensor::_float() const { return gpu::convert(*this, ScalarType::Float); } -Tensor empty(std::vector shape, ScalarType dtype, int device) { - Tensor output(shape, dtype); - output.new_storage_(device); - return output; -} - -Tensor empty(int64_t *shape, int ndim, ScalarType dtype, int device, bool inverse) { - Tensor output(shape, ndim, dtype, inverse); - output.new_storage_(device); - return output; -} - -Tensor empty_like(const Tensor &self) { - auto sizes = self.sizes(); - Tensor output(sizes, self.dtype()); - output.new_storage_(self.device()); - return output; -} - -Tensor empty_strided(std::vector shape, std::vector strides, ScalarType dtype, int device) { - Tensor output(shape, strides, dtype); - output.new_storage_(device); - return output; -} - -Tensor empty_like_reduced(const Tensor &self, int dim, ScalarType dtype) { - auto sizes = self.sizes(); - if (dim >= 0) { - sizes[dim] = 1; - } - Tensor output(sizes, dtype); - output.new_storage_(self.device()); - return output; -} - -Tensor zeros(std::vector shape, ScalarType dtype, int device) { - Tensor output(shape, dtype); - output.new_storage_(device); - dmemset_zeros(output.data_ptr(), output.storage_bytes()); - return output; -} - void print_tensor_(std::ostream &os, const Tensor &t, std::vector indices = {}, int dim = 0) { if (dim == t.dim()) { auto result_ = t.item(indices); diff --git a/src/core/tensor_impl.cpp b/src/core/tensor_impl.cpp new file mode 100644 index 0000000..9223c3d --- /dev/null +++ b/src/core/tensor_impl.cpp @@ -0,0 +1,99 @@ +#include "tensor_impl.h" + +std::ostream &operator<<(std::ostream &os, const dim_t &d) { + os << "dim_t:"; + for (int i = 0; i < MAX_TENSOR_DIMS; i++) + os << d[i] << ", "; + os << "\n"; + return os; +} + +TensorImpl::TensorImpl(std::vector &shape, ScalarType dtype) { + CHECK_FAIL(shape.size() <= MAX_TENSOR_DIMS); + dtype_ = dtype; + dim_ = shape.size(); + numel_ = 1; + for (int i = dim_ - 1; i >= 0; i--) { + stride_[i] = numel_; + numel_ *= shape[i]; + shape_[i] = shape[i]; + } +} + +TensorImpl::TensorImpl(std::vector &shape, std::vector &strides, ScalarType dtype) { + CHECK_FAIL(shape.size() <= MAX_TENSOR_DIMS); + CHECK_FAIL(strides.size() <= MAX_TENSOR_DIMS); + dtype_ = dtype; + dim_ = shape.size(); + CHECK_FAIL(dim_ == strides.size()); + numel_ = 1; + for (int i = dim_ - 1; i >= 0; i--) { + stride_[i] = strides[i]; + numel_ *= shape[i]; + shape_[i] = shape[i]; + } +} + +TensorImpl::TensorImpl(int64_t *shape, int ndim, ScalarType dtype, bool inverse) { + CHECK_FAIL(ndim <= MAX_TENSOR_DIMS); + dtype_ = dtype; + dim_ = ndim; + numel_ = 1; + int is; + for (int i = dim_ - 1; i >= 0; i--) { + stride_[i] = numel_; + if (!inverse) + is = i; + else + is = dim_ - 1 - i; + numel_ *= shape[is]; + shape_[i] = shape[is]; + } +} + +void TensorImpl::new_storage_(int device) { + bool has_defined_storage = storage_.get() && storage_.get()->defined(); + CHECK_FAIL(!has_defined_storage); + auto [min_offset, max_offset] = compute_offset_range(shape_, stride_, dim_); + size_t offset_range = max_offset - min_offset + 1; + size_t bytes = offset_range * element_size(dtype_); + auto ptr = new TensorStorage(bytes, device); + storage_.unsafe_set_ptr(ptr); +} + +void TensorImpl::as_strided_(std::vector sizes, std::vector strides, int64_t storage_offset) { + auto ndim = sizes.size(); + bool has_strides = strides.size() > 0; + if (has_strides) { + CHECK_FAIL(ndim == strides.size()); + } else { + strides.reserve(ndim); + int64_t cumprod = 1; + for (int i = ndim - 1; i >= 0; i--) { + strides[i] = cumprod; + cumprod *= sizes[i]; + } + } + // in-bounds check + auto [min_offset, max_offset] = compute_offset_range(sizes, strides, ndim); + min_offset += storage_offset; + max_offset += storage_offset; + CHECK_FAIL(min_offset >= 0); + CHECK_FAIL(max_offset * element_size(dtype_) < this->storage_bytes()); + // create tensor view + dim_ = ndim; + int64_t numel = 1; + for (int i = 0; i < ndim; i++) { + shape_[i] = sizes[i]; + numel *= sizes[i]; + stride_[i] = strides[i]; + } + numel_ = numel; + is_contiguous_ = false; + storage_offset_ = storage_offset; + // TODO: remove it + for (int i = ndim; i < MAX_TENSOR_DIMS; i++) { + shape_[i] = 0; + stride_[i] = 0; + } +} diff --git a/src/register.cpp b/src/register.cpp index 93aa0d4..e3a41cc 100644 --- a/src/register.cpp +++ b/src/register.cpp @@ -117,6 +117,7 @@ PYBIND11_MODULE(kfunca, m) { return reinterpret_cast(self.data_ptr()); }) .def("storage_ref_count", &Tensor::storage_ref_count) + .def("impl_ref_count", &Tensor::impl_ref_count) .def("contiguous", &Tensor::contiguous) .def("permute", [](Tensor &self, py::args args) { CHECK_FAIL(args.size() == self.dim()); diff --git a/test/test_tensor.py b/test/test_tensor.py index c111d41..c43a517 100644 --- a/test/test_tensor.py +++ b/test/test_tensor.py @@ -70,17 +70,18 @@ def test_inplace_op(self): def test_data_ptr(self): import copy arr_ = np.random.uniform(-10, 10, size=(3,4)).astype(np.float32) - arr_x = kfunca.from_numpy(arr_, 0) - arr_x_ref = kfunca.from_numpy(arr_, 0) - arr_x_ref = arr_x - arr_x_deep = copy.deepcopy(arr_x) + arr_x = kfunca.from_numpy(arr_, 0) # x-1 + arr_x_ref = kfunca.from_numpy(arr_, 0) # y-1 + arr_x_ref = arr_x # x-1, y-0 + arr_x_deep = copy.deepcopy(arr_x) # x-2 assert arr_x.data_ptr() == arr_x_ref.data_ptr() == arr_x_deep.data_ptr() - assert arr_x.storage_ref_count() == arr_x_ref.storage_ref_count() == arr_x_deep.storage_ref_count() == 2 + assert arr_x.storage_ref_count() == arr_x_ref.storage_ref_count() == arr_x_deep.storage_ref_count() == 1 + assert arr_x.impl_ref_count() == arr_x_ref.impl_ref_count() == arr_x_deep.impl_ref_count() == 2 del arr_x - assert arr_x_deep.storage_ref_count() == 2 - assert arr_x_ref.storage_ref_count() == 2 + assert arr_x_deep.impl_ref_count() == 2 + assert arr_x_ref.impl_ref_count() == 2 del arr_x_ref - assert arr_x_deep.storage_ref_count() == 1 + assert arr_x_deep.impl_ref_count() == 1 def test_broadcast_basic_binary(self): shapes = [