Skip to content

Commit 98bb579

Browse files
reeselevinereeselevine
authored
ggml-webgpu: fix buffer aliasing for ssm_scan and refactor aliasing logic (ggml-org#22456)
* Refactor buffer aliasing to be part of shader lib decisions * cleanup * formatting
1 parent f42e29f commit 98bb579

6 files changed

Lines changed: 325 additions & 237 deletions

File tree

‎ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp‎

Lines changed: 106 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -26,21 +26,21 @@
2626
// Matrix multiplication parameters
2727

2828
// Register tiling parameters
29-
#define WEBGPU_MUL_MAT_TILE_M 4
30-
#define WEBGPU_MUL_MAT_TILE_N 4
31-
#define WEBGPU_MUL_MAT_WG_SIZE_M 8
32-
#define WEBGPU_MUL_MAT_WG_SIZE_N 8
29+
#define WEBGPU_MUL_MAT_TILE_M 4
30+
#define WEBGPU_MUL_MAT_TILE_N 4
31+
#define WEBGPU_MUL_MAT_WG_SIZE_M 8
32+
#define WEBGPU_MUL_MAT_WG_SIZE_N 8
3333
#define WEBGPU_MUL_MAT_REG_TILE_K_FLOAT 8
3434
#define WEBGPU_MUL_MAT_REG_TILE_K_QUANT 32
3535

3636
// Subgroup matrix parameters
3737
// The number of subgroups in the M dimension
38-
#define WEBGPU_MUL_MAT_SUBGROUP_M 2
38+
#define WEBGPU_MUL_MAT_SUBGROUP_M 2
3939
// The number of subgroups in the N dimension
40-
#define WEBGPU_MUL_MAT_SUBGROUP_N 4
40+
#define WEBGPU_MUL_MAT_SUBGROUP_N 4
4141
// The number of subgroup matrices each subgroup accumulates over
42-
#define WEBGPU_MUL_MAT_SUBGROUP_MATRIX_M 4
43-
#define WEBGPU_MUL_MAT_SUBGROUP_MATRIX_N 2
42+
#define WEBGPU_MUL_MAT_SUBGROUP_MATRIX_M 4
43+
#define WEBGPU_MUL_MAT_SUBGROUP_MATRIX_N 2
4444
#define WEBGPU_MUL_MAT_SUBGROUP_TILE_K_FLOAT 32
4545
#define WEBGPU_MUL_MAT_SUBGROUP_TILE_K_QUANT 32
4646

@@ -59,19 +59,32 @@ template <typename T> inline void ggml_webgpu_hash_combine(size_t & seed, const
5959
seed ^= std::hash<T>{}(value) + 0x9e3779b9 + (seed << 6) + (seed >> 2);
6060
}
6161

62+
// Calculates base address of a tensor ignoring the fake base pointer
63+
inline uintptr_t ggml_webgpu_tensor_addr(const ggml_tensor * tensor) {
64+
const ggml_tensor * base_tensor = tensor->view_src ? tensor->view_src : tensor;
65+
return (uintptr_t) base_tensor->data + tensor->view_offs;
66+
}
67+
68+
inline bool ggml_webgpu_tensor_equal(const ggml_tensor * a, const ggml_tensor * b) {
69+
return a->buffer == b->buffer && ggml_webgpu_tensor_addr(a) == ggml_webgpu_tensor_addr(b);
70+
}
71+
72+
inline bool ggml_webgpu_tensor_overlap(const ggml_tensor * a, const ggml_tensor * b) {
73+
return a->buffer == b->buffer && ggml_webgpu_tensor_addr(a) < ggml_webgpu_tensor_addr(b) + ggml_nbytes(b) &&
74+
ggml_webgpu_tensor_addr(b) < ggml_webgpu_tensor_addr(a) + ggml_nbytes(a);
75+
}
76+
6277
struct ggml_webgpu_shader_lib_context {
6378
ggml_tensor * src0;
6479
ggml_tensor * src1;
6580
ggml_tensor * src2;
6681
ggml_tensor * src3;
6782
ggml_tensor * src4;
83+
ggml_tensor * src5;
6884
ggml_tensor * dst;
6985

7086
uint32_t max_wg_size;
7187
size_t wg_mem_limit_bytes = 0;
72-
bool inplace = false;
73-
bool overlap = false;
74-
bool src_overlap = false;
7588
bool supports_subgroups = false;
7689
bool supports_subgroup_matrix = false;
7790
uint32_t sg_mat_m = 0;
@@ -88,6 +101,14 @@ struct webgpu_pipeline {
88101

89102
struct ggml_webgpu_generic_shader_decisions {
90103
uint32_t wg_size = 0;
104+
bool inplace = false;
105+
};
106+
107+
struct ggml_webgpu_binary_shader_decisions {
108+
uint32_t wg_size = 0;
109+
bool inplace = false;
110+
bool overlap = false;
111+
bool src_overlap = false;
91112
};
92113

93114
struct ggml_webgpu_processed_shader {
@@ -102,11 +123,12 @@ struct ggml_webgpu_ssm_conv_shader_decisions {
102123
};
103124

104125
struct ggml_webgpu_ssm_scan_pipeline_key {
105-
int type;
106-
int d_state;
126+
int type;
127+
int d_state;
128+
bool xbc_overlap;
107129

108130
bool operator==(const ggml_webgpu_ssm_scan_pipeline_key & other) const {
109-
return type == other.type && d_state == other.d_state;
131+
return type == other.type && d_state == other.d_state && xbc_overlap == other.xbc_overlap;
110132
}
111133
};
112134

@@ -115,13 +137,15 @@ struct ggml_webgpu_ssm_scan_pipeline_key_hash {
115137
size_t seed = 0;
116138
ggml_webgpu_hash_combine(seed, key.type);
117139
ggml_webgpu_hash_combine(seed, key.d_state);
140+
ggml_webgpu_hash_combine(seed, key.xbc_overlap);
118141
return seed;
119142
}
120143
};
121144

122145
struct ggml_webgpu_ssm_scan_shader_decisions {
123146
uint32_t wg_size;
124147
uint32_t tokens_per_tile;
148+
bool xbc_overlap = false;
125149
};
126150

127151
/** Argsort **/
@@ -242,6 +266,13 @@ struct ggml_webgpu_rms_norm_mul_pipeline_key_hash {
242266
}
243267
};
244268

269+
struct ggml_webgpu_rms_norm_mul_shader_decisions {
270+
uint32_t wg_size = 0;
271+
bool inplace = false;
272+
bool overlap = false;
273+
bool src_overlap = false;
274+
};
275+
245276
/** Pad **/
246277
struct ggml_webgpu_pad_pipeline_key {
247278
bool circular;
@@ -503,11 +534,12 @@ struct ggml_webgpu_flash_attn_pipeline_key_hash {
503534
};
504535

505536
struct ggml_webgpu_flash_attn_decisions {
506-
uint32_t path = GGML_WEBGPU_FLASH_ATTN_PATH_SUBGROUP_MATRIX;
507-
uint32_t q_tile = 0;
508-
uint32_t kv_tile = 0;
509-
uint32_t wg_size = 0;
510-
bool kv_direct = false;
537+
uint32_t path = GGML_WEBGPU_FLASH_ATTN_PATH_SUBGROUP_MATRIX;
538+
uint32_t q_tile = 0;
539+
uint32_t kv_tile = 0;
540+
uint32_t wg_size = 0;
541+
bool kv_direct = false;
542+
bool kv_overlap = false;
511543
};
512544

513545
inline constexpr uint32_t GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH = 4u;
@@ -552,7 +584,7 @@ inline ggml_webgpu_flash_attn_pipeline_key ggml_webgpu_flash_attn_make_pipeline_
552584
key.head_dim_qk = (uint32_t) context.src0->ne[0];
553585
key.head_dim_v = (uint32_t) context.src2->ne[0];
554586
key.kv_direct = kv_direct;
555-
key.kv_overlap = context.src_overlap;
587+
key.kv_overlap = ggml_webgpu_tensor_overlap(context.src1, context.src2);
556588
key.has_mask = has_mask;
557589
key.has_sinks = has_sinks;
558590
key.uses_logit_softcap = ggml_get_op_params_f32(context.dst, 2) != 0.0f;
@@ -1021,7 +1053,7 @@ class ggml_webgpu_shader_lib {
10211053
webgpu_pipeline get_row_norm_pipeline(const ggml_webgpu_shader_lib_context & context) {
10221054
ggml_webgpu_row_norm_pipeline_key key = {};
10231055
key.op = context.dst->op;
1024-
key.inplace = context.inplace;
1056+
key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst);
10251057

10261058
auto it = row_norm_pipelines.find(key);
10271059
if (it != row_norm_pipelines.end()) {
@@ -1051,8 +1083,12 @@ class ggml_webgpu_shader_lib {
10511083
const uint32_t row_norm_wg_size = 128u;
10521084
uint32_t wg_size = std::min(context.max_wg_size, row_norm_wg_size);
10531085
defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size));
1054-
auto processed = preprocessor.preprocess(wgsl_row_norm, defines);
1055-
row_norm_pipelines[key] = ggml_webgpu_create_pipeline(device, processed, variant);
1086+
auto processed = preprocessor.preprocess(wgsl_row_norm, defines);
1087+
auto decisions = std::make_shared<ggml_webgpu_generic_shader_decisions>();
1088+
decisions->wg_size = wg_size;
1089+
decisions->inplace = key.inplace;
1090+
row_norm_pipelines[key] = ggml_webgpu_create_pipeline(device, processed, variant);
1091+
row_norm_pipelines[key].context = decisions;
10561092
return row_norm_pipelines[key];
10571093
}
10581094

@@ -1127,7 +1163,7 @@ class ggml_webgpu_shader_lib {
11271163
webgpu_pipeline get_set_pipeline(const ggml_webgpu_shader_lib_context & context) {
11281164
ggml_webgpu_set_pipeline_key key = {};
11291165
key.type = context.dst->type;
1130-
key.inplace = context.inplace;
1166+
key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst);
11311167

11321168
auto it = set_pipelines.find(key);
11331169
if (it != set_pipelines.end()) {
@@ -1160,6 +1196,7 @@ class ggml_webgpu_shader_lib {
11601196
auto processed = preprocessor.preprocess(wgsl_set, defines);
11611197
auto decisions = std::make_shared<ggml_webgpu_generic_shader_decisions>();
11621198
decisions->wg_size = context.max_wg_size;
1199+
decisions->inplace = key.inplace;
11631200
webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant);
11641201
pipeline.context = decisions;
11651202
set_pipelines[key] = pipeline;
@@ -1355,7 +1392,7 @@ class ggml_webgpu_shader_lib {
13551392

13561393
webgpu_pipeline get_scale_pipeline(const ggml_webgpu_shader_lib_context & context) {
13571394
ggml_webgpu_scale_pipeline_key key = {};
1358-
key.inplace = context.inplace;
1395+
key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst);
13591396

13601397
auto it = scale_pipelines.find(key);
13611398
if (it != scale_pipelines.end()) {
@@ -1375,6 +1412,7 @@ class ggml_webgpu_shader_lib {
13751412
auto processed = preprocessor.preprocess(wgsl_scale, defines);
13761413
auto decisions = std::make_shared<ggml_webgpu_generic_shader_decisions>();
13771414
decisions->wg_size = context.max_wg_size;
1415+
decisions->inplace = key.inplace;
13781416
webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant);
13791417
pipeline.context = decisions;
13801418
scale_pipelines[key] = pipeline;
@@ -1468,6 +1506,8 @@ class ggml_webgpu_shader_lib {
14681506
ggml_webgpu_ssm_scan_pipeline_key key = {};
14691507
key.type = context.dst->type;
14701508
key.d_state = (int) context.src0->ne[0];
1509+
key.xbc_overlap = ggml_webgpu_tensor_overlap(context.src1, context.src4) &&
1510+
ggml_webgpu_tensor_overlap(context.src1, context.src5);
14711511

14721512
auto it = ssm_scan_pipelines.find(key);
14731513
if (it != ssm_scan_pipelines.end()) {
@@ -1499,12 +1539,17 @@ class ggml_webgpu_shader_lib {
14991539
variant += "_wg_reduce";
15001540
}
15011541

1542+
if (key.xbc_overlap) {
1543+
defines.push_back("XBC_OVERLAP");
1544+
}
1545+
15021546
variant += "_d" + std::to_string(key.d_state);
15031547

15041548
auto processed = preprocessor.preprocess(wgsl_ssm_scan, defines);
15051549
auto decisions = std::make_shared<ggml_webgpu_ssm_scan_shader_decisions>();
15061550
decisions->wg_size = wg_size;
15071551
decisions->tokens_per_tile = tokens_per_tile;
1552+
decisions->xbc_overlap = key.xbc_overlap;
15081553
webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant);
15091554
pipeline.context = decisions;
15101555
ssm_scan_pipelines[key] = pipeline;
@@ -1764,11 +1809,9 @@ class ggml_webgpu_shader_lib {
17641809

17651810
uint32_t tile_k;
17661811
if (key.use_subgroup_matrix) {
1767-
tile_k = is_quant ? WEBGPU_MUL_MAT_SUBGROUP_TILE_K_QUANT
1768-
: WEBGPU_MUL_MAT_SUBGROUP_TILE_K_FLOAT;
1812+
tile_k = is_quant ? WEBGPU_MUL_MAT_SUBGROUP_TILE_K_QUANT : WEBGPU_MUL_MAT_SUBGROUP_TILE_K_FLOAT;
17691813
} else {
1770-
tile_k = is_quant ? WEBGPU_MUL_MAT_REG_TILE_K_QUANT
1771-
: WEBGPU_MUL_MAT_REG_TILE_K_FLOAT;
1814+
tile_k = is_quant ? WEBGPU_MUL_MAT_REG_TILE_K_QUANT : WEBGPU_MUL_MAT_REG_TILE_K_FLOAT;
17721815
}
17731816

17741817
// Tiles
@@ -2001,9 +2044,8 @@ class ggml_webgpu_shader_lib {
20012044
defines.push_back("SCALAR");
20022045

20032046
// mul_mat_id is register-tile only.
2004-
const uint32_t tile_k = ggml_is_quantized(context.src0->type)
2005-
? WEBGPU_MUL_MAT_REG_TILE_K_QUANT
2006-
: WEBGPU_MUL_MAT_REG_TILE_K_FLOAT;
2047+
const uint32_t tile_k =
2048+
ggml_is_quantized(context.src0->type) ? WEBGPU_MUL_MAT_REG_TILE_K_QUANT : WEBGPU_MUL_MAT_REG_TILE_K_FLOAT;
20072049

20082050
// Tiles
20092051
defines.push_back("TILE_M=" + std::to_string(WEBGPU_MUL_MAT_TILE_M) + "u");
@@ -2039,8 +2081,8 @@ class ggml_webgpu_shader_lib {
20392081
key.type = context.dst->type;
20402082
key.op = op;
20412083
key.is_unary = is_unary;
2042-
key.inplace = context.inplace;
2043-
key.ttype = (ggml_tri_type) ggml_get_op_params_i32(context.dst, 0);
2084+
key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst) || context.dst->op == GGML_OP_FILL;
2085+
key.ttype = (ggml_tri_type) ggml_get_op_params_i32(context.dst, 0);
20442086

20452087
auto it = unary_pipelines.find(key);
20462088
if (it != unary_pipelines.end()) {
@@ -2098,6 +2140,7 @@ class ggml_webgpu_shader_lib {
20982140
auto processed = preprocessor.preprocess(wgsl_unary, defines);
20992141
auto decisions = std::make_shared<ggml_webgpu_generic_shader_decisions>();
21002142
decisions->wg_size = context.max_wg_size;
2143+
decisions->inplace = key.inplace;
21012144
webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant);
21022145
pipeline.context = decisions;
21032146
unary_pipelines[key] = pipeline;
@@ -2106,9 +2149,9 @@ class ggml_webgpu_shader_lib {
21062149

21072150
webgpu_pipeline get_rms_norm_mul_pipeline(const ggml_webgpu_shader_lib_context & context) {
21082151
ggml_webgpu_rms_norm_mul_pipeline_key key = {};
2109-
key.inplace = context.inplace;
2110-
key.overlap = context.overlap;
2111-
key.src_overlap = context.src_overlap;
2152+
key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst);
2153+
key.overlap = ggml_webgpu_tensor_equal(context.src1, context.dst);
2154+
key.src_overlap = ggml_webgpu_tensor_overlap(context.src0, context.src1);
21122155

21132156
auto it = rms_norm_mul_pipelines.find(key);
21142157
if (it != rms_norm_mul_pipelines.end()) {
@@ -2132,22 +2175,25 @@ class ggml_webgpu_shader_lib {
21322175

21332176
defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size));
21342177

2135-
auto processed = preprocessor.preprocess(wgsl_rms_norm_mul, defines);
2136-
auto decisions = std::make_shared<ggml_webgpu_generic_shader_decisions>();
2137-
decisions->wg_size = context.max_wg_size;
2138-
webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant);
2139-
pipeline.context = decisions;
2140-
rms_norm_mul_pipelines[key] = pipeline;
2178+
auto processed = preprocessor.preprocess(wgsl_rms_norm_mul, defines);
2179+
auto pipeline_decisions = std::make_shared<ggml_webgpu_rms_norm_mul_shader_decisions>();
2180+
pipeline_decisions->wg_size = context.max_wg_size;
2181+
pipeline_decisions->inplace = key.inplace;
2182+
pipeline_decisions->overlap = key.overlap;
2183+
pipeline_decisions->src_overlap = key.src_overlap;
2184+
webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant);
2185+
pipeline.context = pipeline_decisions;
2186+
rms_norm_mul_pipelines[key] = pipeline;
21412187
return rms_norm_mul_pipelines[key];
21422188
}
21432189

21442190
webgpu_pipeline get_binary_pipeline(const ggml_webgpu_shader_lib_context & context) {
21452191
ggml_webgpu_binary_pipeline_key key = {};
21462192
key.type = context.dst->type;
21472193
key.op = context.dst->op;
2148-
key.inplace = context.inplace;
2149-
key.overlap = context.overlap;
2150-
key.src_overlap = context.src_overlap;
2194+
key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst);
2195+
key.overlap = ggml_webgpu_tensor_equal(context.src1, context.dst);
2196+
key.src_overlap = ggml_webgpu_tensor_overlap(context.src0, context.src1);
21512197

21522198
auto it = binary_pipelines.find(key);
21532199
if (it != binary_pipelines.end()) {
@@ -2186,11 +2232,15 @@ class ggml_webgpu_shader_lib {
21862232

21872233
defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size));
21882234

2189-
auto processed = preprocessor.preprocess(wgsl_binary, defines);
2190-
auto decisions = std::make_shared<ggml_webgpu_generic_shader_decisions>();
2191-
decisions->wg_size = context.max_wg_size;
2235+
auto processed = preprocessor.preprocess(wgsl_binary, defines);
2236+
auto pipeline_decisions = std::make_shared<ggml_webgpu_binary_shader_decisions>();
2237+
pipeline_decisions->wg_size = context.max_wg_size;
2238+
pipeline_decisions->inplace = key.inplace;
2239+
pipeline_decisions->overlap = key.overlap;
2240+
pipeline_decisions->src_overlap = key.src_overlap;
2241+
21922242
webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant);
2193-
pipeline.context = decisions;
2243+
pipeline.context = pipeline_decisions;
21942244
binary_pipelines[key] = pipeline;
21952245
return binary_pipelines[key];
21962246
}
@@ -2351,7 +2401,8 @@ class ggml_webgpu_shader_lib {
23512401
defines.push_back(std::string("SG_MAT_K=") + std::to_string(context.sg_mat_k));
23522402
}
23532403

2354-
auto pipeline_decisions = std::make_shared<ggml_webgpu_flash_attn_decisions>(decisions);
2404+
auto pipeline_decisions = std::make_shared<ggml_webgpu_flash_attn_decisions>(decisions);
2405+
pipeline_decisions->kv_overlap = key.kv_overlap;
23552406
defines.push_back(std::string("Q_TILE=") + std::to_string(decisions.q_tile));
23562407
defines.push_back(std::string("KV_TILE=") + std::to_string(decisions.kv_tile));
23572408
defines.push_back(std::string("WG_SIZE=") + std::to_string(decisions.wg_size));
@@ -2543,7 +2594,7 @@ class ggml_webgpu_shader_lib {
25432594
webgpu_pipeline get_rope_pipeline(const ggml_webgpu_shader_lib_context & context) {
25442595
ggml_webgpu_rope_pipeline_key key = {};
25452596
key.type = context.dst->type;
2546-
key.inplace = context.inplace;
2597+
key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst);
25472598
key.has_ff = (context.src2 != nullptr);
25482599

25492600
auto it = rope_pipelines.find(key);
@@ -2582,6 +2633,7 @@ class ggml_webgpu_shader_lib {
25822633
auto processed = preprocessor.preprocess(wgsl_rope, defines);
25832634
auto decisions = std::make_shared<ggml_webgpu_generic_shader_decisions>();
25842635
decisions->wg_size = context.max_wg_size;
2636+
decisions->inplace = key.inplace;
25852637
webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant);
25862638
pipeline.context = decisions;
25872639
rope_pipelines[key] = pipeline;
@@ -2593,7 +2645,7 @@ class ggml_webgpu_shader_lib {
25932645
key.mask_type = context.src1 ? context.src1->type : GGML_TYPE_F32;
25942646
key.has_mask = (context.src1 != nullptr);
25952647
key.has_sink = (context.src2 != nullptr);
2596-
key.inplace = context.inplace;
2648+
key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst);
25972649

25982650
auto it = soft_max_pipelines.find(key);
25992651
if (it != soft_max_pipelines.end()) {
@@ -2634,6 +2686,7 @@ class ggml_webgpu_shader_lib {
26342686
auto processed = preprocessor.preprocess(wgsl_soft_max, defines);
26352687
auto decisions = std::make_shared<ggml_webgpu_generic_shader_decisions>();
26362688
decisions->wg_size = context.max_wg_size;
2689+
decisions->inplace = key.inplace;
26372690
webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant);
26382691
pipeline.context = decisions;
26392692
soft_max_pipelines[key] = pipeline;

0 commit comments

Comments
 (0)