File tree Expand file tree Collapse file tree 1 file changed +7
-3
lines changed
ggml/src/ggml-vulkan/vulkan-shaders Expand file tree Collapse file tree 1 file changed +7
-3
lines changed Original file line number Diff line number Diff line change @@ -801,7 +801,7 @@ void main() {
801
801
}
802
802
#else
803
803
const uint row_i = ic * BN + loadc_b + l;
804
- if (row_i < _ne1) {
804
+ if (row_i < _ne1 && block + loadr_b < end_k ) {
805
805
const u16vec2 row_idx = row_ids[row_i];
806
806
buf_b[(loadc_b + l) * SHMEM_STRIDE + loadr_b] = TO_FLOAT_TYPE(data_b[pos_b + row_idx.y * p.batch_stride_b + (row_idx.x % p.ne11) * p.stride_b + loadr_b]);
807
807
} else {
@@ -875,7 +875,9 @@ void main() {
875
875
876
876
const u16vec2 row_idx = row_ids[row_i];
877
877
878
- data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr + cm_row * TM + store_r] = D_TYPE(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]);
878
+ if (dr + cm_row * TM + store_r < p.M) {
879
+ data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr + cm_row * TM + store_r] = D_TYPE(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]);
880
+ }
879
881
}
880
882
}
881
883
}
@@ -925,7 +927,9 @@ void main() {
925
927
#endif // MUL_MAT_ID
926
928
[[unroll]] for (uint cr = 0; cr < TM; cr++) {
927
929
#ifdef MUL_MAT_ID
928
- data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + cr] = D_TYPE(sums[(wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr]);
930
+ if (dr_warp + cr < p.M) {
931
+ data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + cr] = D_TYPE(sums[(wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr]);
932
+ }
929
933
#else
930
934
if (dr_warp + cr < p.M && dc_warp + cc < p.N) {
931
935
data_d[offsets + (dc_warp + cc) * p.stride_d + dr_warp + cr] = D_TYPE(sums[(wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr]);
You can’t perform that action at this time.
0 commit comments