Skip to content

Commit d3e5f1b

Browse files
SS-JIASS-JIA
committed
[ET-VK][runtime] Replace resize update set with generation stamps
Pull Request resolved: pytorch#22496 Dense generation stamps replace allocation-heavy `std::unordered_set` update tracking while preserving recursive `ValueList` semantics. Authored with Codex. ghstack-source-id: 423933048 @exported-using-ghexport Differential Revision: [D118543865](https://our.internmc.facebook.com/intern/diff/D118543865/)
1 parent 843f77e commit d3e5f1b

3 files changed

Lines changed: 133 additions & 12 deletions

File tree

‎backends/vulkan/runtime/graph/ComputeGraph.cpp‎

Lines changed: 31 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111

1212
#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
1313

14+
#include <algorithm>
15+
1416
#include <executorch/backends/vulkan/runtime/api/containers/StagingBuffer.h>
1517

1618
#include <executorch/backends/vulkan/runtime/graph/ops/impl/Staging.h>
@@ -245,12 +247,12 @@ bool ComputeGraph::was_value_updated(const ValueRef idx) const noexcept {
245247
return false;
246248
}
247249

248-
// Check if this ValueRef itself was updated
249-
if (updated_values_.find(idx) != updated_values_.end()) {
250+
const size_t value_idx = static_cast<size_t>(idx);
251+
if (value_idx < value_update_generations_.size() &&
252+
value_update_generations_[value_idx] == current_update_generation_) {
250253
return true;
251254
}
252255

253-
// If this is a ValueList, check each ValueRef in the list
254256
if (val_is_value_list(idx)) {
255257
const auto& value_list = values_.at(idx).toConstValueList();
256258
for (const auto& nested_idx : value_list) {
@@ -263,6 +265,26 @@ bool ComputeGraph::was_value_updated(const ValueRef idx) const noexcept {
263265
return false;
264266
}
265267

268+
void ComputeGraph::mark_value_updated(const ValueRef idx) {
269+
if (!is_valid_value_idx(idx)) {
270+
return;
271+
}
272+
if (value_update_generations_.size() < values_.size()) {
273+
value_update_generations_.resize(values_.size());
274+
}
275+
value_update_generations_[static_cast<size_t>(idx)] =
276+
current_update_generation_;
277+
}
278+
279+
void ComputeGraph::advance_update_generation() noexcept {
280+
current_update_generation_++;
281+
if (current_update_generation_ == 0) {
282+
std::fill(
283+
value_update_generations_.begin(), value_update_generations_.end(), 0);
284+
current_update_generation_ = 1;
285+
}
286+
}
287+
266288
utils::GPUMemoryLayout ComputeGraph::suggested_memory_layout(
267289
const std::vector<int64_t>& sizes) {
268290
if (config_.enable_memory_layout_override) {
@@ -775,8 +797,7 @@ void ComputeGraph::set_symint(const ValueRef idx, const int32_t val) {
775797
int32_t cur_val = read_symint(idx);
776798
if (cur_val != val) {
777799
get_symint(idx)->set(val);
778-
// Track that this ValueRef was updated
779-
updated_values_.insert(idx);
800+
mark_value_updated(idx);
780801
}
781802
}
782803

@@ -1047,6 +1068,8 @@ void ComputeGraph::maybe_cast_and_copy_from_staging(
10471068
}
10481069

10491070
void ComputeGraph::prepare() {
1071+
value_update_generations_.resize(values_.size());
1072+
10501073
#define MERGE_FIELD(field) \
10511074
static_cast<uint32_t>(std::ceil( \
10521075
std::max( \
@@ -1258,8 +1281,7 @@ void ComputeGraph::execute() {
12581281

12591282
execute_count_++;
12601283

1261-
// Clear the set of updated values at the end of inference
1262-
updated_values_.clear();
1284+
advance_update_generation();
12631285

12641286
// Reset the re-encoding flag at the end of inference
12651287
requires_reencode_ = false;
@@ -1281,7 +1303,7 @@ void ComputeGraph::resize_input(
12811303
const std::vector<int64_t>& new_sizes) {
12821304
IOValueRef io_val = inputs_.at(idx);
12831305
virtual_resize(io_val.value, new_sizes);
1284-
updated_values_.insert(io_val.staging);
1306+
mark_value_updated(io_val.staging);
12851307
}
12861308

12871309
void ComputeGraph::virtual_resize(
@@ -1290,8 +1312,7 @@ void ComputeGraph::virtual_resize(
12901312
std::vector<int64_t> cur_sizes = sizes_of(idx);
12911313
if (cur_sizes != new_sizes) {
12921314
get_tensor(idx)->virtual_resize(new_sizes);
1293-
// Track that this ValueRef was updated
1294-
updated_values_.insert(idx);
1315+
mark_value_updated(idx);
12951316
}
12961317
}
12971318

‎backends/vulkan/runtime/graph/ComputeGraph.h‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -204,8 +204,8 @@ class ComputeGraph final {
204204
// List of command buffers deferred for submission
205205
std::vector<vkapi::CommandBuffer> deferred_cmd_list_;
206206

207-
// Set to track which ValueRefs were updated during inference
208-
std::unordered_set<ValueRef> updated_values_;
207+
std::vector<uint32_t> value_update_generations_;
208+
uint32_t current_update_generation_ = 1;
209209

210210
// Cache to prevent duplicate prepacking of the same weight tensor with the
211211
// same kernel. Key is (inputValueRef, kernel_name).
@@ -1222,6 +1222,10 @@ class ComputeGraph final {
12221222

12231223
void print_readable();
12241224

1225+
private:
1226+
void mark_value_updated(const ValueRef idx);
1227+
void advance_update_generation() noexcept;
1228+
12251229
//
12261230
// Friend classes
12271231
//

‎backends/vulkan/test/vulkan_compute_api_test.cpp‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2152,6 +2152,102 @@ TEST(VulkanComputeGraphTest, test_simple_graph_with_symint) {
21522152
}
21532153
}
21542154

2155+
TEST(VulkanComputeGraphTest, was_value_updated_tracks_tensor_changes) {
2156+
GraphConfig config;
2157+
ComputeGraph graph(config);
2158+
2159+
const ValueRef tensor = graph.add_tensor({2, 4}, vkapi::kFloat);
2160+
2161+
EXPECT_FALSE(graph.was_value_updated(kDummyValueRef));
2162+
EXPECT_FALSE(graph.was_value_updated(tensor));
2163+
2164+
graph.virtual_resize(tensor, {2, 4});
2165+
EXPECT_FALSE(graph.was_value_updated(tensor));
2166+
2167+
graph.virtual_resize(tensor, {1, 4});
2168+
EXPECT_TRUE(graph.was_value_updated(tensor));
2169+
}
2170+
2171+
TEST(VulkanComputeGraphTest, was_value_updated_tracks_symint_changes) {
2172+
GraphConfig config;
2173+
ComputeGraph graph(config);
2174+
2175+
const ValueRef symint = graph.add_symint(3);
2176+
2177+
EXPECT_FALSE(graph.was_value_updated(symint));
2178+
2179+
graph.set_symint(symint, 3);
2180+
EXPECT_FALSE(graph.was_value_updated(symint));
2181+
2182+
graph.set_symint(symint, 5);
2183+
EXPECT_TRUE(graph.was_value_updated(symint));
2184+
}
2185+
2186+
TEST(VulkanComputeGraphTest, was_value_updated_checks_nested_value_lists) {
2187+
GraphConfig config;
2188+
ComputeGraph graph(config);
2189+
2190+
const ValueRef unchanged = graph.add_symint(1);
2191+
const ValueRef changed = graph.add_symint(2);
2192+
const ValueRef inner_list = graph.add_value_list({unchanged, changed});
2193+
const ValueRef outer_list =
2194+
graph.add_value_list({kDummyValueRef, inner_list});
2195+
2196+
EXPECT_FALSE(graph.was_value_updated(inner_list));
2197+
EXPECT_FALSE(graph.was_value_updated(outer_list));
2198+
2199+
graph.set_symint(changed, 3);
2200+
2201+
EXPECT_FALSE(graph.was_value_updated(unchanged));
2202+
EXPECT_TRUE(graph.was_value_updated(changed));
2203+
EXPECT_TRUE(graph.was_value_updated(inner_list));
2204+
EXPECT_TRUE(graph.was_value_updated(outer_list));
2205+
}
2206+
2207+
TEST(VulkanComputeGraphTest, resize_input_marks_staging_value_updated) {
2208+
GraphConfig config;
2209+
ComputeGraph graph(config);
2210+
2211+
const IOValueRef input = graph.add_input_tensor({2, 4}, vkapi::kFloat);
2212+
2213+
EXPECT_FALSE(graph.was_value_updated(input.value));
2214+
EXPECT_FALSE(graph.was_value_updated(input.staging));
2215+
2216+
graph.resize_input(0, {2, 4});
2217+
2218+
EXPECT_FALSE(graph.was_value_updated(input.value));
2219+
EXPECT_TRUE(graph.was_value_updated(input.staging));
2220+
}
2221+
2222+
TEST(VulkanComputeGraphTest, execute_advances_value_update_generation) {
2223+
GraphConfig config;
2224+
ComputeGraph graph(config);
2225+
2226+
const ValueRef symint = graph.add_symint(1);
2227+
const ValueRef values = graph.add_value_list({symint});
2228+
2229+
graph.prepare();
2230+
graph.set_symint(symint, 2);
2231+
2232+
EXPECT_TRUE(graph.was_value_updated(symint));
2233+
EXPECT_TRUE(graph.was_value_updated(values));
2234+
2235+
graph.execute();
2236+
2237+
EXPECT_FALSE(graph.was_value_updated(symint));
2238+
EXPECT_FALSE(graph.was_value_updated(values));
2239+
2240+
graph.set_symint(symint, 3);
2241+
2242+
EXPECT_TRUE(graph.was_value_updated(symint));
2243+
EXPECT_TRUE(graph.was_value_updated(values));
2244+
2245+
graph.execute();
2246+
2247+
EXPECT_FALSE(graph.was_value_updated(symint));
2248+
EXPECT_FALSE(graph.was_value_updated(values));
2249+
}
2250+
21552251
#define CREATE_WEIGHT_TENSOR(name, sizes, dtype, val) \
21562252
std::vector<float> data_##name(utils::multiply_integers(sizes)); \
21572253
std::fill(data_##name.begin(), data_##name.end(), val); \

0 commit comments

Comments
 (0)