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_active_storage_type(STORAGE)} 14 15layout(std430) buffer; 16 17${layout_declare_tensor(B, "w", "t_out", DTYPE, STORAGE)} 18${layout_declare_tensor(B, "r", "t_in", DTYPE, STORAGE)} 19 20${layout_declare_ubo(B, "ivec3", "range", "ivec3", "src_offset", "ivec3", "dst_offset")} 21 22#include "indexing_utils.h" 23 24layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; 25 26${layout_declare_spec_const(C, "int", "out_layout", "DEFAULT_LAYOUT")} 27const lowp ivec4 out_axis_map = unhash_axis_map(out_layout); 28 29${layout_declare_spec_const(C, "int", "in_layout", "DEFAULT_LAYOUT")} 30const lowp ivec4 in_axis_map = unhash_axis_map(in_layout); 31 32void main() { 33 const ivec3 pos = ivec3(gl_GlobalInvocationID); 34 35 const ivec3 out_pos = pos + dst_offset; 36 const ivec3 in_pos = pos + src_offset; 37 38 if (any(greaterThanEqual(pos, range))) { 39 return; 40 } 41 42 write_texel_lpos( 43 t_out, 44 out_pos, 45 load_texel_lpos(t_in, in_pos, in_axis_map), 46 out_axis_map); 47} 48