Skip to content

Commit fe5d8d6

Browse files
SS-JIASS-JIA
committed
[ET-VK][runtime] Resolve invariant quantization shaders once
Pull Request resolved: pytorch#22499 Resolve resize-invariant shader names during graph construction while preserving dynamic workgroups and the M==1 quantize gate. Authored with Codex. ghstack-source-id: 423933078 @exported-using-ghexport Differential Revision: [D118543874](https://our.internmc.facebook.com/intern/diff/D118543874/)
1 parent 8d33e93 commit fe5d8d6

3 files changed

Lines changed: 228 additions & 46 deletions

File tree

‎backends/vulkan/runtime/graph/ops/impl/ChooseQParams.cpp‎

Lines changed: 5 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -34,23 +34,6 @@ void resize_choose_qparams_per_row(
3434
graph->virtual_resize(input_zeros, new_sizes);
3535
}
3636

37-
vkapi::ShaderInfo pick_choose_qparams_per_row_shader(
38-
ComputeGraph* graph,
39-
const std::vector<ArgGroup>& args,
40-
const std::vector<ValueRef>& resize_args) {
41-
(void)resize_args;
42-
43-
const ValueRef input = args.at(1).refs.at(0);
44-
const ValueRef input_zps = args.at(0).refs.at(1);
45-
46-
std::string kernel_name = "choose_qparams_per_row";
47-
add_storage_type_suffix(kernel_name, graph->storage_type_of(input));
48-
add_dtype_suffix(kernel_name, graph->dtype_of(input));
49-
add_zp_dtype_mode_suffix(kernel_name, graph->dtype_of(input_zps));
50-
51-
return VK_KERNEL_FROM_STR(kernel_name);
52-
}
53-
5437
GlobalWorkGrid pick_choose_qparams_per_row_gwg(
5538
ComputeGraph* graph,
5639
const vkapi::ShaderInfo& shader,
@@ -99,9 +82,13 @@ void add_choose_qparams_per_row_node(
9982
PushConstantDataInfo(&quant_max_val, sizeof(int32_t)),
10083
};
10184

85+
std::string kernel_name = "choose_qparams_per_row";
86+
add_storage_type_suffix(kernel_name, graph.storage_type_of(input));
87+
add_dtype_suffix(kernel_name, graph.dtype_of(input));
88+
add_zp_dtype_mode_suffix(kernel_name, graph.dtype_of(input_zps));
10289
graph.execute_nodes().emplace_back(new DynamicDispatchNode(
10390
graph,
104-
pick_choose_qparams_per_row_shader,
91+
VK_KERNEL_FROM_STR(kernel_name),
10592
pick_choose_qparams_per_row_gwg,
10693
pick_required_lwg,
10794
// Inputs and Outputs

‎backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp‎

Lines changed: 12 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -61,33 +61,6 @@ GlobalWorkGrid quantize_and_pack_4h4w_gwg(
6161
kTiledWorkGrid);
6262
}
6363

64-
vkapi::ShaderInfo pick_quantize_and_pack_4h4w_with_group_sums_shader(
65-
ComputeGraph* graph,
66-
const std::vector<ArgGroup>& args,
67-
const std::vector<ValueRef>& resize_args) {
68-
const ValueRef packed_int_input = args.at(0).refs.at(0);
69-
const ValueRef fp_input = args.at(1).refs.at(0);
70-
const ValueRef packed_input_zps = args.at(1).refs.at(2);
71-
const ValueRef group_size = resize_args.at(0);
72-
73-
const int64_t group_size_val = graph->extract_scalar<int64_t>(group_size);
74-
75-
std::string shader_name = "quantize_and_pack_4h4w_with_group_sums";
76-
if (group_size_val >= 128) {
77-
shader_name += "_o2w32";
78-
} else {
79-
shader_name += "_o4w16";
80-
}
81-
82-
add_storage_type_suffix(
83-
shader_name, graph->storage_type_of(packed_int_input));
84-
add_storage_type_suffix(shader_name, graph->storage_type_of(fp_input));
85-
add_dtype_suffix(shader_name, graph->dtype_of(fp_input));
86-
add_zp_dtype_mode_suffix(shader_name, graph->dtype_of(packed_input_zps));
87-
88-
return VK_KERNEL_FROM_STR(shader_name);
89-
}
90-
9164
GlobalWorkGrid pick_quantize_and_pack_4h4w_with_group_sums_gwg(
9265
ComputeGraph* graph,
9366
const vkapi::ShaderInfo& shader,
@@ -199,9 +172,20 @@ void add_quantize_and_pack_4h4w_with_group_sums_node(
199172
const int32_t group_size_val = graph.extract_scalar<int32_t>(group_size);
200173
const int32_t blocks_per_group = utils::div_up(group_size_val, int32_t(4));
201174

175+
std::string shader_name = "quantize_and_pack_4h4w_with_group_sums";
176+
if (group_size_val >= 128) {
177+
shader_name += "_o2w32";
178+
} else {
179+
shader_name += "_o4w16";
180+
}
181+
add_storage_type_suffix(shader_name, graph.storage_type_of(packed_int_input));
182+
add_storage_type_suffix(shader_name, graph.storage_type_of(fp_input));
183+
add_dtype_suffix(shader_name, graph.dtype_of(fp_input));
184+
add_zp_dtype_mode_suffix(shader_name, graph.dtype_of(packed_input_zps));
185+
202186
graph.execute_nodes().emplace_back(new DynamicDispatchNode(
203187
graph,
204-
pick_quantize_and_pack_4h4w_with_group_sums_shader,
188+
VK_KERNEL_FROM_STR(shader_name),
205189
pick_quantize_and_pack_4h4w_with_group_sums_gwg,
206190
pick_required_lwg,
207191
// Inputs and Outputs

‎backends/vulkan/test/vulkan_compute_api_test.cpp‎

Lines changed: 211 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
#include <executorch/backends/vulkan/runtime/graph/ops/DispatchNode.h>
3333

3434
#include <executorch/backends/vulkan/runtime/graph/ops/impl/Int8x4Staging.h>
35+
#include <executorch/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.h>
3536

3637
#include <executorch/backends/vulkan/runtime/utils/VecUtils.h>
3738

@@ -2379,6 +2380,216 @@ TEST(VulkanComputeGraphTest, execute_advances_value_update_generation) {
23792380
EXPECT_FALSE(graph.was_value_updated(values));
23802381
}
23812382

2383+
TEST(VulkanComputeGraphTest, choose_qparams_handles_dynamic_row_counts) {
2384+
constexpr int64_t kMaxM = 8;
2385+
constexpr int64_t kK = 128;
2386+
2387+
GraphConfig config;
2388+
config.enable_querypool = true;
2389+
config.expect_dynamic_shapes = true;
2390+
ComputeGraph graph(config);
2391+
2392+
const IOValueRef input =
2393+
graph.add_input_tensor({kMaxM, kK}, vkapi::kFloat, utils::kBuffer);
2394+
const ValueRef quant_min = graph.add_scalar<int64_t>(-128);
2395+
const ValueRef quant_max = graph.add_scalar<int64_t>(127);
2396+
const ValueRef scales = graph.add_tensor(
2397+
{kMaxM}, vkapi::kFloat, utils::kTexture3D, utils::kWidthPacked);
2398+
const ValueRef zero_points = graph.add_tensor(
2399+
{kMaxM}, vkapi::kChar, utils::kTexture3D, utils::kWidthPacked);
2400+
2401+
VK_GET_OP_FN("etvk.choose_qparams_per_row.default")
2402+
(graph, {input.value, quant_min, quant_max, scales, zero_points});
2403+
2404+
const ValueRef scales_staging = graph.set_output_tensor(scales);
2405+
const ValueRef zero_points_staging = graph.set_output_tensor(zero_points);
2406+
2407+
graph.prepare();
2408+
graph.prepack();
2409+
2410+
for (const int64_t M : std::vector<int64_t>{kMaxM, 4, 1, kMaxM}) {
2411+
graph.resize_input(0, {M, kK});
2412+
graph.propagate_resize();
2413+
2414+
EXPECT_EQ(graph.sizes_of(scales), std::vector<int64_t>({M}));
2415+
EXPECT_EQ(graph.sizes_of(zero_points), std::vector<int64_t>({M}));
2416+
2417+
std::vector<float> input_data(M * kK);
2418+
for (int64_t m = 0; m < M; ++m) {
2419+
std::fill_n(input_data.begin() + m * kK, kK, float(m + 1));
2420+
}
2421+
graph.maybe_cast_and_copy_into_staging(
2422+
input.staging, input_data.data(), input_data.size(), vkapi::kFloat);
2423+
2424+
graph.execute();
2425+
2426+
std::vector<float> scale_data(M);
2427+
std::vector<int8_t> zero_point_data(M);
2428+
graph.maybe_cast_and_copy_from_staging(
2429+
scales_staging, scale_data.data(), scale_data.size(), vkapi::kFloat);
2430+
graph.maybe_cast_and_copy_from_staging(
2431+
zero_points_staging,
2432+
zero_point_data.data(),
2433+
zero_point_data.size(),
2434+
vkapi::kChar);
2435+
2436+
for (int64_t m = 0; m < M; ++m) {
2437+
EXPECT_NEAR(scale_data[m], float(m + 1) / 255.0f, 1e-6f);
2438+
EXPECT_EQ(zero_point_data[m], -128);
2439+
}
2440+
2441+
graph.context()->querypool().extract_results();
2442+
const auto shader_results =
2443+
graph.context()->querypool().get_shader_timestamp_data();
2444+
const auto choose_result = std::find_if(
2445+
shader_results.begin(), shader_results.end(), [](const auto& result) {
2446+
return result.kernel_name.find("choose_qparams_per_row") !=
2447+
std::string::npos;
2448+
});
2449+
ASSERT_NE(choose_result, shader_results.end());
2450+
EXPECT_EQ(choose_result->metadata.gwg[0], 1u);
2451+
EXPECT_EQ(
2452+
choose_result->metadata.gwg[1],
2453+
utils::div_up_4(utils::safe_downcast<uint32_t>(M)));
2454+
EXPECT_EQ(choose_result->metadata.gwg[2], 1u);
2455+
EXPECT_EQ(choose_result->metadata.lwg[0], 64u);
2456+
EXPECT_EQ(choose_result->metadata.lwg[1], 1u);
2457+
EXPECT_EQ(choose_result->metadata.lwg[2], 1u);
2458+
}
2459+
}
2460+
2461+
void test_quantize_and_pack_handles_dynamic_row_counts(
2462+
const int64_t group_size_value,
2463+
const utils::uvec3& expected_local_wg_size) {
2464+
if (!api::context()->adapter_ptr()->supports_int8_dot_product()) {
2465+
GTEST_SKIP() << "Quantize and pack requires integer dot product support";
2466+
}
2467+
2468+
constexpr int64_t kMaxM = 8;
2469+
constexpr int64_t kK = 128;
2470+
const int64_t num_groups = kK / group_size_value;
2471+
const int64_t max_m4 = utils::div_up(kMaxM, int64_t(4));
2472+
2473+
GraphConfig config;
2474+
config.enable_querypool = true;
2475+
config.expect_dynamic_shapes = true;
2476+
ComputeGraph graph(config);
2477+
2478+
const IOValueRef input =
2479+
graph.add_input_tensor({kMaxM, kK}, vkapi::kFloat, utils::kBuffer);
2480+
const ValueRef quant_min = graph.add_scalar<int64_t>(-128);
2481+
const ValueRef quant_max = graph.add_scalar<int64_t>(127);
2482+
const ValueRef scales = graph.add_tensor(
2483+
{kMaxM}, vkapi::kFloat, utils::kTexture3D, utils::kWidthPacked);
2484+
const ValueRef zero_points = graph.add_tensor(
2485+
{kMaxM}, vkapi::kChar, utils::kTexture3D, utils::kWidthPacked);
2486+
2487+
VK_GET_OP_FN("etvk.choose_qparams_per_row.default")
2488+
(graph, {input.value, quant_min, quant_max, scales, zero_points});
2489+
2490+
const ValueRef packed_input = graph.add_tensor(
2491+
{kMaxM, kK}, vkapi::kInt8x4, utils::kBuffer, utils::kPackedInt8_4H4W);
2492+
const ValueRef input_sums = graph.add_tensor(
2493+
{num_groups * max_m4 * 4},
2494+
vkapi::kInt,
2495+
utils::kBuffer,
2496+
utils::kWidthPacked);
2497+
const ValueRef group_size = graph.add_scalar<int64_t>(group_size_value);
2498+
const QuantizationConfig input_quant_config(
2499+
8, kPerChannel, {1, kK}, false, true);
2500+
2501+
add_quantize_and_pack_4h4w_with_group_sums_node(
2502+
graph,
2503+
input_quant_config,
2504+
input.value,
2505+
input_sums,
2506+
scales,
2507+
zero_points,
2508+
packed_input,
2509+
group_size);
2510+
2511+
const ValueRef packed_input_staging = graph.set_output_tensor(packed_input);
2512+
const ValueRef input_sums_staging = graph.set_output_tensor(input_sums);
2513+
2514+
graph.prepare();
2515+
graph.prepack();
2516+
2517+
for (const int64_t M : std::vector<int64_t>{kMaxM, 4, 1, kMaxM}) {
2518+
graph.resize_input(0, {M, kK});
2519+
graph.propagate_resize();
2520+
2521+
std::vector<float> input_data(M * kK);
2522+
for (int64_t m = 0; m < M; ++m) {
2523+
std::fill_n(input_data.begin() + m * kK, kK, float(m + 1));
2524+
}
2525+
graph.maybe_cast_and_copy_into_staging(
2526+
input.staging, input_data.data(), input_data.size(), vkapi::kFloat);
2527+
2528+
graph.execute();
2529+
2530+
graph.context()->querypool().extract_results();
2531+
const auto shader_results =
2532+
graph.context()->querypool().get_shader_timestamp_data();
2533+
const auto quantize_result = std::find_if(
2534+
shader_results.begin(), shader_results.end(), [](const auto& result) {
2535+
return result.kernel_name.find(
2536+
"quantize_and_pack_4h4w_with_group_sums") !=
2537+
std::string::npos;
2538+
});
2539+
2540+
if (M == 1) {
2541+
EXPECT_EQ(quantize_result, shader_results.end());
2542+
continue;
2543+
}
2544+
2545+
ASSERT_NE(quantize_result, shader_results.end());
2546+
EXPECT_EQ(
2547+
quantize_result->metadata.gwg[0],
2548+
utils::safe_downcast<uint32_t>(num_groups));
2549+
EXPECT_EQ(
2550+
quantize_result->metadata.gwg[1],
2551+
utils::div_up_4(utils::safe_downcast<uint32_t>(M)));
2552+
EXPECT_EQ(quantize_result->metadata.gwg[2], 1u);
2553+
EXPECT_EQ(quantize_result->metadata.lwg[0], expected_local_wg_size[0]);
2554+
EXPECT_EQ(quantize_result->metadata.lwg[1], expected_local_wg_size[1]);
2555+
EXPECT_EQ(quantize_result->metadata.lwg[2], expected_local_wg_size[2]);
2556+
2557+
const size_t packed_numel = graph.staging_buffer_numel_of(packed_input);
2558+
std::vector<int32_t> packed_data(packed_numel);
2559+
graph.maybe_cast_and_copy_from_staging(
2560+
packed_input_staging,
2561+
packed_data.data(),
2562+
packed_data.size(),
2563+
vkapi::kInt8x4);
2564+
for (int64_t i = 0; i < M * kK / 4; ++i) {
2565+
EXPECT_EQ(packed_data[i], 0x7f7f7f7f);
2566+
}
2567+
2568+
std::vector<int32_t> sums_data(num_groups * max_m4 * 4);
2569+
graph.maybe_cast_and_copy_from_staging(
2570+
input_sums_staging, sums_data.data(), sums_data.size(), vkapi::kInt);
2571+
const int64_t current_m4 = utils::div_up(M, int64_t(4));
2572+
for (int64_t group = 0; group < num_groups; ++group) {
2573+
for (int64_t m = 0; m < M; ++m) {
2574+
EXPECT_EQ(
2575+
sums_data[group * current_m4 * 4 + m], 127 * group_size_value);
2576+
}
2577+
}
2578+
}
2579+
}
2580+
2581+
TEST(
2582+
VulkanComputeGraphTest,
2583+
quantize_and_pack_handles_dynamic_row_counts_with_small_groups) {
2584+
test_quantize_and_pack_handles_dynamic_row_counts(32, {4u, 1u, 16u});
2585+
}
2586+
2587+
TEST(
2588+
VulkanComputeGraphTest,
2589+
quantize_and_pack_handles_dynamic_row_counts_with_large_groups) {
2590+
test_quantize_and_pack_handles_dynamic_row_counts(128, {2u, 1u, 32u});
2591+
}
2592+
23822593
#define CREATE_WEIGHT_TENSOR(name, sizes, dtype, val) \
23832594
std::vector<float> data_##name(utils::multiply_integers(sizes)); \
23842595
std::fill(data_##name.begin(), data_##name.end(), val); \

0 commit comments

Comments
 (0)