Skip to content

[ET-VK] Introduce vTensorPtr to prevent reference invalidation and remove get_val() API #2978

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

Closed
wants to merge 5 commits into from
Closed

[ET-VK] Introduce vTensorPtr to prevent reference invalidation and remove get_val() API #2978

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
04975bd
[ET-VK] Introduce `vTensorPtr` to prevent reference invalidation and …
SS-JIA Apr 10, 2024
7795ca2
Update on "[ET-VK] Introduce `vTensorPtr` to prevent reference invali…
SS-JIA Apr 10, 2024
5c421ba
Update on "[ET-VK] Introduce `vTensorPtr` to prevent reference invali…
SS-JIA Apr 11, 2024
ff3213f
Update on "[ET-VK] Introduce `vTensorPtr` to prevent reference invali…
SS-JIA Apr 11, 2024
b99f1d8
Update on "[ET-VK] Introduce `vTensorPtr` to prevent reference invali…
SS-JIA Apr 11, 2024
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
18 changes: 9 additions & 9 deletions backends/vulkan/runtime/VulkanBackend.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -334,18 +334,18 @@ bool maybe_resize_input(
const size_t input_i,
exec_aten::Tensor& et_tensor) {
ValueRef in_tensor_ref = graph->inputs()[input_i].value;
vTensor& in_tensor = graph->get_val(in_tensor_ref).toTensor();
vTensorPtr in_tensor = graph->get_tensor(in_tensor_ref);

ET_CHECK_MSG(
et_tensor.dim() == in_tensor.sizes().size(),
et_tensor.dim() == in_tensor->sizes().size(),
"Cannot resize input tensor: old ndim %zu does not match new ndim %zu",
static_cast<size_t>(in_tensor.sizes().size()),
static_cast<size_t>(in_tensor->sizes().size()),
static_cast<size_t>(et_tensor.dim()));

bool should_resize = false;
std::vector<int64_t> new_sizes(et_tensor.dim());
for (size_t i = 0; i < et_tensor.dim(); i++) {
if (in_tensor.sizes()[i] != et_tensor.sizes()[i]) {
if (in_tensor->sizes()[i] != et_tensor.sizes()[i]) {
should_resize = true;
}
new_sizes.at(i) = et_tensor.sizes()[i];
Expand All @@ -356,9 +356,9 @@ bool maybe_resize_input(
}

ET_CHECK_MSG(
in_tensor.numel() == et_tensor.numel(),
in_tensor->numel() == et_tensor.numel(),
"Vulkan tensor numel %zu does not match ET tensor numel %zu",
static_cast<size_t>(in_tensor.numel()),
static_cast<size_t>(in_tensor->numel()),
static_cast<size_t>(et_tensor.numel()));

return should_resize;
Expand All @@ -369,12 +369,12 @@ void maybe_resize_output(
const size_t output_i,
exec_aten::Tensor& et_tensor) {
ValueRef out_tensor_ref = graph->outputs()[output_i].value;
vTensor& out_tensor = graph->get_val(out_tensor_ref).toTensor();
vTensorPtr out_tensor = graph->get_tensor(out_tensor_ref);

exec_aten::SizesType new_output_size[kTensorDimensionLimit];
size_t ndim = out_tensor.sizes().size();
size_t ndim = out_tensor->sizes().size();
for (int i = 0; i < ndim; ++i) {
new_output_size[i] = out_tensor.sizes()[i];
new_output_size[i] = out_tensor->sizes()[i];
}

exec_aten::ArrayRef<exec_aten::SizesType> output_size{new_output_size, ndim};
Expand Down
105 changes: 86 additions & 19 deletions backends/vulkan/runtime/graph/ComputeGraph.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,39 @@

namespace vkcompute {

//
// VTensorPtr
//

#define VALUE_PTR_CLASS_IMPL(classname, ctype, type_name) \
classname::classname(ComputeGraph* const graph, const ValueRef idx) \
: graph_(graph), ptr_(&(graph_->values_.at(idx).to##type_name())) { \
graph_->values_in_use_++; \
} \
ctype* classname::operator->() const { \
return ptr_; \
} \
ctype& classname::operator*() const { \
return *ptr_; \
} \
classname::~classname() { \
graph_->values_in_use_--; \
}

VALUE_PTR_CLASS_IMPL(vTensorPtr, vTensor, Tensor)
VALUE_PTR_CLASS_IMPL(TensorRefPtr, TensorRef, TensorRef)
VALUE_PTR_CLASS_IMPL(StagingPtr, api::StorageBuffer, Staging)
VALUE_PTR_CLASS_IMPL(IntListPtr, std::vector<int64_t>, IntList)
VALUE_PTR_CLASS_IMPL(DoubleListPtr, std::vector<double>, DoubleList)
VALUE_PTR_CLASS_IMPL(BoolListPtr, std::vector<bool>, BoolList)
VALUE_PTR_CLASS_IMPL(ValueListPtr, std::vector<ValueRef>, ValueList)

#undef VALUE_PTR_CLASS_IMPL

//
// ComputeGraph
//

ComputeGraph::ComputeGraph(GraphConfig config)
: config_{config},
prepack_descriptor_counts_{},
Expand Down Expand Up @@ -105,6 +138,35 @@ api::GPUMemoryLayout ComputeGraph::suggested_memory_layout(
return api::kChannelsPacked;
}

void ComputeGraph::check_no_active_value_ptrs() {
VK_CHECK_COND(
values_in_use_ == 0,
"Make sure that there are no pointers stored from the return values of "
"`ComputeGraph::get_*()` functions in scope before adding Values to the "
"graph. Modifying the graph's values may cause existing pointers to be "
"invalidated.");
}

std::vector<int64_t> ComputeGraph::get_sizes_of(ValueRef idx) {
Value& val = values_.at(idx);
if (val.isTensor()) {
return val.toTensor().sizes();
} else if (val.isTensorRef()) {
return val.toTensorRef().sizes;
}
VK_THROW("Could not get sizes of value with type ", val.type());
}

api::ScalarType ComputeGraph::get_dtype_of(ValueRef idx) {
Value& val = values_.at(idx);
if (val.isTensor()) {
return val.toTensor().dtype();
} else if (val.isTensorRef()) {
return val.toTensorRef().dtype;
}
VK_THROW("Could not get dtype of value with type ", val.type());
}

ValueRef ComputeGraph::add_tensor(
const std::vector<int64_t>& sizes,
const api::ScalarType dtype,
Expand All @@ -114,6 +176,7 @@ ValueRef ComputeGraph::add_tensor(
bool allocate_memory = shared_object_idx < 0;

ValueRef idx(static_cast<int>(values_.size()));
check_no_active_value_ptrs();
values_.emplace_back(vTensor(
context(), sizes, dtype, storage_type, memory_layout, allocate_memory));

Expand All @@ -133,18 +196,17 @@ ValueRef ComputeGraph::add_tensor(
}

ValueRef ComputeGraph::add_tensor_like(
const ValueRef vref,
const ValueRef idx,
const api::StorageType storage_type,
const api::GPUMemoryLayout memory_layout) {
TensorRef& tref = get_val(vref).toTensorRef();
return add_tensor(tref.sizes, tref.dtype, storage_type, memory_layout);
return add_tensor(
get_sizes_of(idx), get_dtype_of(idx), storage_type, memory_layout);
}

ValueRef ComputeGraph::add_tensor_like(
const ValueRef vref,
const ValueRef idx,
const api::GPUMemoryLayout memory_layout) {
TensorRef& tref = get_val(vref).toTensorRef();
return add_tensor(tref.sizes, tref.dtype, memory_layout);
return add_tensor(get_sizes_of(idx), get_dtype_of(idx), memory_layout);
}

ValueRef ComputeGraph::add_tensor(
Expand All @@ -160,6 +222,7 @@ ValueRef ComputeGraph::add_tensorref(
const api::ScalarType dtype,
const void* const data) {
ValueRef idx(static_cast<int>(values_.size()));
check_no_active_value_ptrs();
values_.emplace_back(TensorRef(sizes, dtype, data));
return idx;
}
Expand All @@ -168,24 +231,28 @@ ValueRef ComputeGraph::add_staging(
const api::ScalarType dtype,
const size_t numel) {
ValueRef idx(static_cast<int>(values_.size()));
check_no_active_value_ptrs();
values_.emplace_back(api::StorageBuffer(context(), dtype, numel));
return idx;
}

ValueRef ComputeGraph::add_none() {
ValueRef idx(static_cast<int>(values_.size()));
check_no_active_value_ptrs();
values_.emplace_back();
return idx;
}

ValueRef ComputeGraph::add_value_list(std::vector<ValueRef>&& value) {
ValueRef idx(static_cast<int>(values_.size()));
check_no_active_value_ptrs();
values_.emplace_back(std::move(value));
return idx;
}

ValueRef ComputeGraph::add_string(std::string&& str) {
ValueRef idx(static_cast<int>(values_.size()));
check_no_active_value_ptrs();
values_.emplace_back(std::move(str));
return idx;
}
Expand All @@ -194,8 +261,9 @@ ValueRef ComputeGraph::set_input_tensor(
const ValueRef idx,
const bool use_staging) {
if (use_staging) {
vTensor& tensor = get_val(idx).toTensor();
ValueRef staging_idx = add_staging(tensor.dtype(), tensor.gpu_numel());
api::ScalarType dtype = get_tensor(idx)->dtype();
size_t gpu_numel = get_tensor(idx)->gpu_numel();
ValueRef staging_idx = add_staging(dtype, gpu_numel);
add_staging_to_tensor_node(*this, staging_idx, idx);
inputs_.push_back({idx, staging_idx});
return staging_idx;
Expand All @@ -208,8 +276,9 @@ ValueRef ComputeGraph::set_output_tensor(
const ValueRef idx,
const bool use_staging) {
if (use_staging) {
vTensor& tensor = get_val(idx).toTensor();
ValueRef staging_idx = add_staging(tensor.dtype(), tensor.gpu_numel());
api::ScalarType dtype = get_tensor(idx)->dtype();
size_t gpu_numel = get_tensor(idx)->gpu_numel();
ValueRef staging_idx = add_staging(dtype, gpu_numel);
add_tensor_to_staging_node(*this, idx, staging_idx);
outputs_.push_back({idx, staging_idx});
return staging_idx;
Expand All @@ -229,20 +298,18 @@ void ComputeGraph::copy_into_staging(
const ValueRef idx,
const void* data,
const size_t numel) {
Value& in_val = get_val(idx);
api::StorageBuffer& staging = in_val.toStaging();
size_t nbytes = numel * api::element_size(staging.dtype());
copy_ptr_to_staging(data, staging, nbytes);
StagingPtr staging = get_staging(idx);
size_t nbytes = numel * api::element_size(staging->dtype());
copy_ptr_to_staging(data, *staging, nbytes);
}

void ComputeGraph::copy_from_staging(
const ValueRef idx,
void* data,
const size_t numel) {
Value& out_val = get_val(idx);
api::StorageBuffer& staging = out_val.toStaging();
size_t nbytes = numel * api::element_size(staging.dtype());
copy_staging_to_ptr(staging, data, nbytes);
StagingPtr staging = get_staging(idx);
size_t nbytes = numel * api::element_size(staging->dtype());
copy_staging_to_ptr(*staging, data, nbytes);
}

void ComputeGraph::prepare() {
Expand Down Expand Up @@ -308,7 +375,7 @@ void ComputeGraph::resize_input(
const int64_t idx,
const std::vector<int64_t>& new_sizes) {
IOValueRef io_val = inputs_.at(idx);
get_val(io_val.value).toTensor().virtual_resize(new_sizes);
get_tensor(io_val.value)->virtual_resize(new_sizes);
}

void ComputeGraph::propagate_resize() {
Expand Down
Loading