diff --git a/backends/vulkan/runtime/graph/ComputeGraph.cpp b/backends/vulkan/runtime/graph/ComputeGraph.cpp index a6b3c4f6347..e3997b26988 100644 --- a/backends/vulkan/runtime/graph/ComputeGraph.cpp +++ b/backends/vulkan/runtime/graph/ComputeGraph.cpp @@ -426,10 +426,10 @@ bool ComputeGraph::is_standard_channels_packed_texture_tensor( bool ComputeGraph::is_2d_matrix(const ValueRef idx) const { std::vector sizes = sizes_of(idx); const size_t ndim = sizes.size(); - if (sizes.size() < 2) { + if (ndim == 0) { return false; } - if (sizes.size() == 2) { + if (ndim <= 2) { return true; } diff --git a/backends/vulkan/runtime/graph/ComputeGraph.h b/backends/vulkan/runtime/graph/ComputeGraph.h index 6b78d8b50fb..eb0a0e168d2 100644 --- a/backends/vulkan/runtime/graph/ComputeGraph.h +++ b/backends/vulkan/runtime/graph/ComputeGraph.h @@ -440,9 +440,9 @@ class ComputeGraph final { bool is_standard_channels_packed_texture_tensor(const ValueRef idx) const; /* - * Checks that the value at `idx` is either a 2D tensor, or if the tensor has - * more than 2 dims, the outermost dims have size of 1, i.e. can be squeezed - * to be a 2D tensor. + * Checks that the value at `idx` is a 1D tensor (a single row), a 2D tensor, + * or a tensor whose outermost dims have size of 1, i.e. can be squeezed to + * be a 2D tensor. */ bool is_2d_matrix(const ValueRef idx) const; diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp index 7cf03d98a94..5834f2551e1 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp @@ -32,11 +32,10 @@ bool is_gemv(ComputeGraph* graph, const ValueRef& fp_input) { std::tuple get_quantized_input_num_blocks( ComputeGraph& graph, const ValueRef input) { - std::vector input_sizes = graph.sizes_of(input); - const int64_t ndim = graph.dim_of(input); + const std::vector input_sizes = graph.sizes_of(input); - const int64_t M = input_sizes.at(ndim - 2); - const int64_t K = input_sizes.at(ndim - 1); + const int64_t M = utils::val_at(-2, input_sizes); + const int64_t K = utils::val_at(-1, input_sizes); const int64_t num_blocks_M = utils::div_up(M, int64_t(4)); const int64_t num_blocks_K = utils::div_up(K, int64_t(4)); diff --git a/backends/vulkan/test/vulkan_compute_api_test.cpp b/backends/vulkan/test/vulkan_compute_api_test.cpp index 0374505d0bb..b087b0e4baa 100644 --- a/backends/vulkan/test/vulkan_compute_api_test.cpp +++ b/backends/vulkan/test/vulkan_compute_api_test.cpp @@ -2465,6 +2465,84 @@ TEST(VulkanComputeGraphTest, choose_qparams_handles_dynamic_row_counts) { } } +TEST(VulkanComputeGraphTest, choose_qparams_affine_handles_1d_input) { + constexpr int64_t kK = 128; + + ComputeGraph graph(GraphConfig{}); + const IOValueRef input = + graph.add_input_tensor({kK}, vkapi::kFloat, utils::kBuffer); + const ValueRef scales = graph.add_tensor( + {1}, vkapi::kFloat, utils::kTexture3D, utils::kWidthPacked); + const ValueRef zero_points = graph.add_tensor( + {1}, vkapi::kChar, utils::kTexture3D, utils::kWidthPacked); + + VK_GET_OP_FN("torchao.choose_qparams_affine.default") + (graph, + {input.value, + graph.add_none(), // mapping_type + graph.add_scalar_list({kK}), + graph.add_none(), // target_dtype + graph.add_scalar(-128), + graph.add_scalar(127), + graph.add_none(), // eps + graph.add_none(), // scale_dtype + graph.add_none(), // zero_point_dtype + graph.add_value_list({scales, zero_points})}); + + const ValueRef scales_staging = graph.set_output_tensor(scales); + const ValueRef zero_points_staging = graph.set_output_tensor(zero_points); + + graph.prepare(); + graph.prepack(); + + // The row's maximum is its last element, so the whole row must be reduced. + std::vector input_data(kK); + for (int64_t k = 0; k < kK; ++k) { + input_data[k] = 0.5f * float(k); + } + graph.maybe_cast_and_copy_into_staging( + input.staging, input_data.data(), input_data.size(), vkapi::kFloat); + + graph.execute(); + + float scale = 0; + int8_t zero_point = 0; + graph.maybe_cast_and_copy_from_staging( + scales_staging, &scale, 1, vkapi::kFloat); + graph.maybe_cast_and_copy_from_staging( + zero_points_staging, &zero_point, 1, vkapi::kChar); + EXPECT_NEAR(scale, input_data.back() / 255.0f, 1e-6f); + EXPECT_EQ(zero_point, -128); +} + +TEST(VulkanComputeGraphTest, quantize_and_pack_accepts_1d_input) { + constexpr int64_t kK = 128; + + ComputeGraph graph(GraphConfig{}); + const IOValueRef input = + graph.add_input_tensor({kK}, vkapi::kFloat, utils::kBuffer); + const ValueRef scales = graph.add_tensor( + {1}, vkapi::kFloat, utils::kTexture3D, utils::kWidthPacked); + const ValueRef zero_points = graph.add_tensor( + {1}, vkapi::kChar, utils::kTexture3D, utils::kWidthPacked); + const ValueRef packed_input = graph.add_tensor( + {1, kK}, vkapi::kInt8x4, utils::kBuffer, utils::kPackedInt8_4H4W); + const ValueRef input_sums = + graph.add_tensor({4}, vkapi::kInt, utils::kBuffer, utils::kWidthPacked); + const QuantizationConfig input_quant_config( + 8, kPerChannel, {1, kK}, false, true); + + EXPECT_NO_THROW(add_quantize_and_pack_4h4w_with_group_sums_node( + graph, + input_quant_config, + input.value, + input_sums, + scales, + zero_points, + packed_input, + graph.add_scalar(kK))); +} + void test_quantize_and_pack_handles_dynamic_row_counts( const int64_t group_size_value, const utils::uvec3& expected_local_wg_size) {