[executorch][vulkan] Accept 1-D activations in 8da4w linear

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<int64_t> 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;
   }
 

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;
 

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

@@ -32,11 +32,10 @@ bool is_gemv(ComputeGraph* graph, const ValueRef& fp_input) {
 std::tuple<int64_t, int64_t> get_quantized_input_num_blocks(
     ComputeGraph& graph,
     const ValueRef input) {
-  std::vector<int64_t> input_sizes = graph.sizes_of(input);
-  const int64_t ndim = graph.dim_of(input);
+  const std::vector<int64_t> 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));

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<int64_t>({kK}),
+    graph.add_none(), // target_dtype
+    graph.add_scalar<int64_t>(-128),
+    graph.add_scalar<int64_t>(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<float> 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<int64_t>(kK)));
+}
+
 void test_quantize_and_pack_handles_dynamic_row_counts(
     const int64_t group_size_value,
     const utils::uvec3& expected_local_wg_size) {