[ET-VK] Enable buffer implementation of aten.linear

## Changes

As title. Extend the existing buffer implementation of `matmul` to support the linear operator as well.

Differential Revision: [D65277712](https://our.internmc.facebook.com/intern/diff/D65277712/)

[ghstack-poisoned]

Author: SS-JIA

backends/vulkan/runtime/graph/ops/glsl/matmul_naive_buffer.glsl

@@ -16,21 +16,23 @@ ${define_required_extensions(DTYPE)}
 
 layout(std430) buffer;
 
-${layout_declare_tensor(0, "w", "t_out", DTYPE, "buffer")}
-${layout_declare_tensor(1, "r", "t_mat1", DTYPE, "buffer")}
-${layout_declare_tensor(2, "r", "t_mat2", DTYPE, "buffer")}
-${layout_declare_ubo(3, "ivec4", "out_sizes")}
-${layout_declare_ubo(4, "ivec4", "out_strides")}
-${layout_declare_ubo(5, "ivec4", "mat1_sizes")}
-${layout_declare_ubo(6, "ivec4", "mat1_strides")}
-${layout_declare_ubo(7, "ivec4", "mat2_sizes")}
-${layout_declare_ubo(8, "ivec4", "mat2_strides")}
-${layout_declare_ubo(9, "int", "out_numel")}
+${layout_declare_tensor(B, "w", "t_out", DTYPE, "buffer")}
+${layout_declare_tensor(B, "r", "t_mat1", DTYPE, "buffer")}
+${layout_declare_tensor(B, "r", "t_mat2", DTYPE, "buffer")}
+${layout_declare_ubo(B, "ivec4", "out_sizes")}
+${layout_declare_ubo(B, "ivec4", "out_strides")}
+${layout_declare_ubo(B, "ivec4", "mat1_sizes")}
+${layout_declare_ubo(B, "ivec4", "mat1_strides")}
+${layout_declare_ubo(B, "ivec4", "mat2_sizes")}
+${layout_declare_ubo(B, "ivec4", "mat2_strides")}
+${layout_declare_ubo(B, "int", "out_numel")}
 
 #include "indexing_utils.h"
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
+${layout_declare_spec_const(C, "int", "mat2_is_transposed", "0")}
+
 void main() {
   const ivec4 out_bufix = ivec4(
       gl_GlobalInvocationID.x,
@@ -44,15 +46,28 @@ void main() {
 
   int mat1_bufi = tidx_to_bufi(
       ivec4(0, out_bufix.y, out_bufix.z, out_bufix.w), mat1_strides);
-  int mat2_bufi = tidx_to_bufi(
-      ivec4(out_bufix.x, 0, out_bufix.z, out_bufix.w), mat2_strides);
+  int mat2_bufi;
+  if (mat2_is_transposed > 0) {
+    mat2_bufi = tidx_to_bufi(
+        ivec4(0, out_bufix.x, 0, 0), mat2_strides);
+  } else {
+    mat2_bufi = tidx_to_bufi(
+        ivec4(out_bufix.x, 0, out_bufix.z, out_bufix.w), mat2_strides);
+  }
+
+  int mat2_stride;
+  if (mat2_is_transposed > 0) {
+    mat2_stride = mat2_strides.x;
+  } else {
+    mat2_stride = mat2_strides.y;
+  }
 
   T sum = T(0.0);
   for (int i = 0; i < mat1_sizes.x; ++i) {
     sum += t_mat1[mat1_bufi] * t_mat2[mat2_bufi];
 
     mat1_bufi += mat1_strides.x;
-    mat2_bufi += mat2_strides.y;
+    mat2_bufi += mat2_stride;
   }
 
   const int out_bufi = tidx_to_bufi(out_bufix, out_strides);

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

@@ -77,6 +77,12 @@ void add_matmul_naive_buffer_node(
       graph.size_at<uint32_t>(-2, out),
       graph.size_at<uint32_t>(-3, out) * graph.size_at<uint32_t>(-4, out)};
 
+  int mat2_is_transposed_val = 0;
+  if (mat2_is_transposed != kDummyValueRef &&
+      graph.get_bool(mat2_is_transposed)) {
+    mat2_is_transposed_val = 1;
+  }
+
   graph.execute_nodes().emplace_back(new DispatchNode(
       graph,
       VK_KERNEL_FROM_STR(kernel_name),
@@ -96,7 +102,7 @@ void add_matmul_naive_buffer_node(
           graph.numel_ubo(out),
       },
       // Specialization Constants
-      {},
+      {mat2_is_transposed_val},
       // Resizing Logic
       resize_matmul_node,
       {mat2_is_transposed}));

backends/vulkan/test/op_tests/cases.py

@@ -126,8 +126,7 @@ def get_addmm_inputs():
 ]
 
 
-@register_test_suite("aten.linear.default")
-def get_linear_inputs():
+def get_linear_texture_inputs():
     MKN_list = common_MKN_list
 
     inputs_list = [((M, K), (N, K), None) for M, K, N in MKN_list]
@@ -142,9 +141,32 @@ def get_linear_inputs():
         "utils::kWidthPacked",
         "utils::kChannelsPacked",
     ]
+    test_suite.test_name_suffix = "texture"
+    return test_suite
+
+
+def get_linear_buffer_inputs():
+    MKN_list = common_MKN_list
+
+    inputs_list = [((M, K), (N, K), None) for M, K, N in MKN_list]
+    inputs_list += [((3, M, K), (N, K), None) for M, K, N in MKN_list]
+
+    test_suite = VkTestSuite(inputs_list)
+    test_suite.dtypes = ["at::kFloat"]
+    test_suite.layouts = [
+        "utils::kWidthPacked",
+        "utils::kChannelsPacked",
+    ]
+    test_suite.storage_types = ["utils::kBuffer"]
+    test_suite.test_name_suffix = "buffer"
     return test_suite
 
 
+@register_test_suite("aten.linear.default")
+def get_linear_test_suites():
+    return [get_linear_texture_inputs(), get_linear_buffer_inputs()]
+
+
 @register_test_suite("aten._weight_int8pack_mm.default")
 def get_weight_int8pack_mm_inputs():
     MKN_list = common_MKN_list