Skip to content

Add TensorImpl #35

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
xytpai merged 1 commit into from
Aug 14, 2025
Merged

Add TensorImpl #35

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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})

Expand Down
177 changes: 31 additions & 146 deletions src/core/include/tensor.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,26 +1,11 @@
#pragma once

#include <iostream>
#include <vector>
#include <string>
#include <tuple>
#include <optional>

#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<int64_t> shape, ScalarType dtype, int device = 0);
Tensor empty(int64_t *shape, int ndim, ScalarType dtype, int device, bool inverse = false);
Expand All @@ -30,191 +15,91 @@ 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, int vec_size>
struct d_array {
T val[vec_size] = {0};
d_array() {
}
d_array(double value) {
*reinterpret_cast<double *>(&val[0]) = value;
}
operator double() const {
return *reinterpret_cast<double *>(const_cast<T *>(&val[0]));
}
T &operator[](int i) {
return val[i];
}
T const &operator[](int i) const {
return val[i];
}
bool equals(d_array<T, vec_size> &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<int64_t, MAX_TENSOR_DIMS> dim_t;
using any_t = d_array<char, 256>;
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 <typename T>
T *data_ptr() const {
return reinterpret_cast<T *>(ptr_.get());
}
bool defined() const {
return static_cast<bool>(ptr_);
}
};

class Tensor {
int dim_;
dim_t shape_;
dim_t stride_;
ScalarType dtype_;
int64_t numel_;
intrusive_ptr<TensorStorage> storage_;
int64_t storage_offset_ = 0;
bool is_contiguous_ = true;

void new_storage_(int device);

private:
intrusive_ptr<TensorImpl> impl_;
friend Tensor empty(std::vector<int64_t> 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);
friend Tensor empty_strided(std::vector<int64_t> shape, std::vector<int64_t> strides, ScalarType dtype, int device);
friend Tensor empty_like_reduced(const Tensor &self, int dim, ScalarType dtype);
friend Tensor zeros(std::vector<int64_t> shape, ScalarType dtype, int device);

Tensor(std::vector<int64_t> &shape, ScalarType dtype);
Tensor(std::vector<int64_t> &shape, std::vector<int64_t> &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<int64_t> sizes() const {
std::vector<int64_t> vec(shape_.val, shape_.val + dim_);
return vec;
return impl_.get()->sizes();
}
std::vector<int64_t> strides() const {
std::vector<int64_t> 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 <typename T>
T *data_ptr() const {
return reinterpret_cast<T *>(storage_.get()->data_ptr()) + storage_offset_;
return impl_.get()->data_ptr<T>();
}
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<TensorStorage> 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;
oss << *this;
return oss.str();
}
bool is_contiguous() const {
return is_contiguous_;
return impl_.get()->is_contiguous();
}

void copy_from_cpu_ptr(void *ptr);
Expand Down
Loading