diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index 22c9da4aab1..019c141b6fc 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -417,7 +417,7 @@ def register_softmax_cpp_ops(): # ============================================================================= -# MatMul.cpp +# gemm/matmul/MatMul.cpp # ============================================================================= @@ -437,7 +437,7 @@ def register_matmul_cpp_ops(): # ============================================================================= -# Linear.cpp +# gemm/linear/Linear.cpp # ============================================================================= @@ -457,7 +457,7 @@ def register_linear_cpp_ops(): # ============================================================================= -# QuantizedLinearQCSNW.cpp +# gemm/linear/qw/{q8csw,legacy}/*.cpp # ============================================================================= @@ -477,7 +477,7 @@ def register_quantizedlinearqcsnw_cpp_ops(): # ============================================================================= -# QuantizedLinear.cpp +# gemm/linear/qw/{q4gsw,dq8ca_q4gsw,legacy}/**/*.cpp # ============================================================================= @@ -1092,7 +1092,7 @@ def register_q8ta_conv2d_transposed_op(): # ============================================================================= -# Q8taLinear.cpp +# gemm/linear/q8ta/Q8taLinear.cpp # ============================================================================= diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_fp_input_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/conv2d_fp_input_tile_load.glslh index 4456043bb9f..f1cafc503e3 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/conv2d_fp_input_tile_load.glslh +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_fp_input_tile_load.glslh @@ -11,7 +11,7 @@ #extension GL_EXT_control_flow_attributes : require -#include "linear_fp_input_tile.glslh" +#include "gemm/tile_utils/linear_fp_input_tile.glslh" VEC4_T load_fp_input_texel(const Conv2dTensorIndex tidx) { #ifdef INPUT_BUFFER diff --git a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv1d/conv1d_pw.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv1d/conv1d_pw.glsl index 4ebca126e03..41d422515f2 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv1d/conv1d_pw.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv1d/conv1d_pw.glsl @@ -67,11 +67,11 @@ layout(push_constant) uniform restrict Block { layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; -#include "linear_fp_input_tile.glslh" -#include "linear_fp_weight_tile.glslh" -#include "linear_fp_output_tile.glslh" -#include "linear_fp_packed_weight_tile_load.glslh" -#include "linear_fp_output_tile_fp_compute.glslh" +#include "gemm/tile_utils/linear_fp_input_tile.glslh" +#include "gemm/tile_utils/linear_fp_weight_tile.glslh" +#include "gemm/tile_utils/linear_fp_output_tile.glslh" +#include "gemm/tile_utils/linear_fp_packed_weight_tile_load.glslh" +#include "gemm/tile_utils/linear_fp_output_tile_fp_compute.glslh" // Conv1d pointwise is matrix multiplication with swapped texture coordinates. // Linear: input ivec3(k4, m, b), output ivec3(n4, m, b) [width-packed] diff --git a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_pw_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_pw_tiled.glsl index bd12d149a2c..e4e8ae7a5ff 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_pw_tiled.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_pw_tiled.glsl @@ -47,9 +47,9 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; ${layout_declare_spec_const(C, "int", "stride_1_padding_0", "0")} ${layout_declare_spec_const(C, "int", "activation_type", "0")} -#include "linear_fp_input_tile.glslh" -#include "linear_fp_packed_weight_tile_load.glslh" -#include "linear_fp_output_tile_fp_compute.glslh" +#include "gemm/tile_utils/linear_fp_input_tile.glslh" +#include "gemm/tile_utils/linear_fp_packed_weight_tile_load.glslh" +#include "gemm/tile_utils/linear_fp_output_tile_fp_compute.glslh" // Texel coordinates for the TILE_M spatial positions this invocation covers. // diff --git a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/im2col/conv2d_gemm.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/im2col/conv2d_gemm.glsl index 6eebdb6b25f..d704b2363bf 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/im2col/conv2d_gemm.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/im2col/conv2d_gemm.glsl @@ -114,9 +114,9 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; ${layout_declare_spec_const(C, "int", "activation_type", "0")} -#include "linear_fp_input_tile.glslh" -#include "linear_fp_packed_weight_tile_load.glslh" -#include "linear_fp_output_tile_fp_compute.glslh" +#include "gemm/tile_utils/linear_fp_input_tile.glslh" +#include "gemm/tile_utils/linear_fp_packed_weight_tile_load.glslh" +#include "gemm/tile_utils/linear_fp_output_tile_fp_compute.glslh" /* * Load TILE_M rows × TILE_K4 K-tiles of the im2col'd input. diff --git a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_fp_im2col_block_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_fp_im2col_block_load.glslh index b94bdaf506f..d156adc6e6a 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_fp_im2col_block_load.glslh +++ b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_fp_im2col_block_load.glslh @@ -26,7 +26,7 @@ #include "indexing.glslh" #include "conv2d_common.glslh" #include "conv2d_fp_im2col_block.glslh" -#include "linear_fp_input_tile.glslh" +#include "gemm/tile_utils/linear_fp_input_tile.glslh" VEC4_T load_input_texel(const TensorIndex4D tidx) { // Assumes batch size is 1 and channels packing diff --git a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_fp_im2col_block_store.glslh b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_fp_im2col_block_store.glslh index 6c4dd7f0b52..8859b6c02e3 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_fp_im2col_block_store.glslh +++ b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_fp_im2col_block_store.glslh @@ -23,7 +23,7 @@ #include "indexing.glslh" #include "conv2d_common.glslh" #include "conv2d_fp_im2col_block.glslh" -#include "linear_fp_output_tile.glslh" +#include "gemm/tile_utils/linear_fp_output_tile.glslh" // TODO: implement buffer support void write_output_texel(const VEC4_T out_texel, const TensorIndex4D tidx) { diff --git a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_int8_input_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_int8_input_tile_load.glslh index 44aa09912ec..69407633e1f 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_int8_input_tile_load.glslh +++ b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_int8_input_tile_load.glslh @@ -11,7 +11,7 @@ #extension GL_EXT_control_flow_attributes : require -#include "linear_int8_input_tile.glslh" +#include "gemm/tile_utils/linear_int8_input_tile.glslh" struct Int8InputTileIndex { #ifdef PACKED_INT8_INPUT_BUFFER diff --git a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_int8_output_tile_store.glslh b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_int8_output_tile_store.glslh index 27244f67953..e6c48d99462 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_int8_output_tile_store.glslh +++ b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_int8_output_tile_store.glslh @@ -12,7 +12,7 @@ #extension GL_EXT_control_flow_attributes : require #include "conv2d_common.glslh" -#include "linear_int8_output_tile.glslh" +#include "gemm/tile_utils/linear_int8_output_tile.glslh" void store_packed_int8_output_tile( const Int8OutTile int8_tile, diff --git a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_pw_q8ta_q8csw_q8to_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_pw_q8ta_q8csw_q8to_tiled.glsl index f52241ee60f..bff8267c412 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_pw_q8ta_q8csw_q8to_tiled.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_pw_q8ta_q8csw_q8to_tiled.glsl @@ -74,12 +74,12 @@ ${layout_declare_spec_const(C, "int", "conv2d_params_groups", "1")} #include "conv2d_int8_input_tile_load.glslh" -#include "linear_int8_weight_tile_load.glslh" -#include "linear_fp_output_tile_int8_int8_compute.glslh" -#include "linear_int_weight_sums_load.glslh" -#include "linear_fp_weight_scales_load.glslh" -#include "linear_fp_bias_load.glslh" -#include "linear_int8_output_tile_compute.glslh" +#include "gemm/tile_utils/linear_int8_weight_tile_load.glslh" +#include "gemm/tile_utils/linear_fp_output_tile_int8_int8_compute.glslh" +#include "gemm/tile_utils/linear_int_weight_sums_load.glslh" +#include "gemm/tile_utils/linear_fp_weight_scales_load.glslh" +#include "gemm/tile_utils/linear_fp_bias_load.glslh" +#include "gemm/tile_utils/linear_int8_output_tile_compute.glslh" #include "conv2d_int8_output_tile_store.glslh" void main() { diff --git a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8_utils.glslh b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8_utils.glslh index 3fa03978fed..c30cf23aad8 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8_utils.glslh +++ b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8_utils.glslh @@ -12,7 +12,7 @@ #extension GL_EXT_control_flow_attributes : require #extension GL_EXT_integer_dot_product : require -#include "linear_int_accumulator.glslh" +#include "gemm/tile_utils/linear_int_accumulator.glslh" struct Int8InputWindow1D { int[MAX_WINDOW_WIDTH] data; diff --git a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8csw_linear_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8csw_linear_tiled.glsl index a93b34a3e6e..f0dcaa88273 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8csw_linear_tiled.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8csw_linear_tiled.glsl @@ -63,12 +63,12 @@ ${layout_declare_spec_const(C, "int", "conv2d_params_logical_K_per_group", "1")} ${layout_declare_spec_const(C, "int", "conv2d_params_groups", "1")} -#include "linear_fp_input_tile_load.glslh" -#include "linear_int8_weight_tile_load.glslh" -#include "linear_fp_weight_scales_load.glslh" -#include "linear_fp_bias_load.glslh" -#include "linear_fp_output_tile_fp_int8_compute.glslh" -#include "linear_fp_output_tile_fp_compute.glslh" +#include "gemm/tile_utils/linear_fp_input_tile_load.glslh" +#include "gemm/tile_utils/linear_int8_weight_tile_load.glslh" +#include "gemm/tile_utils/linear_fp_weight_scales_load.glslh" +#include "gemm/tile_utils/linear_fp_bias_load.glslh" +#include "gemm/tile_utils/linear_fp_output_tile_fp_int8_compute.glslh" +#include "gemm/tile_utils/linear_fp_output_tile_fp_compute.glslh" #include "conv2d_fp_im2col_block_store.glslh" void main() { diff --git a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_linear_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_linear_tiled.glsl index 3500f8392ad..dbbf3a08d2c 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_linear_tiled.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_linear_tiled.glsl @@ -71,12 +71,12 @@ ${layout_declare_spec_const(C, "int", "conv2d_params_logical_K_per_group", "1")} ${layout_declare_spec_const(C, "int", "conv2d_params_groups", "1")} -#include "linear_int8_input_tile_load.glslh" -#include "linear_int8_weight_tile_load.glslh" -#include "linear_int_weight_sums_load.glslh" -#include "linear_fp_weight_scales_load.glslh" -#include "linear_fp_output_tile_int8_int8_compute.glslh" -#include "linear_fp_bias_load.glslh" +#include "gemm/tile_utils/linear_int8_input_tile_load.glslh" +#include "gemm/tile_utils/linear_int8_weight_tile_load.glslh" +#include "gemm/tile_utils/linear_int_weight_sums_load.glslh" +#include "gemm/tile_utils/linear_fp_weight_scales_load.glslh" +#include "gemm/tile_utils/linear_fp_output_tile_int8_int8_compute.glslh" +#include "gemm/tile_utils/linear_fp_bias_load.glslh" #include "conv2d_fp_im2col_block_store.glslh" void main() { diff --git a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_q8to.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_q8to.glsl index 2c7f0aab981..4e3da29355b 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_q8to.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_q8to.glsl @@ -77,12 +77,12 @@ ${layout_declare_spec_const(C, "int", "conv2d_params_groups", "1")} #include "im2col_packed_int8_utils.glslh" #include "conv2d_int8_input_tile_load.glslh" -#include "linear_int8_weight_tile_load.glslh" -#include "linear_fp_output_tile_int8_int8_compute.glslh" -#include "linear_int_weight_sums_load.glslh" -#include "linear_fp_weight_scales_load.glslh" -#include "linear_fp_bias_load.glslh" -#include "linear_int8_output_tile_compute.glslh" +#include "gemm/tile_utils/linear_int8_weight_tile_load.glslh" +#include "gemm/tile_utils/linear_fp_output_tile_int8_int8_compute.glslh" +#include "gemm/tile_utils/linear_int_weight_sums_load.glslh" +#include "gemm/tile_utils/linear_fp_weight_scales_load.glslh" +#include "gemm/tile_utils/linear_fp_bias_load.glslh" +#include "gemm/tile_utils/linear_int8_output_tile_compute.glslh" #include "conv2d_int8_output_tile_store.glslh" #include "conv2d_q8_utils.glslh" diff --git a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_q8to_linear_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_q8to_linear_tiled.glsl index f136c6ce99a..b46a0630862 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_q8to_linear_tiled.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_q8to_linear_tiled.glsl @@ -74,12 +74,12 @@ ${layout_declare_spec_const(C, "int", "conv2d_params_groups", "1")} #include "conv2d_int8_input_tile_load.glslh" -#include "linear_int8_weight_tile_load.glslh" -#include "linear_fp_output_tile_int8_int8_compute.glslh" -#include "linear_int_weight_sums_load.glslh" -#include "linear_fp_weight_scales_load.glslh" -#include "linear_fp_bias_load.glslh" -#include "linear_int8_output_tile_compute.glslh" +#include "gemm/tile_utils/linear_int8_weight_tile_load.glslh" +#include "gemm/tile_utils/linear_fp_output_tile_int8_int8_compute.glslh" +#include "gemm/tile_utils/linear_int_weight_sums_load.glslh" +#include "gemm/tile_utils/linear_fp_weight_scales_load.glslh" +#include "gemm/tile_utils/linear_fp_bias_load.glslh" +#include "gemm/tile_utils/linear_int8_output_tile_compute.glslh" #include "conv2d_int8_output_tile_store.glslh" void main() { diff --git a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/quantize_and_pack_im2col.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/quantize_and_pack_im2col.glsl index a008866a779..1011b6d41ac 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/quantize_and_pack_im2col.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/quantize_and_pack_im2col.glsl @@ -67,7 +67,7 @@ layout(push_constant) uniform restrict Block { layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; #include "conv2d_fp_im2col_block_load.glslh" -#include "linear_int8_input_block.glslh" +#include "gemm/tile_utils/linear_int8_input_block.glslh" void main() { // The quantized and packed im2col matrix can be conceptualized as a 2D matrix diff --git a/backends/vulkan/runtime/graph/ops/glsl/coopmat_mm.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/coopmat_mm.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/coopmat_mm.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/coopmat_mm.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/coopmat_mm.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/coopmat_mm.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/coopmat_mm.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/coopmat_mm.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dW.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/backward/linear_dW.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_dW.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/backward/linear_dW.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dW.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/backward/linear_dW.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_dW.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/backward/linear_dW.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_scalar.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/linear_scalar.glsl similarity index 90% rename from backends/vulkan/runtime/graph/ops/glsl/linear_scalar.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/linear_scalar.glsl index 9f5c2ee6a50..4ac5b9c88ee 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_scalar.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/linear_scalar.glsl @@ -61,11 +61,11 @@ $else: layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; -#include "matmul_fp_mat1_tile_load.glslh" -#include "linear_fp_packed_weight_tile_load.glslh" -#include "linear_fp_output_tile_fp_compute.glslh" -#include "matmul_fp_out_tile_store.glslh" -#include "matmul_fp_bias_apply.glslh" +#include "tile_utils/matmul_fp_mat1_tile_load.glslh" +#include "tile_utils/linear_fp_packed_weight_tile_load.glslh" +#include "tile_utils/linear_fp_output_tile_fp_compute.glslh" +#include "tile_utils/matmul_fp_out_tile_store.glslh" +#include "tile_utils/matmul_fp_bias_apply.glslh" void main() { const int tile_idx_n = int(gl_GlobalInvocationID.x); diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_scalar.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/linear_scalar.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_scalar.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/linear_scalar.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_vec.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/linear_vec.glsl similarity index 90% rename from backends/vulkan/runtime/graph/ops/glsl/linear_vec.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/linear_vec.glsl index adf5272673e..f17520777f6 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_vec.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/linear_vec.glsl @@ -62,11 +62,11 @@ $else: layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; -#include "matmul_fp_mat1_tile_load.glslh" -#include "linear_fp_packed_weight_tile_load.glslh" -#include "linear_fp_output_tile_fp_compute.glslh" -#include "matmul_fp_out_tile_store.glslh" -#include "matmul_fp_bias_apply.glslh" +#include "tile_utils/matmul_fp_mat1_tile_load.glslh" +#include "tile_utils/linear_fp_packed_weight_tile_load.glslh" +#include "tile_utils/linear_fp_output_tile_fp_compute.glslh" +#include "tile_utils/matmul_fp_out_tile_store.glslh" +#include "tile_utils/matmul_fp_bias_apply.glslh" void main() { const int tile_idx_n = int(gl_GlobalInvocationID.x); diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_vec.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/linear_vec.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_vec.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/linear_vec.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_fp_linear_weight.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/pack_weights/pack_fp_linear_weight.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/pack_fp_linear_weight.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/pack_weights/pack_fp_linear_weight.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_fp_linear_weight.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/pack_weights/pack_fp_linear_weight.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/pack_fp_linear_weight.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/pack_weights/pack_fp_linear_weight.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_linear.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/q8ta/q8ta_linear.glsl similarity index 92% rename from backends/vulkan/runtime/graph/ops/glsl/q8ta_linear.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/q8ta/q8ta_linear.glsl index 98b0895d759..b09b593990c 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q8ta_linear.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/q8ta/q8ta_linear.glsl @@ -57,12 +57,12 @@ layout(push_constant) uniform restrict Block { layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; #include "common.glslh" -#include "linear_int8_input_tile_load.glslh" -#include "linear_int8_weight_tile_load.glslh" -#include "linear_fp_output_tile_int8_int8_compute.glslh" -#include "linear_fp_weight_scales_load.glslh" -#include "linear_int_weight_sums_load.glslh" -#include "linear_fp_bias_load.glslh" +#include "tile_utils/linear_int8_input_tile_load.glslh" +#include "tile_utils/linear_int8_weight_tile_load.glslh" +#include "tile_utils/linear_fp_output_tile_int8_int8_compute.glslh" +#include "tile_utils/linear_fp_weight_scales_load.glslh" +#include "tile_utils/linear_int_weight_sums_load.glslh" +#include "tile_utils/linear_fp_bias_load.glslh" void main() { const int out_tile_x = int(gl_GlobalInvocationID.x); diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_linear.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/q8ta/q8ta_linear.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q8ta_linear.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/q8ta/q8ta_linear.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_linear_gemv.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/q8ta/q8ta_linear_gemv.glsl similarity index 94% rename from backends/vulkan/runtime/graph/ops/glsl/q8ta_linear_gemv.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/q8ta/q8ta_linear_gemv.glsl index 9153cb0e294..7ae85518563 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q8ta_linear_gemv.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/q8ta/q8ta_linear_gemv.glsl @@ -62,11 +62,11 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; #include "dispatch.glslh" #include "common.glslh" -#include "linear_int8_weight_tile_load.glslh" -#include "linear_fp_output_tile_int8_int8_compute.glslh" -#include "linear_fp_weight_scales_load.glslh" -#include "linear_int_weight_sums_load.glslh" -#include "linear_fp_bias_load.glslh" +#include "tile_utils/linear_int8_weight_tile_load.glslh" +#include "tile_utils/linear_fp_output_tile_int8_int8_compute.glslh" +#include "tile_utils/linear_fp_weight_scales_load.glslh" +#include "tile_utils/linear_int_weight_sums_load.glslh" +#include "tile_utils/linear_fp_bias_load.glslh" // Array-of-arrays shared memory layout: partial_accums[lid][tile_n4]. // Each element is exactly 16 bytes (ivec4). This avoids the Samsung S25 diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_linear_gemv.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/q8ta/q8ta_linear_gemv.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q8ta_linear_gemv.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/q8ta/q8ta_linear_gemv.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_tiled.glsl similarity index 85% rename from backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_tiled.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_tiled.glsl index a68a56c3713..5acab8e9a12 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_tiled.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_tiled.glsl @@ -60,19 +60,19 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; ${layout_declare_spec_const(C, "int", "apply_bias", "0")} ${layout_declare_spec_const(C, "int", "K4_per_group", "0")} -#include "linear_fp_input_tile_load.glslh" -#include "linear_int8_input_tile_load.glslh" -#include "linear_int8_input_scales_zps_load.glslh" -#include "linear_int4_weight_tile_load.glslh" -#include "linear_int_weight_sums_load.glslh" -#include "linear_fp_weight_scales_load.glslh" -#include "linear_int8_input_sums_load.glslh" -#include "linear_fp_output_tile_int8_int8_compute.glslh" -#include "linear_fp_output_tile_int8_int4_compute.glslh" -#include "linear_fp_output_tile_fp_int4_compute.glslh" -#include "linear_fp_output_tile_fp_compute.glslh" -#include "linear_fp_output_tile_store.glslh" -#include "linear_fp_bias_load.glslh" +#include "tile_utils/linear_fp_input_tile_load.glslh" +#include "tile_utils/linear_int8_input_tile_load.glslh" +#include "tile_utils/linear_int8_input_scales_zps_load.glslh" +#include "tile_utils/linear_int4_weight_tile_load.glslh" +#include "tile_utils/linear_int_weight_sums_load.glslh" +#include "tile_utils/linear_fp_weight_scales_load.glslh" +#include "tile_utils/linear_int8_input_sums_load.glslh" +#include "tile_utils/linear_fp_output_tile_int8_int8_compute.glslh" +#include "tile_utils/linear_fp_output_tile_int8_int4_compute.glslh" +#include "tile_utils/linear_fp_output_tile_fp_int4_compute.glslh" +#include "tile_utils/linear_fp_output_tile_fp_compute.glslh" +#include "tile_utils/linear_fp_output_tile_store.glslh" +#include "tile_utils/linear_fp_bias_load.glslh" void main() { const int out_tile_x = int(gl_GlobalInvocationID.x); diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_tiled.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_tiled.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_tiled.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_qw_coopmat.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_qw_coopmat.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_qw_coopmat.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_qw_coopmat.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q8ta_q8csw_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/legacy/linear_q8ta_q8csw_tiled.glsl similarity index 89% rename from backends/vulkan/runtime/graph/ops/glsl/linear_q8ta_q8csw_tiled.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/legacy/linear_q8ta_q8csw_tiled.glsl index e3f49fac475..20003489dca 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_q8ta_q8csw_tiled.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/legacy/linear_q8ta_q8csw_tiled.glsl @@ -50,13 +50,13 @@ layout(push_constant) uniform restrict Block { layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; -#include "linear_int8_input_tile_load.glslh" -#include "linear_int8_weight_tile_load.glslh" -#include "linear_fp_output_tile_int8_int8_compute.glslh" -#include "linear_fp_output_tile_store.glslh" -#include "linear_fp_weight_scales_load.glslh" -#include "linear_int_weight_sums_load.glslh" -#include "linear_fp_bias_load.glslh" +#include "tile_utils/linear_int8_input_tile_load.glslh" +#include "tile_utils/linear_int8_weight_tile_load.glslh" +#include "tile_utils/linear_fp_output_tile_int8_int8_compute.glslh" +#include "tile_utils/linear_fp_output_tile_store.glslh" +#include "tile_utils/linear_fp_weight_scales_load.glslh" +#include "tile_utils/linear_int_weight_sums_load.glslh" +#include "tile_utils/linear_fp_bias_load.glslh" void main() { // Each thread writes out a 4 wide x 4 high tile of output values diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q8ta_q8csw_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/legacy/linear_q8ta_q8csw_tiled.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_q8ta_q8csw_tiled.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/legacy/linear_q8ta_q8csw_tiled.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_qcsnw.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/legacy/linear_qcsnw.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_qcsnw.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/legacy/linear_qcsnw.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_qcsnw.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/legacy/linear_qcsnw.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_qcsnw.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/legacy/linear_qcsnw.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_qcsnw_coop.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/legacy/linear_qcsnw_coop.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_qcsnw_coop.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/legacy/linear_qcsnw_coop.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_qcsnw_coop.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/legacy/linear_qcsnw_coop.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_qcsnw_coop.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/legacy/linear_qcsnw_coop.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_qcsnw_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/legacy/linear_qcsnw_tiled.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_qcsnw_tiled.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/legacy/linear_qcsnw_tiled.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_qcsnw_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/legacy/linear_qcsnw_tiled.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_qcsnw_tiled.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/legacy/linear_qcsnw_tiled.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/pack_weights/pack_q8_linear_weight.glsl similarity index 97% rename from backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/pack_weights/pack_q8_linear_weight.glsl index 7b17dbe4990..765c8a164e1 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/pack_weights/pack_q8_linear_weight.glsl @@ -26,7 +26,7 @@ layout(push_constant) uniform restrict Block { layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; #include "common.glslh" -#include "linear_int8_weight_block.glslh" +#include "tile_utils/linear_int8_weight_block.glslh" void main() { // The size of the source weight tensor is [W=K, H=N]. Each shader invocation diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/pack_weights/pack_q8_linear_weight.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/pack_weights/pack_q8_linear_weight.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/q4gsw_backward.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/backward/q4gsw_backward.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q4gsw_backward.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/backward/q4gsw_backward.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/q4gsw_backward.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/backward/q4gsw_backward.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q4gsw_backward.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/backward/q4gsw_backward.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/q4gsw_requant.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/backward/q4gsw_requant.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q4gsw_requant.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/backward/q4gsw_requant.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/q4gsw_requant.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/backward/q4gsw_requant.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q4gsw_requant.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/backward/q4gsw_requant.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coop.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_coop.glsl similarity index 92% rename from backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coop.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_coop.glsl index 0f88dc59702..8a2c9f8bdcf 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coop.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_coop.glsl @@ -63,13 +63,13 @@ ${layout_declare_spec_const(C, "int", "apply_bias", "0")} ${layout_declare_spec_const(C, "int", "K4_per_group", "0")} #include "common.glslh" -#include "linear_fp_input_tile_load.glslh" -#include "linear_int4_weight_tile_load.glslh" -#include "linear_fp_weight_scales_load.glslh" -#include "linear_fp_output_tile_fp_int4_compute.glslh" -#include "linear_fp_output_tile_fp_compute.glslh" -#include "linear_fp_output_tile_store.glslh" -#include "linear_fp_bias_load.glslh" +#include "tile_utils/linear_fp_input_tile_load.glslh" +#include "tile_utils/linear_int4_weight_tile_load.glslh" +#include "tile_utils/linear_fp_weight_scales_load.glslh" +#include "tile_utils/linear_fp_output_tile_fp_int4_compute.glslh" +#include "tile_utils/linear_fp_output_tile_fp_compute.glslh" +#include "tile_utils/linear_fp_output_tile_store.glslh" +#include "tile_utils/linear_fp_bias_load.glslh" shared FPOutTile partial_sums[WGS]; diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coop.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_coop.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coop.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_coop.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_tiled.glsl similarity index 89% rename from backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_tiled.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_tiled.glsl index f79512ac7ee..9bfa875d58d 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_tiled.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_tiled.glsl @@ -49,13 +49,13 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; ${layout_declare_spec_const(C, "int", "apply_bias", "0")} ${layout_declare_spec_const(C, "int", "K4_per_group", "0")} -#include "linear_fp_input_tile_load.glslh" -#include "linear_int4_weight_tile_load.glslh" -#include "linear_fp_weight_scales_load.glslh" -#include "linear_fp_bias_load.glslh" -#include "linear_fp_output_tile_fp_int4_compute.glslh" -#include "linear_fp_output_tile_fp_compute.glslh" -#include "linear_fp_output_tile_store.glslh" +#include "tile_utils/linear_fp_input_tile_load.glslh" +#include "tile_utils/linear_int4_weight_tile_load.glslh" +#include "tile_utils/linear_fp_weight_scales_load.glslh" +#include "tile_utils/linear_fp_bias_load.glslh" +#include "tile_utils/linear_fp_output_tile_fp_int4_compute.glslh" +#include "tile_utils/linear_fp_output_tile_fp_compute.glslh" +#include "tile_utils/linear_fp_output_tile_store.glslh" void main() { const int out_tile_x = int(gl_GlobalInvocationID.x); diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_tiled.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_tiled.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_q4gsw_tiled.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_qw_coopmat.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_qw_coopmat.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_qw_coopmat.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_qw_coopmat.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_qw_coopmat.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_qw_coopmat.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_qw_coopmat.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/linear_qw_coopmat.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_q4_linear_weight.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/pack_weights/pack_q4_linear_weight.glsl similarity index 97% rename from backends/vulkan/runtime/graph/ops/glsl/pack_q4_linear_weight.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/pack_weights/pack_q4_linear_weight.glsl index b9f5c994910..2b2cf023b5c 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/pack_q4_linear_weight.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/pack_weights/pack_q4_linear_weight.glsl @@ -25,7 +25,7 @@ layout(push_constant) uniform restrict Block { layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; #include "common.glslh" -#include "linear_int4_weight_block.glslh" +#include "tile_utils/linear_int4_weight_block.glslh" void main() { const int k8 = int(gl_GlobalInvocationID.x); diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_q4_linear_weight.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/pack_weights/pack_q4_linear_weight.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/pack_q4_linear_weight.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/pack_weights/pack_q4_linear_weight.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_q4_linear_weight__w_4x8.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/pack_weights/pack_q4_linear_weight__w_4x8.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/pack_q4_linear_weight__w_4x8.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/pack_weights/pack_q4_linear_weight__w_4x8.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_q4_linear_weight__w_4x8.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/pack_weights/pack_q4_linear_weight__w_4x8.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/pack_q4_linear_weight__w_4x8.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/pack_weights/pack_q4_linear_weight__w_4x8.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/q4gsw_linear_gemm__tin__w_4x8.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/q4gsw_linear_gemm__tin__w_4x8.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q4gsw_linear_gemm__tin__w_4x8.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/q4gsw_linear_gemm__tin__w_4x8.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/q4gsw_linear_gemm__tin__w_4x8.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/q4gsw_linear_gemm__tin__w_4x8.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q4gsw_linear_gemm__tin__w_4x8.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/q4gsw_linear_gemm__tin__w_4x8.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/q4gsw_linear_gemm__w_4x8.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/q4gsw_linear_gemm__w_4x8.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q4gsw_linear_gemm__w_4x8.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/q4gsw_linear_gemm__w_4x8.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/q4gsw_linear_gemm__w_4x8.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/q4gsw_linear_gemm__w_4x8.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q4gsw_linear_gemm__w_4x8.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/q4gsw_linear_gemm__w_4x8.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/q4gsw_linear_gemv_coop__w_4x8.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/q4gsw_linear_gemv_coop__w_4x8.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q4gsw_linear_gemv_coop__w_4x8.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/q4gsw_linear_gemv_coop__w_4x8.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/q4gsw_linear_gemv_coop__w_4x8.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/q4gsw_linear_gemv_coop__w_4x8.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q4gsw_linear_gemv_coop__w_4x8.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q4gsw/q4gsw_linear_gemv_coop__w_4x8.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q8csw_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q8csw/linear_q8csw_tiled.glsl similarity index 87% rename from backends/vulkan/runtime/graph/ops/glsl/linear_q8csw_tiled.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q8csw/linear_q8csw_tiled.glsl index c8dc6bfeec4..51a73d1126f 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_q8csw_tiled.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q8csw/linear_q8csw_tiled.glsl @@ -44,14 +44,14 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; ${layout_declare_spec_const(C, "uint", "apply_bias", "0")} -#include "linear_fp_input_tile_load.glslh" -#include "linear_int8_weight_tile_load.glslh" -#include "linear_fp_weight_tile.glslh" -#include "linear_fp_output_tile_fp_compute.glslh" -#include "linear_fp_output_tile_fp_int8_compute.glslh" -#include "linear_fp_output_tile_store.glslh" -#include "linear_fp_weight_scales_load.glslh" -#include "linear_fp_bias_load.glslh" +#include "tile_utils/linear_fp_input_tile_load.glslh" +#include "tile_utils/linear_int8_weight_tile_load.glslh" +#include "tile_utils/linear_fp_weight_tile.glslh" +#include "tile_utils/linear_fp_output_tile_fp_compute.glslh" +#include "tile_utils/linear_fp_output_tile_fp_int8_compute.glslh" +#include "tile_utils/linear_fp_output_tile_store.glslh" +#include "tile_utils/linear_fp_weight_scales_load.glslh" +#include "tile_utils/linear_fp_bias_load.glslh" void main() { // Each thread writes out a 4 wide x 4 high tile of output values diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q8csw_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q8csw/linear_q8csw_tiled.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_q8csw_tiled.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/q8csw/linear_q8csw_tiled.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/matmul_scalar.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/matmul/matmul_scalar.glsl similarity index 90% rename from backends/vulkan/runtime/graph/ops/glsl/matmul_scalar.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/matmul/matmul_scalar.glsl index a3f68dda839..bfffe5fa38d 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/matmul_scalar.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/matmul/matmul_scalar.glsl @@ -54,11 +54,11 @@ $if HAS_BIAS: layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; -#include "matmul_fp_mat1_tile_load.glslh" -#include "matmul_fp_mat2_tile_load.glslh" -#include "linear_fp_output_tile_fp_compute.glslh" -#include "matmul_fp_out_tile_store.glslh" -#include "matmul_fp_bias_apply.glslh" +#include "tile_utils/matmul_fp_mat1_tile_load.glslh" +#include "tile_utils/matmul_fp_mat2_tile_load.glslh" +#include "tile_utils/linear_fp_output_tile_fp_compute.glslh" +#include "tile_utils/matmul_fp_out_tile_store.glslh" +#include "tile_utils/matmul_fp_bias_apply.glslh" void main() { const int tile_idx_n = int(gl_GlobalInvocationID.x); diff --git a/backends/vulkan/runtime/graph/ops/glsl/matmul_scalar.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/matmul/matmul_scalar.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/matmul_scalar.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/matmul/matmul_scalar.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/matmul_vec.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/matmul/matmul_vec.glsl similarity index 90% rename from backends/vulkan/runtime/graph/ops/glsl/matmul_vec.glsl rename to backends/vulkan/runtime/graph/ops/glsl/gemm/matmul/matmul_vec.glsl index 9d4f6f04aba..c14861170c2 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/matmul_vec.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/matmul/matmul_vec.glsl @@ -55,11 +55,11 @@ $if HAS_BIAS: layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; -#include "matmul_fp_mat1_tile_load.glslh" -#include "matmul_fp_mat2_tile_load.glslh" -#include "linear_fp_output_tile_fp_compute.glslh" -#include "matmul_fp_out_tile_store.glslh" -#include "matmul_fp_bias_apply.glslh" +#include "tile_utils/matmul_fp_mat1_tile_load.glslh" +#include "tile_utils/matmul_fp_mat2_tile_load.glslh" +#include "tile_utils/linear_fp_output_tile_fp_compute.glslh" +#include "tile_utils/matmul_fp_out_tile_store.glslh" +#include "tile_utils/matmul_fp_bias_apply.glslh" void main() { const int tile_idx_n = int(gl_GlobalInvocationID.x); diff --git a/backends/vulkan/runtime/graph/ops/glsl/matmul_vec.yaml b/backends/vulkan/runtime/graph/ops/glsl/gemm/matmul/matmul_vec.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/matmul_vec.yaml rename to backends/vulkan/runtime/graph/ops/glsl/gemm/matmul/matmul_vec.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_common.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_common.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_common.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_common.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_bias_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_bias_load.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_fp_bias_load.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_bias_load.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_input_tile.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_input_tile.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_fp_input_tile.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_input_tile.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_input_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_input_tile_load.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_fp_input_tile_load.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_input_tile_load.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_output_tile.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_output_tile.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_fp_compute.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_output_tile_fp_compute.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_fp_compute.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_output_tile_fp_compute.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_fp_int4_compute.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_output_tile_fp_int4_compute.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_fp_int4_compute.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_output_tile_fp_int4_compute.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_fp_int8_compute.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_output_tile_fp_int8_compute.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_fp_int8_compute.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_output_tile_fp_int8_compute.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_int8_int4_compute.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_output_tile_int8_int4_compute.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_int8_int4_compute.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_output_tile_int8_int4_compute.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_int8_int8_compute.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_output_tile_int8_int8_compute.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_int8_int8_compute.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_output_tile_int8_int8_compute.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_store.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_output_tile_store.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_store.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_output_tile_store.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_packed_weight_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_packed_weight_tile_load.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_fp_packed_weight_tile_load.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_packed_weight_tile_load.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_per_out_channel_params.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_per_out_channel_params.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_fp_per_out_channel_params.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_per_out_channel_params.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_weight_scales_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_weight_scales_load.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_fp_weight_scales_load.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_weight_scales_load.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_weight_tile.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_weight_tile.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_fp_weight_tile.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_fp_weight_tile.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int4_weight_block.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int4_weight_block.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_int4_weight_block.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int4_weight_block.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int4_weight_tile.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int4_weight_tile.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_int4_weight_tile.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int4_weight_tile.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int4_weight_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int4_weight_tile_load.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_int4_weight_tile_load.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int4_weight_tile_load.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_block.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_block.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_block.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_block.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_scales_zps.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_scales_zps.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_scales_zps.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_scales_zps.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_scales_zps_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_scales_zps_load.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_scales_zps_load.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_scales_zps_load.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_sums_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_sums_load.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_sums_load.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_sums_load.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_sums_store.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_sums_store.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_sums_store.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_sums_store.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_tile.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_tile.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_tile.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_tile.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_tile_load.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_tile_load.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_input_tile_load.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_output_tile.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_output_tile.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_int8_output_tile.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_output_tile.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_output_tile_compute.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_output_tile_compute.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_int8_output_tile_compute.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_output_tile_compute.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_weight_block.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_weight_block.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_int8_weight_block.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_weight_block.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_weight_tile.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_weight_tile.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_int8_weight_tile.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_weight_tile.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_weight_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_weight_tile_load.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_int8_weight_tile_load.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int8_weight_tile_load.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int_accumulator.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int_accumulator.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_int_accumulator.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int_accumulator.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int_per_in_channel_params.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int_per_in_channel_params.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_int_per_in_channel_params.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int_per_in_channel_params.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int_per_out_channel_params.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int_per_out_channel_params.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_int_per_out_channel_params.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int_per_out_channel_params.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int_weight_sums_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int_weight_sums_load.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/linear_int_weight_sums_load.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/linear_int_weight_sums_load.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/matmul_fp_bias_apply.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/matmul_fp_bias_apply.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/matmul_fp_bias_apply.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/matmul_fp_bias_apply.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/matmul_fp_mat1_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/matmul_fp_mat1_tile_load.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/matmul_fp_mat1_tile_load.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/matmul_fp_mat1_tile_load.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/matmul_fp_mat2_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/matmul_fp_mat2_tile_load.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/matmul_fp_mat2_tile_load.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/matmul_fp_mat2_tile_load.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/matmul_fp_out_tile_store.glslh b/backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/matmul_fp_out_tile_store.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/matmul_fp_out_tile_store.glslh rename to backends/vulkan/runtime/graph/ops/glsl/gemm/tile_utils/matmul_fp_out_tile_store.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_int4_linear_weight_transposed_interleaved.glsl b/backends/vulkan/runtime/graph/ops/glsl/pack_weights/pack_int4_linear_weight_transposed_interleaved.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/pack_int4_linear_weight_transposed_interleaved.glsl rename to backends/vulkan/runtime/graph/ops/glsl/pack_weights/pack_int4_linear_weight_transposed_interleaved.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_int4_linear_weight_transposed_interleaved.yaml b/backends/vulkan/runtime/graph/ops/glsl/pack_weights/pack_int4_linear_weight_transposed_interleaved.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/pack_int4_linear_weight_transposed_interleaved.yaml rename to backends/vulkan/runtime/graph/ops/glsl/pack_weights/pack_int4_linear_weight_transposed_interleaved.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4h4w.glsl b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4h4w.glsl index 9d65d7d534f..41414a22dfd 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4h4w.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4h4w.glsl @@ -38,8 +38,8 @@ layout(push_constant) uniform restrict Block { layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; -#include "linear_int8_input_block.glslh" -#include "linear_fp_input_tile_load.glslh" +#include "gemm/tile_utils/linear_int8_input_block.glslh" +#include "gemm/tile_utils/linear_fp_input_tile_load.glslh" void main() { // Each input block contains 4x4 int8 quantized values, which are packed into diff --git a/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4h4w_with_group_sums.glsl b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4h4w_with_group_sums.glsl index d5f14f70b62..407c49fddf4 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4h4w_with_group_sums.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4h4w_with_group_sums.glsl @@ -48,9 +48,9 @@ shared ivec4 shared_sums[NUM_GROUPS_PER_WG][NUM_WORKERS_PER_GROUP]; #define TILE_M 4 -#include "linear_int8_input_block.glslh" -#include "linear_int8_input_scales_zps_load.glslh" -#include "linear_fp_input_tile_load.glslh" +#include "gemm/tile_utils/linear_int8_input_block.glslh" +#include "gemm/tile_utils/linear_int8_input_scales_zps_load.glslh" +#include "gemm/tile_utils/linear_fp_input_tile_load.glslh" void main() { const int group_idx = int(gl_GlobalInvocationID.x); diff --git a/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4w4c.glsl b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4w4c.glsl index 6f58021be7f..3b5be301425 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4w4c.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4w4c.glsl @@ -43,7 +43,7 @@ layout(push_constant) uniform restrict Block { layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; #include "conv2d_fp_input_tile_load.glslh" -#include "linear_int8_input_block.glslh" +#include "gemm/tile_utils/linear_int8_input_block.glslh" void store_packed_int8_block( const Conv2dBlockIndex block_idx, diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.glsl index b7f14f435fa..ea1a80d6388 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.glsl @@ -51,7 +51,7 @@ ${layout_declare_spec_const(C, "float", "inv_scale", "1.0")} #include "sdpa_fp_q_projected_tile_load.glslh" #include "sdpa_fp_k_cache_tile_load.glslh" -#include "linear_fp_output_tile_fp_compute.glslh" +#include "gemm/tile_utils/linear_fp_output_tile_fp_compute.glslh" #include "sdpa_fp_attn_weight_tile_store.glslh" shared FPOutTile partial_sums[NUM_WORKERS_PER_OUT]; diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_tiled.glsl index 662d5edde68..a0feba5bb15 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_tiled.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_tiled.glsl @@ -70,7 +70,7 @@ ${layout_declare_spec_const(C, "float", "inv_scale", "1.0")} #include "sdpa_fp_q_projected_tile_load.glslh" #include "sdpa_fp_k_cache_tile_load.glslh" -#include "linear_fp_output_tile_fp_compute.glslh" +#include "gemm/tile_utils/linear_fp_output_tile_fp_compute.glslh" #include "sdpa_fp_attn_weight_tile_store.glslh" /* diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.glsl index 7f873ce0718..d1aab16290f 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.glsl @@ -64,7 +64,7 @@ $if GQA: #include "sdpa_fp_attn_weight_tile_load.glslh" #include "sdpa_fp_v_cache_tile_load.glslh" -#include "linear_fp_output_tile_fp_compute.glslh" +#include "gemm/tile_utils/linear_fp_output_tile_fp_compute.glslh" #include "sdpa_fp_out_tile_store.glslh" #ifdef GQA diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_tiled.glsl index 9f8f2dbc231..d191401296c 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_tiled.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_tiled.glsl @@ -56,7 +56,7 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; #include "sdpa_fp_attn_weight_tile_load.glslh" #include "sdpa_fp_v_cache_tile_load.glslh" -#include "linear_fp_output_tile_fp_compute.glslh" +#include "gemm/tile_utils/linear_fp_output_tile_fp_compute.glslh" #include "sdpa_fp_out_tile_store.glslh" /* diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_attn_weight_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_attn_weight_tile_load.glslh index 829f03beb60..af639e77371 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_attn_weight_tile_load.glslh +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_attn_weight_tile_load.glslh @@ -27,7 +27,7 @@ #extension GL_EXT_control_flow_attributes : require -#include "linear_fp_input_tile.glslh" +#include "gemm/tile_utils/linear_fp_input_tile.glslh" LINEAR_FP_INPUT_TILE_VEC4_T load_attn_weight_c4( const int c4, diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_attn_weight_tile_store.glslh b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_attn_weight_tile_store.glslh index 83d81ddf812..78a7df1d4be 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_attn_weight_tile_store.glslh +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_attn_weight_tile_store.glslh @@ -32,7 +32,7 @@ #extension GL_EXT_control_flow_attributes : require -#include "linear_fp_output_tile.glslh" +#include "gemm/tile_utils/linear_fp_output_tile.glslh" // ============================================================ // Shared store helpers (buffer or texture; fp32 or input-dtype) diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_k_cache_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_k_cache_tile_load.glslh index 65d08755528..555561581c7 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_k_cache_tile_load.glslh +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_k_cache_tile_load.glslh @@ -31,7 +31,7 @@ #extension GL_EXT_control_flow_attributes : require -#include "linear_fp_weight_tile.glslh" +#include "gemm/tile_utils/linear_fp_weight_tile.glslh" // Determine whether buffer mode is active. Both K_CACHE_BUFFER (LLM) and // K_BUFFER (fused) activate the buffer path. diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_out_tile_store.glslh b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_out_tile_store.glslh index 382747bf9a4..2c04170bab4 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_out_tile_store.glslh +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_out_tile_store.glslh @@ -30,7 +30,7 @@ #extension GL_EXT_control_flow_attributes : require -#include "linear_fp_output_tile.glslh" +#include "gemm/tile_utils/linear_fp_output_tile.glslh" void store_out_d4( const LINEAR_FP_OUTPUT_TILE_VEC4_T out_texel, diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_q_projected_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_q_projected_tile_load.glslh index 752762b623d..40ed5b9e181 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_q_projected_tile_load.glslh +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_q_projected_tile_load.glslh @@ -28,7 +28,7 @@ #extension GL_EXT_control_flow_attributes : require -#include "linear_fp_input_tile.glslh" +#include "gemm/tile_utils/linear_fp_input_tile.glslh" LINEAR_FP_INPUT_TILE_VEC4_T load_q_projected_d4( const int d4, diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_v_cache_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_v_cache_tile_load.glslh index 98516744b44..fc2b280dedd 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_v_cache_tile_load.glslh +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_v_cache_tile_load.glslh @@ -30,7 +30,7 @@ #extension GL_EXT_control_flow_attributes : require -#include "linear_fp_weight_tile.glslh" +#include "gemm/tile_utils/linear_fp_weight_tile.glslh" #ifndef SDPA_V_BUF #define SDPA_V_BUF t_v diff --git a/backends/vulkan/runtime/graph/ops/glsl/unpack_4w4c_and_dequantize.glsl b/backends/vulkan/runtime/graph/ops/glsl/unpack_4w4c_and_dequantize.glsl index fd7fa0aa7f1..41a106ab5ff 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/unpack_4w4c_and_dequantize.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/unpack_4w4c_and_dequantize.glsl @@ -42,8 +42,8 @@ layout(push_constant) uniform restrict Block { layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; -#include "linear_fp_input_tile.glslh" -#include "linear_int8_input_tile.glslh" +#include "gemm/tile_utils/linear_fp_input_tile.glslh" +#include "gemm/tile_utils/linear_int8_input_tile.glslh" void load_packed_int8_tile( out Int8InputTile int8_tile, diff --git a/backends/vulkan/runtime/graph/ops/impl/EmbeddingQ4gsw.cpp b/backends/vulkan/runtime/graph/ops/impl/EmbeddingQ4gsw.cpp index ad4786d1c60..2d0d8c623f4 100644 --- a/backends/vulkan/runtime/graph/ops/impl/EmbeddingQ4gsw.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/EmbeddingQ4gsw.cpp @@ -9,8 +9,8 @@ #include #include -#include #include +#include #include #include diff --git a/backends/vulkan/runtime/graph/ops/impl/GemmCommon.h b/backends/vulkan/runtime/graph/ops/impl/GemmCommon.h deleted file mode 100644 index eadeeee7080..00000000000 --- a/backends/vulkan/runtime/graph/ops/impl/GemmCommon.h +++ /dev/null @@ -1,45 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -#pragma once - -#include -#include - -namespace vkcompute { - -// Helpers shared by the linear/matmul tiled and coopmat dispatch paths. - -// Prepack a floating-point weight tensor for the linear/matmul kernels. -// Source layout: [N, K] (is_transposed=true) or [K, N] (is_transposed=false), -// optionally batched. Output layout: 4OC x 4IC blocked, packed as -// kWidthPacked. When force_buffer is true, the packed tensor uses buffer -// storage (required by the coopmat shader); otherwise texture2d is used -// when it fits within max_texture2d_dim. -ValueRef prepack_fp_linear_weight( - ComputeGraph& graph, - const ValueRef weight_data, - bool is_transposed, - int64_t B, - bool force_buffer = false); - -// Resize logic for linear-shaped output: takes M from the input's penultimate -// dim, N from resize_args[0], optional batch from input's leading dim. -void resize_linear_node( - ComputeGraph* graph, - const std::vector& args, - const std::vector& resize_args); - -// Resize logic for matmul (mat1 @ mat2): output rows from mat1, cols from -// mat2, batch dims propagated from mat1. -void resize_matmul_tiled_node( - ComputeGraph* graph, - const std::vector& args, - const std::vector& resize_args); - -} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.h b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.h deleted file mode 100644 index 4ab87f8ac33..00000000000 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.h +++ /dev/null @@ -1,30 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -#pragma once - -#include -#include -#include - -namespace vkcompute { - -LocalWorkGroup quantized_linear_lwg( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const GlobalWorkGrid& gwg, - const std::vector& args, - const std::vector& resize_args); - -ValueRef prepack_quantized_linear_weight( - ComputeGraph& graph, - const QuantizationConfig& weight_quant_config, - const ValueRef qmat2_data, - const bool use_unsigned_dot = false); - -} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp index 39854b23d02..bc315727bd2 100644 --- a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp @@ -12,10 +12,9 @@ #include #include -#include #include #include -#include +#include #include diff --git a/backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1dPW.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1dPW.cpp index 225a5118688..22457c37e24 100644 --- a/backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1dPW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1dPW.cpp @@ -9,7 +9,6 @@ #include #include -#include #include #include diff --git a/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/im2col/Q8taConv2dIm2Col.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/im2col/Q8taConv2dIm2Col.cpp index 39f9a44b476..0fcb2105b2e 100644 --- a/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/im2col/Q8taConv2dIm2Col.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/im2col/Q8taConv2dIm2Col.cpp @@ -14,9 +14,9 @@ #include #include -#include #include #include +#include #include #include diff --git a/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/quantized_legacy/QuantizedConvolution.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/quantized_legacy/QuantizedConvolution.cpp index 595395b6f82..d58f669362f 100644 --- a/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/quantized_legacy/QuantizedConvolution.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/quantized_legacy/QuantizedConvolution.cpp @@ -11,13 +11,13 @@ #include #include #include -#include #include #include #include #include #include #include +#include #include namespace vkcompute { diff --git a/backends/vulkan/runtime/graph/ops/impl/GemmCommon.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/GemmCommon.cpp similarity index 98% rename from backends/vulkan/runtime/graph/ops/impl/GemmCommon.cpp rename to backends/vulkan/runtime/graph/ops/impl/gemm/GemmCommon.cpp index b763dcaeb38..61e053e0751 100644 --- a/backends/vulkan/runtime/graph/ops/impl/GemmCommon.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/GemmCommon.cpp @@ -6,7 +6,7 @@ * LICENSE file in the root directory of this source tree. */ -#include +#include #include #include diff --git a/backends/vulkan/runtime/graph/ops/impl/GemmCoopmat.h b/backends/vulkan/runtime/graph/ops/impl/gemm/GemmCommon.h similarity index 59% rename from backends/vulkan/runtime/graph/ops/impl/GemmCoopmat.h rename to backends/vulkan/runtime/graph/ops/impl/gemm/GemmCommon.h index 7be7e8bc157..eb1c1c74c24 100644 --- a/backends/vulkan/runtime/graph/ops/impl/GemmCoopmat.h +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/GemmCommon.h @@ -9,13 +9,40 @@ #pragma once #include +#include namespace vkcompute { -// Cooperative-matrix GEMM dispatch. The coopmat shader (coopmat_mm.glsl) is -// shared across linear/matmul; the two add_*_node entry points below differ -// only in input layout (prepacked weight vs runtime mat2) and bias handling. -// +// Helpers shared by the linear/matmul tiled and coopmat dispatch paths. + +// Prepack a floating-point weight tensor for the linear/matmul kernels. +// Source layout: [N, K] (is_transposed=true) or [K, N] (is_transposed=false), +// optionally batched. Output layout: 4OC x 4IC blocked, packed as +// kWidthPacked. When force_buffer is true, the packed tensor uses buffer +// storage (required by the coopmat shader); otherwise texture2d is used +// when it fits within max_texture2d_dim. +ValueRef prepack_fp_linear_weight( + ComputeGraph& graph, + const ValueRef weight_data, + bool is_transposed, + int64_t B, + bool force_buffer = false); + +// Resize logic for linear-shaped output: takes M from the input's penultimate +// dim, N from resize_args[0], optional batch from input's leading dim. +void resize_linear_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args); + +// Resize logic for matmul (mat1 @ mat2): output rows from mat1, cols from +// mat2, batch dims propagated from mat1. +void resize_matmul_tiled_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args); + +// Cooperative-matrix GEMM (coopmat_mm.glsl, shared by linear and matmul). // Tile dimensions match the coopmat_mm.yaml defaults; they participate in // the eligibility check below because the shader has no partial-tile or // K-tail handling — misaligned shapes must fall back to the tiled path. @@ -52,19 +79,4 @@ inline bool is_coopmat_eligible( N % kCoopmatTileN == 0 && K % kCoopmatTileK == 0; } -void add_linear_coopmat_node( - ComputeGraph& graph, - const ValueRef input, - const ValueRef packed_weight, - const ValueRef packed_bias, - bool has_bias, - const ValueRef out, - int32_t weight_B = 1); - -void add_matmul_coopmat_node( - ComputeGraph& graph, - const ValueRef mat1, - const ValueRef mat2, - const ValueRef out); - } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Linear.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/Linear.cpp similarity index 95% rename from backends/vulkan/runtime/graph/ops/impl/Linear.cpp rename to backends/vulkan/runtime/graph/ops/impl/gemm/linear/Linear.cpp index 86f27b63586..722cde98515 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Linear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/Linear.cpp @@ -9,11 +9,11 @@ #include #include -#include -#include -#include -#include #include +#include +#include +#include +#include #include #include diff --git a/backends/vulkan/runtime/graph/ops/impl/Linear.h b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/Linear.h similarity index 100% rename from backends/vulkan/runtime/graph/ops/impl/Linear.h rename to backends/vulkan/runtime/graph/ops/impl/gemm/linear/Linear.h diff --git a/backends/vulkan/runtime/graph/ops/impl/GemmCoopmat.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/LinearCoopmat.cpp similarity index 53% rename from backends/vulkan/runtime/graph/ops/impl/GemmCoopmat.cpp rename to backends/vulkan/runtime/graph/ops/impl/gemm/linear/LinearCoopmat.cpp index 6b62fa6a9ba..5b3e898cff6 100644 --- a/backends/vulkan/runtime/graph/ops/impl/GemmCoopmat.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/LinearCoopmat.cpp @@ -6,19 +6,15 @@ * LICENSE file in the root directory of this source tree. */ -#include +#include #include -#include +#include -#include -#include #include namespace vkcompute { -// ── Linear coopmat ── - static vkapi::ShaderInfo pick_linear_coopmat_shader( ComputeGraph* graph, const std::vector& args, @@ -40,8 +36,10 @@ static GlobalWorkGrid pick_linear_coopmat_gwg( (void)resize_args; const ValueRef out = args.at(0).refs.at(0); const auto out_sizes = graph->sizes_of(out); - uint32_t M = out_sizes.at(out_sizes.size() - 2); - uint32_t N = out_sizes.at(out_sizes.size() - 1); + uint32_t M = + utils::safe_downcast(out_sizes.at(out_sizes.size() - 2)); + uint32_t N = + utils::safe_downcast(out_sizes.at(out_sizes.size() - 1)); uint32_t num_tiles_n = utils::div_up(N, kCoopmatTileN); uint32_t num_tiles_m = utils::div_up(M, kCoopmatTileM); // Each workgroup processes one WG_TILE_M x WG_TILE_N output tile via @@ -81,7 +79,8 @@ void add_linear_coopmat_node( "linear_coopmat requires buffer storage"); std::vector out_sizes = graph.sizes_of(out); - int32_t orig_N = utils::safe_downcast(out_sizes.back()); + int32_t orig_N = + utils::safe_downcast(out_sizes.at(out_sizes.size() - 1)); ValueRef orig_N_ref = graph.add_scalar(static_cast(orig_N)); ValueRef has_bias_ref = graph.add_scalar(has_bias); @@ -109,79 +108,4 @@ void add_linear_coopmat_node( resize_linear_node)); } -// ── Matmul coopmat ── - -static vkapi::ShaderInfo pick_matmul_coopmat_shader( - ComputeGraph* graph, - const std::vector& args, - const std::vector& resize_args) { - (void)resize_args; - const ValueRef out = args.at(0).refs.at(0); - std::string kernel_name = "matmul_coopmat"; - kernel_name.reserve(kShaderNameReserve); - add_dtype_suffix(kernel_name, graph->dtype_of(out)); - return VK_KERNEL_FROM_STR(kernel_name); -} - -static GlobalWorkGrid pick_matmul_coopmat_gwg( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const std::vector& args, - const std::vector& resize_args) { - (void)shader; - (void)resize_args; - const ValueRef out = args.at(0).refs.at(0); - const auto out_sizes = graph->sizes_of(out); - uint32_t M = out_sizes.at(out_sizes.size() - 2); - uint32_t N = out_sizes.at(out_sizes.size() - 1); - uint32_t num_tiles_n = utils::div_up(N, kCoopmatTileN); - uint32_t num_tiles_m = utils::div_up(M, kCoopmatTileM); - // Each workgroup processes one WG_TILE_M x WG_TILE_N output tile via - // cooperative-matrix MMAs across its 4 subgroups. We want the dispatch - // to launch exactly num_tiles_n x num_tiles_m workgroups. - // - // The framework computes the group count as - // group_count = div_up(gwg, lwg) - // (see Context.cpp + Command.cpp). With lwg = (kCoopmatInvocations, - // 1, 1), multiplying num_tiles_n by kCoopmatInvocations cancels the - // div, yielding group_count.x = num_tiles_n. - return GlobalWorkGrid( - {num_tiles_n * kCoopmatInvocations, num_tiles_m, 1u}, - kTiledWorkGrid, - LocalWorkGroup(kCoopmatInvocations, 1u, 1u)); -} - -void add_matmul_coopmat_node( - ComputeGraph& graph, - const ValueRef mat1, - const ValueRef mat2, - const ValueRef out) { - VK_CHECK_COND(graph.packed_dim_of(mat1) == WHCN::kWidthDim); - VK_CHECK_COND(graph.packed_dim_of(mat2) == WHCN::kWidthDim); - VK_CHECK_COND(graph.packed_dim_of(out) == WHCN::kWidthDim); - VK_CHECK_COND( - graph.storage_type_of(out) == utils::kBuffer, - "matmul_coopmat requires buffer storage"); - - ValueRef has_bias_ref = graph.add_scalar(false); - - graph.execute_nodes().emplace_back(new DynamicDispatchNode( - graph, - pick_matmul_coopmat_shader, - pick_matmul_coopmat_gwg, - pick_required_lwg, - // Inputs and Outputs — same binding order as matmul_vec - {{out, vkapi::kWrite}, {{mat1, mat2}, vkapi::kRead}}, - // Shader params buffers — same UBOs as matmul_vec - {graph.sizes_ubo(mat1), graph.sizes_ubo(mat2)}, - // Push Constants - {}, - // Specialization Constants - {}, - // Resize Args - {has_bias_ref}, - // Resizing Logic - resize_matmul_tiled_node)); -} - } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/LinearCoopmat.h b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/LinearCoopmat.h new file mode 100644 index 00000000000..e5aac00c04e --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/LinearCoopmat.h @@ -0,0 +1,24 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace vkcompute { + +void add_linear_coopmat_node( + ComputeGraph& graph, + const ValueRef input, + const ValueRef packed_weight, + const ValueRef packed_bias, + bool has_bias, + const ValueRef out, + int32_t weight_B = 1); + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/LinearDw.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/backward/LinearDw.cpp similarity index 100% rename from backends/vulkan/runtime/graph/ops/impl/LinearDw.cpp rename to backends/vulkan/runtime/graph/ops/impl/gemm/linear/backward/LinearDw.cpp diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/q8ta/Q8taLinear.cpp similarity index 97% rename from backends/vulkan/runtime/graph/ops/impl/Q8taLinear.cpp rename to backends/vulkan/runtime/graph/ops/impl/gemm/linear/q8ta/Q8taLinear.cpp index ea1a87b910d..fba34214dab 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/q8ta/Q8taLinear.cpp @@ -6,14 +6,14 @@ * LICENSE file in the root directory of this source tree. */ -#include +#include #include #include -#include #include #include +#include #include namespace vkcompute { diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taLinear.h b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/q8ta/Q8taLinear.h similarity index 100% rename from backends/vulkan/runtime/graph/ops/impl/Q8taLinear.h rename to backends/vulkan/runtime/graph/ops/impl/gemm/linear/q8ta/Q8taLinear.h diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taLinearGemv.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/q8ta/Q8taLinearGemv.cpp similarity index 97% rename from backends/vulkan/runtime/graph/ops/impl/Q8taLinearGemv.cpp rename to backends/vulkan/runtime/graph/ops/impl/gemm/linear/q8ta/Q8taLinearGemv.cpp index a3944641839..2861bb1d7b9 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taLinearGemv.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/q8ta/Q8taLinearGemv.cpp @@ -6,14 +6,14 @@ * LICENSE file in the root directory of this source tree. */ -#include +#include #include #include -#include #include #include +#include #include namespace vkcompute { diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taLinearGemv.h b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/q8ta/Q8taLinearGemv.h similarity index 100% rename from backends/vulkan/runtime/graph/ops/impl/Q8taLinearGemv.h rename to backends/vulkan/runtime/graph/ops/impl/gemm/linear/q8ta/Q8taLinearGemv.h diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/QuantizedLinear.cpp similarity index 63% rename from backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp rename to backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/QuantizedLinear.cpp index 3f1ea14f20b..3b3e11c8498 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/QuantizedLinear.cpp @@ -6,13 +6,13 @@ * LICENSE file in the root directory of this source tree. */ -#include +#include #include -#include #include -#include #include +#include +#include #include namespace vkcompute { @@ -51,26 +51,6 @@ void resize_linear_qw_node( graph->virtual_resize(output, new_out_sizes); } -// Per-shader coopmat tile geometry (must match each shader's yaml). -// Workgroup size (wg_size) = SG_GRID_X * SG_GRID_Y * SUBGROUP_SIZE. -// linear_q4gsw_coopmat 128x64x16, 2x2 subgroups x 32 (forced) -> 128 -// linear_dq8ca_q4gsw_coopmat 128x64x32, 2x2 subgroups x 64 -> 256 -// (The int8-MMA shaders stay on wave64: int8 WMMA at forced subgroup 32 -// crashes the Xclipse PAL compiler.) -struct CoopmatTileDims { - uint32_t m; - uint32_t n; - uint32_t k; - // Threads per workgroup = SG_GRID_X * SG_GRID_Y * SUBGROUP_SIZE. MUST match - // the WG_SIZE the shader yaml resolves to, or the launched thread count won't - // match the shader's staging passes (out-of-bounds). - uint32_t wg_size; -}; -// linear_qw_coopmat.yaml: 128x64, 2x2 subgroup grid, sg32 -> WG_SIZE 128. -constexpr CoopmatTileDims kQ4gswCoopmatDims = {128, 64, 16, 128}; -// linear_dq8ca_qw_coopmat.yaml: 128x64, 2x2 grid, sg64 -> WG_SIZE 256. -constexpr CoopmatTileDims kDq8caQ4gswCoopmatDims = {128, 64, 32, 256}; - static CoopmatTileDims coopmat_tile_dims(const std::string& kernel_name) { // Exact prefix matches (the "linear_dq8ca_*" names must not match the // weight-only entries). @@ -88,16 +68,19 @@ GlobalWorkGrid quantized_linear_gwg( const vkapi::ShaderInfo& shader, const std::vector& args, const std::vector& resize_args) { + (void)resize_args; const ValueRef out = args.at(0).refs.at(0); std::vector out_sizes = graph->sizes_of(out); // width - const uint32_t N = utils::val_at(-1, out_sizes); + const uint32_t N = + utils::safe_downcast(utils::val_at(-1, out_sizes)); // height - const uint32_t M = utils::val_at(-2, out_sizes); + const uint32_t M = + utils::safe_downcast(utils::val_at(-2, out_sizes)); // Coopmat variants dispatch a 256-thread WG per 64x64 output tile. Mirrors - // GemmCoopmat.cpp's pick_linear_coopmat_gwg — the multiplication + // LinearCoopmat.cpp's pick_linear_coopmat_gwg — the multiplication // by kCoopmatInvocations cancels the framework's div_up, since // lwg = {256, 1, 1}. if (shader.kernel_name.find("_coopmat") != std::string::npos) { @@ -159,15 +142,15 @@ LocalWorkGroup quantized_linear_lwg( // (M, N, K, dtype, output_storage, group_size) tuple. Preconditions match what // linear_q4gsw_coopmat.glsl assumes; the subgroup_size == 64 check scopes this // to wave64 devices (e.g. AMD RDNA), which the coopmat tiling is tuned for. -static bool can_use_q4gsw_coopmat( +bool can_use_q4gsw_coopmat( ComputeGraph* graph, const ValueRef output, const ValueRef fp_input, int64_t group_size, const ValueRef bias, - int64_t tile_m = kCoopmatTileM, - int64_t tile_n = kCoopmatTileN, - int64_t tile_k = kCoopmatTileK) { + int64_t tile_m, + int64_t tile_n, + int64_t tile_k) { // The coopmat shaders only build HAS_BIAS=false variants, so they would // silently drop a bias. Fall back to the tiled path (which applies bias at // runtime via the apply_bias spec constant) whenever a bias is present. @@ -288,56 +271,6 @@ vkapi::ShaderInfo pick_linear_qw_shader( return VK_KERNEL_FROM_STR(kernel_name); } -vkapi::ShaderInfo pick_linear_dqa_qw_shader( - ComputeGraph* graph, - const std::vector& args, - const std::vector& resize_args) { - (void)resize_args; - - const ValueRef out = args.at(0).refs.at(0); - const ValueRef fp_input = args.at(1).refs.at(0); - const ValueRef int_input = args.at(1).refs.at(1); - (void)int_input; - const ValueRef input_zp = args.at(1).refs.at(4); - const ValueRef int_weight = args.at(1).refs.at(5); - - const bool weight_is_4bit = resize_args.at(0) != kDummyValueRef; - const bool is_gemv_case = is_gemv(graph, fp_input); - - // Use the coopmat shader for 4-bit dq8ca dispatches when the device - // enumerates VK_COMPONENT_TYPE_SINT8_KHR in its cooperative matrix property - // list and the shape aligns; tiled otherwise. - if (weight_is_4bit && !is_gemv_case && - graph->context()->adapter_ptr()->supports_int8_cooperative_matrix()) { - const int64_t group_size = - graph->extract_scalar(resize_args.at(0)); - if (can_use_q4gsw_coopmat( - graph, - out, - fp_input, - group_size, - resize_args.at(2), - kDq8caQ4gswCoopmatDims.m, - kDq8caQ4gswCoopmatDims.n, - kDq8caQ4gswCoopmatDims.k)) { - std::string kernel_name = "linear_dq8ca_q4gsw_coopmat"; - add_storage_type_suffix(kernel_name, graph->storage_type_of(out)); - add_storage_type_suffix(kernel_name, graph->storage_type_of(int_weight)); - add_dtype_suffix(kernel_name, graph->dtype_of(out)); - return VK_KERNEL_FROM_STR(kernel_name); - } - } - - std::string kernel_name = "linear_dq8ca_q4gsw"; - kernel_name += is_gemv_case ? "_coop" : "_tiled"; - add_storage_type_suffix(kernel_name, graph->storage_type_of(out)); - add_storage_type_suffix(kernel_name, graph->storage_type_of(int_weight)); - add_dtype_suffix(kernel_name, graph->dtype_of(out)); - add_zp_dtype_mode_suffix(kernel_name, graph->dtype_of(input_zp)); - - return VK_KERNEL_FROM_STR(kernel_name); -} - // // Prepacking nodes // @@ -500,6 +433,7 @@ void add_linear_qw_node( const ValueRef bias_data, const ValueRef packed_bias, const ValueRef output) { + (void)packed_weight_zeros; // Only certain quantization types supported at the moment VK_CHECK_COND( weight_quant_config.granularity == kPerChannel || @@ -556,160 +490,6 @@ void add_linear_qw_node( resize_linear_qw_node)); } -void add_linear_qa_qw_node( - ComputeGraph& graph, - const QuantizationConfig& input_quant_config, - const QuantizationConfig& weight_quant_config, - const ValueRef fp_input, - const ValueRef packed_int_input, - const ValueRef packed_input_scale, - const ValueRef packed_input_zp, - const ValueRef input_scale_data, - const ValueRef input_zp_data, - const ValueRef weight_data, - const ValueRef packed_weight, - const ValueRef packed_weight_sums, - const ValueRef packed_weight_scales, - const ValueRef group_size, - const ValueRef bias_data, - const ValueRef packed_bias, - const ValueRef output) { - VK_CHECK_COND(input_quant_config.granularity == kPerTensor); - VK_CHECK_COND(input_quant_config.nbits == 8); - VK_CHECK_COND(weight_quant_config.granularity == kPerChannel); - VK_CHECK_COND(weight_quant_config.is_symmetric); - VK_CHECK_COND(weight_quant_config.nbits == 8); - - float scale = graph.extract_scalar(input_scale_data); - int32_t zp = graph.extract_scalar(input_zp_data); - - // Get shader for quantized linear - std::string kernel_name = "linear_q8ta_q8csw_tiled"; - add_storage_type_suffix(kernel_name, graph.storage_type_of(output)); - add_storage_type_suffix(kernel_name, graph.storage_type_of(packed_int_input)); - add_storage_type_suffix(kernel_name, graph.storage_type_of(packed_weight)); - add_dtype_suffix(kernel_name, graph.dtype_of(output)); - vkapi::ShaderInfo shader = VK_KERNEL_FROM_STR(kernel_name); - - vkapi::ParamsBindList param_buffers = { - graph.sizes_ubo(output), graph.sizes_ubo(packed_int_input)}; - - std::vector push_constants = { - PushConstantDataInfo(&scale, sizeof(scale)), - PushConstantDataInfo(&zp, sizeof(zp)), - }; - - uint32_t apply_bias = 1; - if (graph.val_is_none(bias_data)) { - apply_bias = 0; - } - - graph.execute_nodes().emplace_back(new DynamicDispatchNode( - graph, - VK_KERNEL_FROM_STR(kernel_name), - quantized_linear_gwg, - quantized_linear_lwg, - // Inputs and Outputs - {{output, vkapi::kWrite}, - {{packed_int_input, - packed_weight, - packed_weight_sums, - packed_weight_scales, - packed_bias}, - vkapi::kRead}}, - // Shader params buffers - param_buffers, - // Push Constants - push_constants, - // Specialization Constants - {apply_bias}, - // Resize args - {fp_input}, - // Resizing Logic - nullptr)); -} - -void add_linear_dqa_qw_node( - ComputeGraph& graph, - const QuantizationConfig& input_quant_config, - const QuantizationConfig& weight_quant_config, - const ValueRef fp_input, - const ValueRef packed_int_input, - const ValueRef int_input_sums, - const ValueRef packed_input_scale, - const ValueRef packed_input_zp, - const ValueRef input_scale_data, - const ValueRef input_zp_data, - const ValueRef weight_data, - const ValueRef packed_weight, - const ValueRef packed_weight_sums, - const ValueRef packed_weight_scales, - const ValueRef group_size, - const ValueRef bias_data, - const ValueRef packed_bias, - const ValueRef output) { - VK_CHECK_COND(input_quant_config.granularity == kPerChannel); - VK_CHECK_COND(input_quant_config.nbits == 8); - VK_CHECK_COND(input_quant_config.is_dynamic); - - VK_CHECK_COND(weight_quant_config.granularity == kPerGroup); - VK_CHECK_COND(weight_quant_config.is_symmetric); - VK_CHECK_COND(weight_quant_config.nbits == 4); - - vkapi::ParamsBindList param_buffers = { - graph.sizes_ubo(output), graph.sizes_ubo(fp_input)}; - - uint32_t apply_bias = 1; - if (graph.val_is_none(bias_data)) { - apply_bias = 0; - } - - int32_t K4_per_group = 0; - int32_t coopmat_k_iters = 0; - const int32_t K_dim = graph.size_at(-1, fp_input); - if (weight_quant_config.nbits == 4) { - int32_t group_size_val = graph.extract_scalar(group_size); - K4_per_group = utils::div_up(group_size_val, int32_t(4)); - coopmat_k_iters = K_dim / group_size_val; - } - - const ValueRef is_4bit_flag = - weight_quant_config.nbits == 4 ? group_size : kDummyValueRef; - - graph.execute_nodes().emplace_back(new DynamicDispatchNode( - graph, - pick_linear_dqa_qw_shader, - quantized_linear_gwg, - quantized_linear_lwg, - // Inputs and Outputs - {{output, vkapi::kWrite}, - {{fp_input, - packed_int_input, - int_input_sums, - packed_input_scale, - packed_input_zp, - packed_weight, - packed_weight_sums, - packed_weight_scales, - packed_bias}, - vkapi::kRead}}, - // Shader params buffers - param_buffers, - // Push Constants - {}, - // Specialization Constants - // 4th spec const: output width N for coopMatStore (see - // add_linear_qw_node). - {apply_bias, - K4_per_group, - coopmat_k_iters, - graph.size_at(-1, output)}, - // Resize args (resize_args.at(2) = bias_data, read by the coopmat gate) - {is_4bit_flag, weight_data, bias_data}, - // Resizing Logic - resize_linear_qw_node)); -} - // // High level operator impl // @@ -728,6 +508,7 @@ void quantized_linear_impl( const ValueRef group_size, const ValueRef bias_data, const ValueRef output) { + (void)weight_zeros_data; std::vector input_sizes = graph.sizes_of(fp_input); std::vector weight_sizes = graph.sizes_of(weight_data); @@ -898,141 +679,4 @@ void quantized_linear_impl( output); } -void linear_q8ta_q8csw(ComputeGraph& graph, const std::vector& args) { - int32_t idx = 0; - const ValueRef fp_input = args.at(idx++); - const ValueRef input_scale = args.at(idx++); - const ValueRef input_zp = args.at(idx++); - const ValueRef weight_data = args.at(idx++); - const ValueRef weight_sums_data = args.at(idx++); - const ValueRef weight_scales_data = args.at(idx++); - const ValueRef bias_data = args.at(idx++); - const ValueRef output = args.at(idx++); - - const int64_t K = graph.size_at(-1, fp_input); - - QuantizationConfig input_quant_config(8, kPerTensor, {}, false); - QuantizationConfig weight_quant_config(8, kPerChannel, {K}); - - quantized_linear_impl( - graph, - input_quant_config, - weight_quant_config, - fp_input, - input_scale, - input_zp, - weight_data, - weight_sums_data, - weight_scales_data, - kDummyValueRef, // weight_zeros_data - kDummyValueRef, // group_size - bias_data, - output); -} - -void linear_q8csw(ComputeGraph& graph, const std::vector& args) { - int32_t idx = 0; - const ValueRef fp_input = args.at(idx++); - const ValueRef weight_data = args.at(idx++); - const ValueRef weight_scales_data = args.at(idx++); - const ValueRef bias_data = args.at(idx++); - const ValueRef output = args.at(idx++); - - const int64_t K = graph.size_at(-1, fp_input); - - QuantizationConfig input_quant_config(32, kNoQuantization, {}); - QuantizationConfig weight_quant_config(8, kPerChannel, {K}); - - quantized_linear_impl( - graph, - input_quant_config, - weight_quant_config, - fp_input, - kDummyValueRef, // input scale - kDummyValueRef, // input zp - weight_data, - kDummyValueRef, // weight sums - weight_scales_data, - kDummyValueRef, // weight zeros - kDummyValueRef, // group size - bias_data, - output); -} - -// aten._weight_int8pack_mm is what the AOT weight-only int8 fusion -// (FuseQuantizedOpsTransform) emits. It carries the same operands as -// et_vk.linear_q8csw minus the bias, so it runs through the same -// implementation. -void weight_int8pack_mm( - ComputeGraph& graph, - const std::vector& args) { - int32_t idx = 0; - const ValueRef fp_input = args.at(idx++); - const ValueRef weight_data = args.at(idx++); - const ValueRef weight_scales_data = args.at(idx++); - const ValueRef output = args.at(idx++); - - const int64_t K = graph.size_at(-1, fp_input); - - QuantizationConfig input_quant_config(32, kNoQuantization, {}); - QuantizationConfig weight_quant_config(8, kPerChannel, {K}); - - quantized_linear_impl( - graph, - input_quant_config, - weight_quant_config, - fp_input, - kDummyValueRef, // input scale - kDummyValueRef, // input zp - weight_data, - kDummyValueRef, // weight sums - weight_scales_data, - kDummyValueRef, // weight zeros - kDummyValueRef, // group size - kDummyValueRef, // bias - output); -} - -void linear_dq8ca_q4gsw( - ComputeGraph& graph, - const std::vector& args) { - int32_t idx = 0; - const ValueRef fp_input = args.at(idx++); - const ValueRef input_scale = args.at(idx++); - const ValueRef input_zp = args.at(idx++); - const ValueRef weight_data = args.at(idx++); - const ValueRef weight_sums_data = args.at(idx++); - const ValueRef weight_scales_data = args.at(idx++); - const ValueRef group_size = args.at(idx++); - const ValueRef bias_data = args.at(idx++); - const ValueRef output = args.at(idx++); - - const int64_t group_size_val = graph.extract_scalar(group_size); - - QuantizationConfig input_quant_config(8, kPerChannel, {}, false, true); - QuantizationConfig weight_quant_config(4, kPerGroup, {group_size_val}); - - quantized_linear_impl( - graph, - input_quant_config, - weight_quant_config, - fp_input, - input_scale, - input_zp, - weight_data, - weight_sums_data, - weight_scales_data, - kDummyValueRef, // weight_zeros_data - group_size, // group_size - bias_data, - output); -} - -REGISTER_OPERATORS { - VK_REGISTER_OP(et_vk.linear_q8ta_q8csw.default, linear_q8ta_q8csw); - VK_REGISTER_OP(et_vk.linear_q8csw.default, linear_q8csw); - VK_REGISTER_OP(aten._weight_int8pack_mm.default, weight_int8pack_mm); - VK_REGISTER_OP(et_vk.linear_dq8ca_q4gsw.default, linear_dq8ca_q4gsw); -} - } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/QuantizedLinear.h b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/QuantizedLinear.h new file mode 100644 index 00000000000..23f70416754 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/QuantizedLinear.h @@ -0,0 +1,87 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include +#include + +namespace vkcompute { + +// Per-shader coopmat tile geometry (must match each shader's yaml). +// Workgroup size (wg_size) = SG_GRID_X * SG_GRID_Y * SUBGROUP_SIZE. +// linear_q4gsw_coopmat 128x64x16, 2x2 subgroups x 32 (forced) -> 128 +// linear_dq8ca_q4gsw_coopmat 128x64x32, 2x2 subgroups x 64 -> 256 +// (The int8-MMA shaders stay on wave64: int8 WMMA at forced subgroup 32 +// crashes the Xclipse PAL compiler.) +struct CoopmatTileDims { + uint32_t m; + uint32_t n; + uint32_t k; + // Threads per workgroup = SG_GRID_X * SG_GRID_Y * SUBGROUP_SIZE. MUST match + // the WG_SIZE the shader yaml resolves to, or the launched thread count won't + // match the shader's staging passes (out-of-bounds). + uint32_t wg_size; +}; +// linear_qw_coopmat.yaml: 128x64, 2x2 subgroup grid, sg32 -> WG_SIZE 128. +constexpr CoopmatTileDims kQ4gswCoopmatDims = {128, 64, 16, 128}; +// linear_dq8ca_qw_coopmat.yaml: 128x64, 2x2 grid, sg64 -> WG_SIZE 256. +constexpr CoopmatTileDims kDq8caQ4gswCoopmatDims = {128, 64, 32, 256}; + +void resize_linear_qw_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& extra_args); + +GlobalWorkGrid quantized_linear_gwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args); + +LocalWorkGroup quantized_linear_lwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const GlobalWorkGrid& gwg, + const std::vector& args, + const std::vector& resize_args); + +bool can_use_q4gsw_coopmat( + ComputeGraph* graph, + const ValueRef output, + const ValueRef fp_input, + int64_t group_size, + const ValueRef bias, + int64_t tile_m = kCoopmatTileM, + int64_t tile_n = kCoopmatTileN, + int64_t tile_k = kCoopmatTileK); + +ValueRef prepack_quantized_linear_weight( + ComputeGraph& graph, + const QuantizationConfig& weight_quant_config, + const ValueRef qmat2_data, + const bool use_unsigned_dot = false); + +void quantized_linear_impl( + ComputeGraph& graph, + const QuantizationConfig& input_quant_config, + const QuantizationConfig& weight_quant_config, + const ValueRef fp_input, + const ValueRef input_scale, + const ValueRef input_zp, + const ValueRef weight_data, + const ValueRef weight_sums_data, + const ValueRef weight_scales_data, + const ValueRef weight_zeros_data, + const ValueRef group_size, + const ValueRef bias_data, + const ValueRef output); + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinear.cpp new file mode 100644 index 00000000000..dd13ec2ebb5 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinear.cpp @@ -0,0 +1,192 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include + +#include +#include +#include +#include + +namespace vkcompute { + +vkapi::ShaderInfo pick_linear_dqa_qw_shader( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + (void)resize_args; + + const ValueRef out = args.at(0).refs.at(0); + const ValueRef fp_input = args.at(1).refs.at(0); + const ValueRef int_input = args.at(1).refs.at(1); + (void)int_input; + const ValueRef input_zp = args.at(1).refs.at(4); + const ValueRef int_weight = args.at(1).refs.at(5); + + const bool weight_is_4bit = resize_args.at(0) != kDummyValueRef; + const bool is_gemv_case = is_gemv(graph, fp_input); + + // Use the coopmat shader for 4-bit dq8ca dispatches when the device + // enumerates VK_COMPONENT_TYPE_SINT8_KHR in its cooperative matrix property + // list and the shape aligns; tiled otherwise. + if (weight_is_4bit && !is_gemv_case && + graph->context()->adapter_ptr()->supports_int8_cooperative_matrix()) { + const int64_t group_size = + graph->extract_scalar(resize_args.at(0)); + if (can_use_q4gsw_coopmat( + graph, + out, + fp_input, + group_size, + resize_args.at(2), + kDq8caQ4gswCoopmatDims.m, + kDq8caQ4gswCoopmatDims.n, + kDq8caQ4gswCoopmatDims.k)) { + std::string kernel_name = "linear_dq8ca_q4gsw_coopmat"; + add_storage_type_suffix(kernel_name, graph->storage_type_of(out)); + add_storage_type_suffix(kernel_name, graph->storage_type_of(int_weight)); + add_dtype_suffix(kernel_name, graph->dtype_of(out)); + return VK_KERNEL_FROM_STR(kernel_name); + } + } + + std::string kernel_name = "linear_dq8ca_q4gsw"; + kernel_name += is_gemv_case ? "_coop" : "_tiled"; + add_storage_type_suffix(kernel_name, graph->storage_type_of(out)); + add_storage_type_suffix(kernel_name, graph->storage_type_of(int_weight)); + add_dtype_suffix(kernel_name, graph->dtype_of(out)); + add_zp_dtype_mode_suffix(kernel_name, graph->dtype_of(input_zp)); + + return VK_KERNEL_FROM_STR(kernel_name); +} + +void add_linear_dqa_qw_node( + ComputeGraph& graph, + const QuantizationConfig& input_quant_config, + const QuantizationConfig& weight_quant_config, + const ValueRef fp_input, + const ValueRef packed_int_input, + const ValueRef int_input_sums, + const ValueRef packed_input_scale, + const ValueRef packed_input_zp, + const ValueRef input_scale_data, + const ValueRef input_zp_data, + const ValueRef weight_data, + const ValueRef packed_weight, + const ValueRef packed_weight_sums, + const ValueRef packed_weight_scales, + const ValueRef group_size, + const ValueRef bias_data, + const ValueRef packed_bias, + const ValueRef output) { + (void)input_scale_data; + (void)input_zp_data; + VK_CHECK_COND(input_quant_config.granularity == kPerChannel); + VK_CHECK_COND(input_quant_config.nbits == 8); + VK_CHECK_COND(input_quant_config.is_dynamic); + + VK_CHECK_COND(weight_quant_config.granularity == kPerGroup); + VK_CHECK_COND(weight_quant_config.is_symmetric); + VK_CHECK_COND(weight_quant_config.nbits == 4); + + vkapi::ParamsBindList param_buffers = { + graph.sizes_ubo(output), graph.sizes_ubo(fp_input)}; + + uint32_t apply_bias = 1; + if (graph.val_is_none(bias_data)) { + apply_bias = 0; + } + + int32_t K4_per_group = 0; + int32_t coopmat_k_iters = 0; + const int32_t K_dim = graph.size_at(-1, fp_input); + if (weight_quant_config.nbits == 4) { + int32_t group_size_val = graph.extract_scalar(group_size); + K4_per_group = utils::div_up(group_size_val, int32_t(4)); + coopmat_k_iters = K_dim / group_size_val; + } + + const ValueRef is_4bit_flag = + weight_quant_config.nbits == 4 ? group_size : kDummyValueRef; + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + pick_linear_dqa_qw_shader, + quantized_linear_gwg, + quantized_linear_lwg, + // Inputs and Outputs + {{output, vkapi::kWrite}, + {{fp_input, + packed_int_input, + int_input_sums, + packed_input_scale, + packed_input_zp, + packed_weight, + packed_weight_sums, + packed_weight_scales, + packed_bias}, + vkapi::kRead}}, + // Shader params buffers + param_buffers, + // Push Constants + {}, + // Specialization Constants + // 4th spec const: output width N for coopMatStore (see + // add_linear_qw_node). + {apply_bias, + K4_per_group, + coopmat_k_iters, + graph.size_at(-1, output)}, + // Resize args (resize_args.at(2) = bias_data, read by the coopmat gate) + {is_4bit_flag, weight_data, bias_data}, + // Resizing Logic + resize_linear_qw_node)); +} + +void linear_dq8ca_q4gsw( + ComputeGraph& graph, + const std::vector& args) { + int32_t idx = 0; + const ValueRef fp_input = args.at(idx++); + const ValueRef input_scale = args.at(idx++); + const ValueRef input_zp = args.at(idx++); + const ValueRef weight_data = args.at(idx++); + const ValueRef weight_sums_data = args.at(idx++); + const ValueRef weight_scales_data = args.at(idx++); + const ValueRef group_size = args.at(idx++); + const ValueRef bias_data = args.at(idx++); + const ValueRef output = args.at(idx++); + + const int64_t group_size_val = graph.extract_scalar(group_size); + + QuantizationConfig input_quant_config(8, kPerChannel, {}, false, true); + QuantizationConfig weight_quant_config(4, kPerGroup, {group_size_val}); + + quantized_linear_impl( + graph, + input_quant_config, + weight_quant_config, + fp_input, + input_scale, + input_zp, + weight_data, + weight_sums_data, + weight_scales_data, + kDummyValueRef, // weight_zeros_data + group_size, // group_size + bias_data, + output); +} + +REGISTER_OPERATORS { + VK_REGISTER_OP(et_vk.linear_dq8ca_q4gsw.default, linear_dq8ca_q4gsw); +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinear.h b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinear.h new file mode 100644 index 00000000000..0cff158a323 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/dq8ca_q4gsw/Dq8caQ4gswLinear.h @@ -0,0 +1,36 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include + +namespace vkcompute { + +void add_linear_dqa_qw_node( + ComputeGraph& graph, + const QuantizationConfig& input_quant_config, + const QuantizationConfig& weight_quant_config, + const ValueRef fp_input, + const ValueRef packed_int_input, + const ValueRef int_input_sums, + const ValueRef packed_input_scale, + const ValueRef packed_input_zp, + const ValueRef input_scale_data, + const ValueRef input_zp_data, + const ValueRef weight_data, + const ValueRef packed_weight, + const ValueRef packed_weight_sums, + const ValueRef packed_weight_scales, + const ValueRef group_size, + const ValueRef bias_data, + const ValueRef packed_bias, + const ValueRef output); + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/legacy/Q8taQ8cswLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/legacy/Q8taQ8cswLinear.cpp new file mode 100644 index 00000000000..c6c69be6a1c --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/legacy/Q8taQ8cswLinear.cpp @@ -0,0 +1,128 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include + +#include +#include +#include + +namespace vkcompute { + +void add_linear_qa_qw_node( + ComputeGraph& graph, + const QuantizationConfig& input_quant_config, + const QuantizationConfig& weight_quant_config, + const ValueRef fp_input, + const ValueRef packed_int_input, + const ValueRef packed_input_scale, + const ValueRef packed_input_zp, + const ValueRef input_scale_data, + const ValueRef input_zp_data, + const ValueRef weight_data, + const ValueRef packed_weight, + const ValueRef packed_weight_sums, + const ValueRef packed_weight_scales, + const ValueRef group_size, + const ValueRef bias_data, + const ValueRef packed_bias, + const ValueRef output) { + VK_CHECK_COND(input_quant_config.granularity == kPerTensor); + VK_CHECK_COND(input_quant_config.nbits == 8); + VK_CHECK_COND(weight_quant_config.granularity == kPerChannel); + VK_CHECK_COND(weight_quant_config.is_symmetric); + VK_CHECK_COND(weight_quant_config.nbits == 8); + + float scale = graph.extract_scalar(input_scale_data); + int32_t zp = graph.extract_scalar(input_zp_data); + + // Get shader for quantized linear + std::string kernel_name = "linear_q8ta_q8csw_tiled"; + add_storage_type_suffix(kernel_name, graph.storage_type_of(output)); + add_storage_type_suffix(kernel_name, graph.storage_type_of(packed_int_input)); + add_storage_type_suffix(kernel_name, graph.storage_type_of(packed_weight)); + add_dtype_suffix(kernel_name, graph.dtype_of(output)); + vkapi::ShaderInfo shader = VK_KERNEL_FROM_STR(kernel_name); + + vkapi::ParamsBindList param_buffers = { + graph.sizes_ubo(output), graph.sizes_ubo(packed_int_input)}; + + std::vector push_constants = { + PushConstantDataInfo(&scale, sizeof(scale)), + PushConstantDataInfo(&zp, sizeof(zp)), + }; + + uint32_t apply_bias = 1; + if (graph.val_is_none(bias_data)) { + apply_bias = 0; + } + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + quantized_linear_gwg, + quantized_linear_lwg, + // Inputs and Outputs + {{output, vkapi::kWrite}, + {{packed_int_input, + packed_weight, + packed_weight_sums, + packed_weight_scales, + packed_bias}, + vkapi::kRead}}, + // Shader params buffers + param_buffers, + // Push Constants + push_constants, + // Specialization Constants + {apply_bias}, + // Resize args + {fp_input}, + // Resizing Logic + nullptr)); +} + +void linear_q8ta_q8csw(ComputeGraph& graph, const std::vector& args) { + int32_t idx = 0; + const ValueRef fp_input = args.at(idx++); + const ValueRef input_scale = args.at(idx++); + const ValueRef input_zp = args.at(idx++); + const ValueRef weight_data = args.at(idx++); + const ValueRef weight_sums_data = args.at(idx++); + const ValueRef weight_scales_data = args.at(idx++); + const ValueRef bias_data = args.at(idx++); + const ValueRef output = args.at(idx++); + + const int64_t K = graph.size_at(-1, fp_input); + + QuantizationConfig input_quant_config(8, kPerTensor, {}, false); + QuantizationConfig weight_quant_config(8, kPerChannel, {K}); + + quantized_linear_impl( + graph, + input_quant_config, + weight_quant_config, + fp_input, + input_scale, + input_zp, + weight_data, + weight_sums_data, + weight_scales_data, + kDummyValueRef, // weight_zeros_data + kDummyValueRef, // group_size + bias_data, + output); +} + +REGISTER_OPERATORS { + VK_REGISTER_OP(et_vk.linear_q8ta_q8csw.default, linear_q8ta_q8csw); +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/legacy/Q8taQ8cswLinear.h b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/legacy/Q8taQ8cswLinear.h new file mode 100644 index 00000000000..661da132b8c --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/legacy/Q8taQ8cswLinear.h @@ -0,0 +1,35 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include + +namespace vkcompute { + +void add_linear_qa_qw_node( + ComputeGraph& graph, + const QuantizationConfig& input_quant_config, + const QuantizationConfig& weight_quant_config, + const ValueRef fp_input, + const ValueRef packed_int_input, + const ValueRef packed_input_scale, + const ValueRef packed_input_zp, + const ValueRef input_scale_data, + const ValueRef input_zp_data, + const ValueRef weight_data, + const ValueRef packed_weight, + const ValueRef packed_weight_sums, + const ValueRef packed_weight_scales, + const ValueRef group_size, + const ValueRef bias_data, + const ValueRef packed_bias, + const ValueRef output); + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQCSNW.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/legacy/QuantizedLinearQCSNW.cpp similarity index 99% rename from backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQCSNW.cpp rename to backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/legacy/QuantizedLinearQCSNW.cpp index c2e45b96418..99b0fc1e225 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQCSNW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/legacy/QuantizedLinearQCSNW.cpp @@ -298,7 +298,7 @@ void linear_qcs4w(ComputeGraph& graph, const std::vector& args) { } REGISTER_OPERATORS { - // aten._weight_int8pack_mm is registered in QuantizedLinear.cpp, on the + // aten._weight_int8pack_mm is registered in Q8cswLinear.cpp, on the // maintained weight-only quantized linear implementation. The 8-bit path // here dispatched to linear_qcs8w_* shaders that no longer exist. VK_REGISTER_OP(et_vk.linear_qcs4w.default, linear_qcs4w); diff --git a/backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinear.cpp similarity index 99% rename from backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.cpp rename to backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinear.cpp index cf8afc8c8c1..8e7d16092c7 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinear.cpp @@ -6,7 +6,7 @@ * LICENSE file in the root directory of this source tree. */ -#include +#include #include diff --git a/backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.h b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinear.h similarity index 100% rename from backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.h rename to backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/Q4gswLinear.h diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearBackward.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/backward/QuantizedLinearBackward.cpp similarity index 97% rename from backends/vulkan/runtime/graph/ops/impl/QuantizedLinearBackward.cpp rename to backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/backward/QuantizedLinearBackward.cpp index 8bfe1681bbf..01ebb2f692d 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearBackward.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/backward/QuantizedLinearBackward.cpp @@ -11,7 +11,7 @@ #include #include -#include +#include #include diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearRequant.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/backward/QuantizedLinearRequant.cpp similarity index 100% rename from backends/vulkan/runtime/graph/ops/impl/QuantizedLinearRequant.cpp rename to backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q4gsw/backward/QuantizedLinearRequant.cpp diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q8csw/Q8cswLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q8csw/Q8cswLinear.cpp new file mode 100644 index 00000000000..89657534b9b --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/linear/qw/q8csw/Q8cswLinear.cpp @@ -0,0 +1,83 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include + +namespace vkcompute { + +void linear_q8csw(ComputeGraph& graph, const std::vector& args) { + int32_t idx = 0; + const ValueRef fp_input = args.at(idx++); + const ValueRef weight_data = args.at(idx++); + const ValueRef weight_scales_data = args.at(idx++); + const ValueRef bias_data = args.at(idx++); + const ValueRef output = args.at(idx++); + + const int64_t K = graph.size_at(-1, fp_input); + + QuantizationConfig input_quant_config(32, kNoQuantization, {}); + QuantizationConfig weight_quant_config(8, kPerChannel, {K}); + + quantized_linear_impl( + graph, + input_quant_config, + weight_quant_config, + fp_input, + kDummyValueRef, // input scale + kDummyValueRef, // input zp + weight_data, + kDummyValueRef, // weight sums + weight_scales_data, + kDummyValueRef, // weight zeros + kDummyValueRef, // group size + bias_data, + output); +} + +// aten._weight_int8pack_mm is what the AOT weight-only int8 fusion +// (FuseQuantizedOpsTransform) emits. It carries the same operands as +// et_vk.linear_q8csw minus the bias, so it runs through the same +// implementation. +void weight_int8pack_mm( + ComputeGraph& graph, + const std::vector& args) { + int32_t idx = 0; + const ValueRef fp_input = args.at(idx++); + const ValueRef weight_data = args.at(idx++); + const ValueRef weight_scales_data = args.at(idx++); + const ValueRef output = args.at(idx++); + + const int64_t K = graph.size_at(-1, fp_input); + + QuantizationConfig input_quant_config(32, kNoQuantization, {}); + QuantizationConfig weight_quant_config(8, kPerChannel, {K}); + + quantized_linear_impl( + graph, + input_quant_config, + weight_quant_config, + fp_input, + kDummyValueRef, // input scale + kDummyValueRef, // input zp + weight_data, + kDummyValueRef, // weight sums + weight_scales_data, + kDummyValueRef, // weight zeros + kDummyValueRef, // group size + kDummyValueRef, // bias + output); +} + +REGISTER_OPERATORS { + VK_REGISTER_OP(et_vk.linear_q8csw.default, linear_q8csw); + VK_REGISTER_OP(aten._weight_int8pack_mm.default, weight_int8pack_mm); +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Matmul.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/matmul/MatMul.cpp similarity index 95% rename from backends/vulkan/runtime/graph/ops/impl/Matmul.cpp rename to backends/vulkan/runtime/graph/ops/impl/gemm/matmul/MatMul.cpp index fc86255be51..e314436cc7c 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Matmul.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/matmul/MatMul.cpp @@ -9,11 +9,12 @@ #include #include -#include -#include -#include -#include #include +#include +#include +#include +#include +#include #include #include diff --git a/backends/vulkan/runtime/graph/ops/impl/MatMul.h b/backends/vulkan/runtime/graph/ops/impl/gemm/matmul/MatMul.h similarity index 100% rename from backends/vulkan/runtime/graph/ops/impl/MatMul.h rename to backends/vulkan/runtime/graph/ops/impl/gemm/matmul/MatMul.h diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/matmul/MatmulCoopmat.cpp b/backends/vulkan/runtime/graph/ops/impl/gemm/matmul/MatmulCoopmat.cpp new file mode 100644 index 00000000000..1cc9730ac9b --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/matmul/MatmulCoopmat.cpp @@ -0,0 +1,93 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include + +#include + +namespace vkcompute { + +static vkapi::ShaderInfo pick_matmul_coopmat_shader( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + (void)resize_args; + const ValueRef out = args.at(0).refs.at(0); + std::string kernel_name = "matmul_coopmat"; + kernel_name.reserve(kShaderNameReserve); + add_dtype_suffix(kernel_name, graph->dtype_of(out)); + return VK_KERNEL_FROM_STR(kernel_name); +} + +static GlobalWorkGrid pick_matmul_coopmat_gwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + (void)resize_args; + const ValueRef out = args.at(0).refs.at(0); + const auto out_sizes = graph->sizes_of(out); + uint32_t M = + utils::safe_downcast(out_sizes.at(out_sizes.size() - 2)); + uint32_t N = + utils::safe_downcast(out_sizes.at(out_sizes.size() - 1)); + uint32_t num_tiles_n = utils::div_up(N, kCoopmatTileN); + uint32_t num_tiles_m = utils::div_up(M, kCoopmatTileM); + // Each workgroup processes one WG_TILE_M x WG_TILE_N output tile via + // cooperative-matrix MMAs across its 4 subgroups. We want the dispatch + // to launch exactly num_tiles_n x num_tiles_m workgroups. + // + // The framework computes the group count as + // group_count = div_up(gwg, lwg) + // (see Context.cpp + Command.cpp). With lwg = (kCoopmatInvocations, + // 1, 1), multiplying num_tiles_n by kCoopmatInvocations cancels the + // div, yielding group_count.x = num_tiles_n. + return GlobalWorkGrid( + {num_tiles_n * kCoopmatInvocations, num_tiles_m, 1u}, + kTiledWorkGrid, + LocalWorkGroup(kCoopmatInvocations, 1u, 1u)); +} + +void add_matmul_coopmat_node( + ComputeGraph& graph, + const ValueRef mat1, + const ValueRef mat2, + const ValueRef out) { + VK_CHECK_COND(graph.packed_dim_of(mat1) == WHCN::kWidthDim); + VK_CHECK_COND(graph.packed_dim_of(mat2) == WHCN::kWidthDim); + VK_CHECK_COND(graph.packed_dim_of(out) == WHCN::kWidthDim); + VK_CHECK_COND( + graph.storage_type_of(out) == utils::kBuffer, + "matmul_coopmat requires buffer storage"); + + ValueRef has_bias_ref = graph.add_scalar(false); + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + pick_matmul_coopmat_shader, + pick_matmul_coopmat_gwg, + pick_required_lwg, + // Inputs and Outputs — same binding order as matmul_vec + {{out, vkapi::kWrite}, {{mat1, mat2}, vkapi::kRead}}, + // Shader params buffers — same UBOs as matmul_vec + {graph.sizes_ubo(mat1), graph.sizes_ubo(mat2)}, + // Push Constants + {}, + // Specialization Constants + {}, + // Resize Args + {has_bias_ref}, + // Resizing Logic + resize_matmul_tiled_node)); +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/gemm/matmul/MatmulCoopmat.h b/backends/vulkan/runtime/graph/ops/impl/gemm/matmul/MatmulCoopmat.h new file mode 100644 index 00000000000..4cf825fc838 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/gemm/matmul/MatmulCoopmat.h @@ -0,0 +1,21 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace vkcompute { + +void add_matmul_coopmat_node( + ComputeGraph& graph, + const ValueRef mat1, + const ValueRef mat2, + const ValueRef out); + +} // namespace vkcompute diff --git a/backends/vulkan/test/custom_ops/impl/TestFpaQ4gswLinear.cpp b/backends/vulkan/test/custom_ops/impl/TestFpaQ4gswLinear.cpp index a9e5ad96e95..2da2cd6ceb7 100644 --- a/backends/vulkan/test/custom_ops/impl/TestFpaQ4gswLinear.cpp +++ b/backends/vulkan/test/custom_ops/impl/TestFpaQ4gswLinear.cpp @@ -10,9 +10,9 @@ #include #include -#include #include #include +#include #include namespace vkcompute { diff --git a/backends/vulkan/test/custom_ops/impl/TestMatmulLinear.cpp b/backends/vulkan/test/custom_ops/impl/TestMatmulLinear.cpp index 60eeb2b034c..fc5fb20cab4 100644 --- a/backends/vulkan/test/custom_ops/impl/TestMatmulLinear.cpp +++ b/backends/vulkan/test/custom_ops/impl/TestMatmulLinear.cpp @@ -9,11 +9,12 @@ #include #include -#include -#include -#include -#include #include +#include +#include +#include +#include +#include #include namespace vkcompute {