1/* 2 * Copyright (c) Meta Platforms, Inc. and affiliates. 3 * All rights reserved. 4 * 5 * This source code is licensed under the BSD-style license found in the 6 * LICENSE file in the root directory of this source tree. 7 */ 8 9#version 450 core 10 11#define PRECISION ${PRECISION} 12 13#define VEC4_T ${texel_type(DTYPE)} 14 15layout(std430) buffer; 16 17#include "indexing_utils.h" 18 19${layout_declare_tensor(0, "w", "t_out", DTYPE, STORAGE)} 20${layout_declare_ubo(1, "ivec4", "sizes")} 21${layout_declare_ubo(2, "float", "start")} 22${layout_declare_ubo(3, "float", "step")} 23 24layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; 25 26layout(constant_id = 3) const int packed_dim = C_DIM; 27 28void main() { 29 const ivec3 pos = ivec3(gl_GlobalInvocationID); 30 const ivec4 idx = to_tensor_idx(pos, sizes, packed_dim); 31 32 if (pos_out_of_bounds(pos, sizes, packed_dim)) { 33 return; 34 } 35 36 VEC4_T outtex = VEC4_T(start + pos.x * step, 0, 0, 0); 37 38 imageStore(t_out, pos, outtex); 39} 40