[webgpu] Optimize dp4 prefill shader for Qualcomm (microsoft#25578)

This change uses subgroupShuffle for sg_size=64 to perform the matmul.
It also uses a loop instead of loop unrolling to reduce the register
pressure.

Phi4 prefill for 1K tokens becomes 8.8s from 11.32s on Qualcomm Adreno
X1-85 GPU.

Author: qjia7

onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template

@@ -4,6 +4,7 @@
 #param block_size
 #param n_bits
 #param has_zero_points
+#param is_qualcomm
 
 #include "quantization/dp4a_matmul_common.wgsl.template"
 
@@ -138,18 +139,35 @@ $MAIN {
 
     // During the compute phase, we have the 64x64 tile split into
     // subtiles of 16x16. We have a grid of 4x4 subtiles.
-    let subtile_id = u32(local_idx / subtile_size);
-    let subtile_idx = u32(subtile_id / 4);
-    let subtile_idy = u32(subtile_id % 4);
-    let base_A = subtile_idx * 16;
-    let base_B = subtile_idy * 16;
+    var subtile_id = u32(local_idx / subtile_size);
+    var subtile_idx = u32(subtile_id / 4);
+    var subtile_idy = u32(subtile_id % 4);
+    var base_A = subtile_idx * 16;
+    var base_B = subtile_idy * 16;
     // For each subtile we have 16 threads assigned.
-    let a_idx = u32(local_idx % subtile_size);
+    var a_idx = u32(local_idx % subtile_size);
 
+#if is_qualcomm
+    // subtile_idx is always 0
+    // subtile_idy is one of {0,1,2,3}
+    // The subtile is now rectangular 64x16 for qualcomm case and we have 4 subtiles, this way we don't need to
+    // increase the number of lane_output each thread needs to track. That is if we want to use a subtile that is 64x64
+    // we would need var lane_outputs: array<output_element_t, 64>;
+    if (sg_size == 64) {
+        subtile_id = u32(local_idx / sg_size);
+        subtile_idx = u32(subtile_id / 4);
+        subtile_idy = u32(subtile_id % 4);
+        base_A = subtile_idx * sg_size;
+        base_B = subtile_idy * 16;
+        a_idx = sg_id;
+    }
+    var lane_outputs: array<output_element_t, 16>;
+#else
     var lane_output1: vec4<output_element_t>;
     var lane_output2: vec4<output_element_t>;
     var lane_output3: vec4<output_element_t>;
     var lane_output4: vec4<output_element_t>;
+#endif
     // K's vectorization is 16 items per index. See input_a/input_b.
     // tile_size_k_vec - is the k tile size in vectorized space (1/16). That is
     // k tile size is 32. In vectorized space that is 32/16 = 2.
@@ -173,6 +191,34 @@ $MAIN {
         var own_scale_a: output_element_t = scale_A[base_A + a_idx];
 
 #if has_zero_points && n_bits == 8
+ #if is_qualcomm
+        if (sg_size == 64)
+        {
+            var own_b0: vec4<u32>;
+            var own_b1: vec4<u32>;
+            var own_scale_b: output_element_t;
+            var zero: i32;
+            if (sg_id < 16)
+            {
+                own_b0 = tile_B[0][base_B + sg_id];
+                own_b1 = tile_B[1][base_B + sg_id];
+                own_scale_b = scale_B[base_B + sg_id];
+                zero = zeroes[base_B + sg_id];
+            }
+            // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul.
+            for (var i = 0u; i < 16u; i++)
+            {
+                lane_outputs[i] += SDP8AI(own_a0, subgroupShuffle(own_b0, i), own_a1, subgroupShuffle(own_b1, i), subgroupShuffle(own_scale_b, i) * own_scale_a, subgroupShuffle(zero, i));
+            }
+        }
+        else
+        {
+            for (var i = 0u; i < 16u; i++)
+            {
+                lane_outputs[i] += SDP8AI(own_a0, tile_B[0][base_B + i], own_a1, tile_B[1][base_B + i],  own_scale_a * scale_B[base_B + i], zeroes[base_B + i]);
+            }
+        }
+ #else
         if (sg_size == 16)
         {
             var own_b0: vec4<u32> = tile_B[0][base_B + sg_id];
@@ -225,7 +271,34 @@ $MAIN {
             lane_output4[2] += SDP8AI(own_a0, tile_B[0][base_B + 14], own_a1, tile_B[1][base_B + 14],  own_scale_a * scale_B[base_B + 14], zeroes[base_B + 14]);
             lane_output4[3] += SDP8AI(own_a0, tile_B[0][base_B + 15], own_a1, tile_B[1][base_B + 15],  own_scale_a * scale_B[base_B + 15], zeroes[base_B + 15]);
         }
+ #endif
 #else
+ #if is_qualcomm
+        if (sg_size == 64)
+        {
+            var own_b0: vec4<u32>;
+            var own_b1: vec4<u32>;
+            var own_scale_b: output_element_t;
+            if (sg_id < 16)
+            {
+                own_b0 = tile_B[0][base_B + sg_id];
+                own_b1 = tile_B[1][base_B + sg_id];
+                own_scale_b = scale_B[base_B + sg_id];
+            }
+            // Step 2: Access registers across the subgroup using subgroupShuffle and perform the matmul.
+            for (var i = 0u; i < 16u; i++)
+            {
+                lane_outputs[i] += SDP8AI(own_a0, subgroupShuffle(own_b0, i), own_a1, subgroupShuffle(own_b1, i), subgroupShuffle(own_scale_b, i) * own_scale_a);
+            }
+        }
+        else
+        {
+            for (var i = 0u; i < 16u; i++)
+            {
+                lane_outputs[i] += SDP8AI(own_a0, tile_B[0][base_B + i], own_a1, tile_B[1][base_B + i],  own_scale_a * scale_B[base_B + i]);
+            }
+        }
+ #else
         if (sg_size == 16)
         {
             var own_b0: vec4<u32> = tile_B[0][base_B + sg_id];
@@ -277,6 +350,7 @@ $MAIN {
             lane_output4[2] += SDP8AI(own_a0, tile_B[0][base_B + 14], own_a1, tile_B[1][base_B + 14],  own_scale_a * scale_B[base_B + 14]);
             lane_output4[3] += SDP8AI(own_a0, tile_B[0][base_B + 15], own_a1, tile_B[1][base_B + 15],  own_scale_a * scale_B[base_B + 15]);
         }
+ #endif
 #endif
         workgroupBarrier();
     }
@@ -287,9 +361,16 @@ $MAIN {
     // This creates a shader requirement that uniforms.N % 16 == 0
     if (a_global < uniforms.M && b_global < uniforms.N)
     {
+#if is_qualcomm
+        output[output_idx] = vec4<output_element_t>(lane_outputs[0], lane_outputs[1], lane_outputs[2], lane_outputs[3]);
+        output[output_idx+1] = vec4<output_element_t>(lane_outputs[4], lane_outputs[5], lane_outputs[6], lane_outputs[7]);
+        output[output_idx+2] = vec4<output_element_t>(lane_outputs[8], lane_outputs[9], lane_outputs[10], lane_outputs[11]);
+        output[output_idx+3] = vec4<output_element_t>(lane_outputs[12], lane_outputs[13], lane_outputs[14], lane_outputs[15]);
+#else
         output[output_idx] = lane_output1;
         output[output_idx+1] = lane_output2;
         output[output_idx+2] = lane_output3;
         output[output_idx+3] = lane_output4;
+#endif
     }
 }  // MAIN

onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc

@@ -28,6 +28,7 @@ Status DP4AMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const {
   return WGSL_TEMPLATE_APPLY(shader, "quantization/dp4a_matmul.wgsl.template",
                              WGSL_TEMPLATE_PARAMETER(block_size, block_size_),
                              WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_),
+                             WGSL_TEMPLATE_PARAMETER(is_qualcomm, is_qualcomm_),
                              WGSL_TEMPLATE_PARAMETER(n_bits, nbits_),
                              WGSL_TEMPLATE_PARAMETER(output_type_i32, true));
 }
@@ -118,7 +119,8 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor
   TensorShape reshaped_y_shape{1, M, N / kVec4Components};
   uint32_t num_M_tile = (M + kTileSize - 1) / kTileSize;
   uint32_t num_N_tile = (N + kTileSize - 1) / kTileSize;
-  DP4AMatMulNBitsProgram mul_program{block_size, nbits, has_zero_points};
+  bool is_qualcomm = context.AdapterInfo().vendor == std::string_view{"qualcomm"};
+  DP4AMatMulNBitsProgram mul_program{block_size, nbits, has_zero_points, is_qualcomm};
   mul_program.SetWorkgroupSize(256);
   mul_program.SetDispatchGroupSize(num_M_tile * num_N_tile);
   mul_program.AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast<int>(kVec4Components)},
@@ -133,7 +135,7 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor
                             {num_N_tile},
                             {zero_blocks_per_col}})
       .AddOutput({y, ProgramTensorMetadataDependency::TypeAndRank, reshaped_y_shape, static_cast<int>(kVec4Components)})
-      .CacheHint("Block" + std::to_string(block_size), nbits, has_zero_points);
+      .CacheHint("Block" + std::to_string(block_size), nbits, has_zero_points, is_qualcomm);
   if (has_zero_points) {
     mul_program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4});
   }

onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h

@@ -21,10 +21,11 @@ class DP4AMatMulQuantizeProgram final : public Program<DP4AMatMulQuantizeProgram
 
 class DP4AMatMulNBitsProgram final : public Program<DP4AMatMulNBitsProgram> {
  public:
-  DP4AMatMulNBitsProgram(uint32_t block_size, uint32_t nbits, bool has_zero_points) : Program{"DP4AMatMulNBits"},
-                                                                                      block_size_(block_size),
-                                                                                      nbits_(nbits),
-                                                                                      has_zero_points_(has_zero_points) {}
+  DP4AMatMulNBitsProgram(uint32_t block_size, uint32_t nbits, bool has_zero_points, bool is_qualcomm) : Program{"DP4AMatMulNBits"},
+                                                                                                        block_size_(block_size),
+                                                                                                        nbits_(nbits),
+                                                                                                        has_zero_points_(has_zero_points),
+                                                                                                        is_qualcomm_(is_qualcomm) {}
   Status GenerateShaderCode(ShaderHelper& sh) const override;
   WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES(
       {"M", ProgramUniformVariableDataType::Uint32},
@@ -39,6 +40,7 @@ class DP4AMatMulNBitsProgram final : public Program<DP4AMatMulNBitsProgram> {
   uint32_t block_size_;
   uint32_t nbits_;
   bool has_zero_points_;
+  bool is_qualcomm_;
 };
 
 class DP4AMatMulNBitsSmallMProgram final : public Program<DP4AMatMulNBitsSmallMProgram> {