diff --git a/backends/vulkan/runtime/graph/ops/glsl/common.glslh b/backends/vulkan/runtime/graph/ops/glsl/common.glslh index f556ba0a705..3c031ec8913 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/common.glslh +++ b/backends/vulkan/runtime/graph/ops/glsl/common.glslh @@ -35,6 +35,34 @@ #define mod_4(x) ((x) & 3) #define mod_8(x) ((x) & 7) +// Branch access avoids Adreno compiler crashes from runtime-indexed UBO vectors. +int safe_idx(const ivec4 v, const int idx) { + if (idx == 0) return v.x; + if (idx == 1) return v.y; + if (idx == 2) return v.z; + return v.w; +} + +uint safe_idx(const uvec4 v, const int idx) { + if (idx == 0) return v.x; + if (idx == 1) return v.y; + if (idx == 2) return v.z; + return v.w; +} + +int safe_idx(const ivec3 v, const int idx) { + if (idx == 0) return v.x; + if (idx == 1) return v.y; + return v.z; +} + +void safe_set(inout ivec4 v, const int idx, const int val) { + if (idx == 0) { v.x = val; } + else if (idx == 1) { v.y = val; } + else if (idx == 2) { v.z = val; } + else { v.w = val; } +} + int sign_extend_8bit(const int val) { if ((val & 0x80) != 0) { diff --git a/backends/vulkan/runtime/graph/ops/glsl/indexing.glslh b/backends/vulkan/runtime/graph/ops/glsl/indexing.glslh index ddcc6d03de2..86f57ee5526 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/indexing.glslh +++ b/backends/vulkan/runtime/graph/ops/glsl/indexing.glslh @@ -80,40 +80,9 @@ bool is_channels_last(const int hashed_layout) { #define get_outer_packed_dim_block_size(layout) (((layout) >> 28) & 0xF) #define get_batch_concat_dim(layout) (((layout) >> 12) & 0xF) -// Safe ivec4 component access via if/else chain. Avoids dynamic vector -// indexing (v[index]) on UBO-backed ivec4 members, which crashes on Adreno 740 -// when the index is a specialization constant with value 1 or 2. -int safe_idx(const ivec4 v, const int idx) { - if (idx == 0) return v.x; - if (idx == 1) return v.y; - if (idx == 2) return v.z; - return v.w; -} - -// Safe uvec4 component access via if/else chain. Same rationale as safe_idx -// for ivec4 — avoids dynamic vector indexing on UBO-backed uvec4 members. -uint safe_idx(const uvec4 v, const int idx) { - if (idx == 0) return v.x; - if (idx == 1) return v.y; - if (idx == 2) return v.z; - return v.w; -} - -// Safe ivec3 component access via if/else chain. Same rationale as safe_idx -// for ivec4. -int safe_idx(const ivec3 v, const int idx) { - if (idx == 0) return v.x; - if (idx == 1) return v.y; - return v.z; -} - -// Safe ivec4 component write via if/else chain. Companion to safe_idx for -// cases where we need to set a component by a spec-const-derived index. -void safe_set(inout ivec4 v, const int idx, const int val) { - if (idx == 0) { v.x = val; } - else if (idx == 1) { v.y = val; } - else if (idx == 2) { v.z = val; } - else { v.w = val; } +uint safe_idx(const uvec4 lower, const uvec4 upper, const int idx) { + if (idx < 4) return safe_idx(lower, idx); + return safe_idx(upper, idx - 4); } // @@ -140,19 +109,19 @@ uint numel(const BufferMetadata meta) { } uint dim_order_at(const BufferMetadata meta, const int dim) { - return meta.dim_order[div_4(dim)][mod_4(dim)]; + return safe_idx(meta.dim_order[0], meta.dim_order[1], dim); } uint dim_order_at(const BufferMetadata meta, const uint dim) { - return meta.dim_order[div_4(dim)][mod_4(dim)]; + return safe_idx(meta.dim_order[0], meta.dim_order[1], int(dim)); } uint stride_at(const BufferMetadata meta, const int dim) { - return meta.strides[div_4(dim)][mod_4(dim)]; + return safe_idx(meta.strides[0], meta.strides[1], dim); } uint stride_at(const BufferMetadata meta, const uint dim) { - return meta.strides[div_4(dim)][mod_4(dim)]; + return safe_idx(meta.strides[0], meta.strides[1], int(dim)); } uint width(const BufferMetadata meta) { @@ -164,11 +133,11 @@ uint height(const BufferMetadata meta) { } uint size_at(const BufferMetadata meta, const int dim) { - return meta.sizes[div_4(dim)][mod_4(dim)]; + return safe_idx(meta.sizes[0], meta.sizes[1], dim); } uint size_at(const BufferMetadata meta, const uint dim) { - return meta.sizes[div_4(dim)][mod_4(dim)]; + return safe_idx(meta.sizes[0], meta.sizes[1], int(dim)); } bool are_equal(const BufferMetadata meta1, const BufferMetadata meta2) { @@ -322,8 +291,8 @@ TensorIndex linear_idx_to_tensor_idx( int dim = int_ndim(meta); int i = 0; for (int d = max(dim - 1, 0); d >= 0; d--) { - uint dim_idx = meta.dim_order[div_4(d)][mod_4(d)]; - uint dim_stride = meta.strides[div_4(dim_idx)][mod_4(dim_idx)]; + uint dim_idx = dim_order_at(meta, d); + uint dim_stride = stride_at(meta, dim_idx); tidx.data[div_4(dim_idx)][mod_4(dim_idx)] = linear_idx / dim_stride; linear_idx = linear_idx % dim_stride; @@ -424,7 +393,7 @@ TensorIndex4D texel_idx_to_tensor4d_idx( uint remaining = block_idx; [[unroll]] for (int d = 3; d >= 0; d--) { int dim_idx = extract_4b(hashed_layout, d); - uint dim_stride = meta.strides[0][dim_idx]; + uint dim_stride = safe_idx(meta.strides[0], dim_idx); tidx.data[dim_idx] = int(remaining / dim_stride); remaining = remaining % dim_stride; } @@ -562,7 +531,7 @@ uint tensor4d_idx_to_linear_idx( const TensorIndex4D tidx) { uint lin_idx = 0; for (int d = 0; d < 4; ++d) { - lin_idx += meta.strides[0][d] * tidx.data[d]; + lin_idx += safe_idx(meta.strides[0], d) * tidx.data[d]; } return lin_idx; } @@ -637,7 +606,7 @@ int tensor4d_idx_to_block_idx( // Compute block-space linear index int block_idx = 0; [[unroll]] for (int d = 0; d < 4; ++d) { - block_idx += int(meta.strides[0][d]) * tidx.data[d]; + block_idx += int(safe_idx(meta.strides[0], d)) * tidx.data[d]; } return block_idx; } @@ -732,10 +701,10 @@ int tensor_idx_to_block_idx( // Compute block-space linear index over all 8 dims int block_idx = 0; [[unroll]] for (int d = 0; d < 4; ++d) { - block_idx += int(meta.strides[0][d]) * int(tidx.data[0][d]); + block_idx += int(safe_idx(meta.strides[0], d)) * int(tidx.data[0][d]); } [[unroll]] for (int d = 0; d < 4; ++d) { - block_idx += int(meta.strides[1][d]) * int(tidx.data[1][d]); + block_idx += int(safe_idx(meta.strides[1], d)) * int(tidx.data[1][d]); } return block_idx; } diff --git a/backends/vulkan/runtime/graph/ops/glsl/indexing_utils.h b/backends/vulkan/runtime/graph/ops/glsl/indexing_utils.h index fdb6f514a3e..cbd2488017b 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/indexing_utils.h +++ b/backends/vulkan/runtime/graph/ops/glsl/indexing_utils.h @@ -9,6 +9,8 @@ #ifndef INDEXING_UTILS_H #define INDEXING_UTILS_H +#include "common.glslh" + /* * The functions defined in this header file use the following shorthand to * represent tensor related data structures. @@ -94,7 +96,7 @@ ivec4 tidx_to_nchwi(const ivec4 tidx, const ivec4 sizes, const int packed_dim) { int base_i = tidx.x * strides.x + tidx.y * strides.y + tidx.z * strides.z + tidx.w * strides.w; - return base_i + ivec4(0, 1, 2, 3) * strides[packed_dim]; + return base_i + ivec4(0, 1, 2, 3) * safe_idx(strides, packed_dim); } /* @@ -135,9 +137,10 @@ int tidx_to_nchwi(const ivec4 tidx, const ivec4 sizes) { ivec4 bufi_to_tidx(int bufi, const ivec4 strides, const ivec4 dim_order) { ivec4 idx; for (int i = 3; i >= 0; i--) { - int dim = dim_order[i]; - idx[dim] = bufi / strides[dim]; - bufi %= strides[dim]; + int dim = safe_idx(dim_order, i); + int dim_stride = safe_idx(strides, dim); + idx[dim] = bufi / dim_stride; + bufi %= dim_stride; } return idx; } @@ -148,8 +151,9 @@ ivec4 bufi_to_tidx(int bufi, const ivec4 strides, const ivec4 dim_order) { ivec4 contiguous_bufi_to_tidx(int bufi, const ivec4 strides) { ivec4 idx; for (int i = 3; i >= 0; i--) { - idx[i] = bufi / strides[i]; - bufi %= strides[i]; + int dim_stride = safe_idx(strides, i); + idx[i] = bufi / dim_stride; + bufi %= dim_stride; } return idx; } @@ -165,15 +169,16 @@ ivec4 lpos_to_tidx( const int batch_inner_dim, const int packed_dim) { // Align packed dim to next multiple of 4 to account for texel padding - sizes[packed_dim] = alignup4(sizes[packed_dim]); + safe_set(sizes, packed_dim, alignup4(safe_idx(sizes, packed_dim))); // Moving 1 texel along the packed dim traverses 4 tensor elements lpos[packed_dim] *= 4; ivec4 tidx = ivec4(lpos, 0); if (sizes.w > 1) { - tidx.w = tidx[batch_inner_dim] / sizes[batch_inner_dim]; - tidx[batch_inner_dim] %= sizes[batch_inner_dim]; + int batch_inner_size = safe_idx(sizes, batch_inner_dim); + tidx.w = tidx[batch_inner_dim] / batch_inner_size; + tidx[batch_inner_dim] %= batch_inner_size; } return tidx; } @@ -184,13 +189,13 @@ ivec3 tidx_to_lpos( const int batch_inner_dim, const int packed_dim) { // Align packed dim to next multiple of 4 to account for texel padding - sizes[packed_dim] = alignup4(sizes[packed_dim]); + safe_set(sizes, packed_dim, alignup4(safe_idx(sizes, packed_dim))); ivec3 lpos = tidx.xyz; // Adjust batch inner dim by batch index if needed if (sizes.w > 1) { - lpos[batch_inner_dim] += tidx.w * sizes[batch_inner_dim]; + lpos[batch_inner_dim] += tidx.w * safe_idx(sizes, batch_inner_dim); } // Fast division by 4, since moving 1 texel along the packed dim traverses 4 // tensor elements. @@ -204,7 +209,7 @@ ivec3 tidx_to_pos( const ivec4 axis_map, const int packed_dim) { // Align packed dim to next multiple of 4 to account for texel padding - sizes[packed_dim] = alignup4(sizes[packed_dim]); + safe_set(sizes, packed_dim, alignup4(safe_idx(sizes, packed_dim))); ivec3 pos; for (int dim = 0; dim < 3; ++dim) { @@ -213,11 +218,13 @@ ivec3 tidx_to_pos( // Adjust batch inner dim by batch index if needed if (sizes.w > 1) { - pos[axis_map[axis_map.w]] += tidx.w * sizes[axis_map.w]; + int batch_inner_dim = axis_map.w; + pos[safe_idx(axis_map, batch_inner_dim)] += + tidx.w * safe_idx(sizes, batch_inner_dim); } // Fast division by 4, since moving 1 texel along the packed dim traverses 4 // tensor elements. - pos[axis_map[packed_dim]] >>= 2; + pos[safe_idx(axis_map, packed_dim)] >>= 2; return pos; } diff --git a/backends/vulkan/runtime/graph/ops/glsl/nchw_to_bitw8_image_nobitw8buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/nchw_to_bitw8_image_nobitw8buffer.glsl index 1a2c257baec..f4bbd064fc8 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/nchw_to_bitw8_image_nobitw8buffer.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/nchw_to_bitw8_image_nobitw8buffer.glsl @@ -58,7 +58,7 @@ ivec4 read_texel(ivec4 tidx) { ivec4 out_tex = ivec4(0); [[unroll]] for (int i = 0; i < 4; ++i) { - if (tidx[packed_dim] + i < sizes[packed_dim]) { + if (tidx[packed_dim] + i < safe_idx(sizes, packed_dim)) { const int in_texel = nchw_in[buf_indices[i] >> 2]; int extracted_val = (in_texel >> (8 * (buf_indices[i] & 3))) & mask; extracted_val = extend_sign(extracted_val); diff --git a/backends/vulkan/runtime/graph/ops/glsl/pad_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/pad_buffer.glsl index 4b8d1558652..22d3adaa2de 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/pad_buffer.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/pad_buffer.glsl @@ -49,7 +49,7 @@ void main() { // value that fails the out_of_bounds check below. TensorIndex in_tidx = out_tidx; [[unroll]] for (int d = 0; d < 4; d++) { - in_tidx.data[0][d] -= uint(pad_per_dim[d]); + in_tidx.data[0][d] -= uint(safe_idx(pad_per_dim, d)); } if (out_of_bounds(in_tidx, inp)) { diff --git a/backends/vulkan/runtime/graph/ops/glsl/reduce.yaml b/backends/vulkan/runtime/graph/ops/glsl/reduce.yaml index 21a7132b8db..c2f1c78dfdd 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/reduce.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/reduce.yaml @@ -18,7 +18,7 @@ reduce: shader_variants: - NAME: sum - NAME: mean - POSTPROCESS: (accum / tin_sizes[reduce_dim]) + POSTPROCESS: (accum / safe_idx(tin_sizes, reduce_dim)) - NAME: amax INIT_ACCUM: first_val UPDATE_ACCUM: max(accum, new_val) diff --git a/backends/vulkan/runtime/graph/ops/glsl/reduce2d.yaml b/backends/vulkan/runtime/graph/ops/glsl/reduce2d.yaml index fdc5eb9f105..967ba449986 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/reduce2d.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/reduce2d.yaml @@ -18,7 +18,7 @@ reduce2d: shader_variants: - NAME: sum2d - NAME: mean2d - POSTPROCESS: (accum / (tin_sizes[reduce_dim1] * tin_sizes[reduce_dim2])) + POSTPROCESS: (accum / (safe_idx(tin_sizes, reduce_dim1) * safe_idx(tin_sizes, reduce_dim2))) - NAME: amax2d INIT_ACCUM: first_val UPDATE_ACCUM: max(accum, new_val) diff --git a/backends/vulkan/runtime/graph/ops/glsl/var_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/var_buffer.glsl index 65e97ea1c9a..758395aa38c 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/var_buffer.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/var_buffer.glsl @@ -54,7 +54,7 @@ void main() { shared_count[tid] = 0; barrier(); - const int R = in_sizes[reduce_dim]; + const int R = safe_idx(in_sizes, reduce_dim); const uint N = gl_WorkGroupSize[reduce_dim]; // Each workgroup processes a contiguous chunk of the input tensor