|
| 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_required_extensions("texture3d", DTYPE)} |
| 12 | + |
| 13 | +#define PRECISION ${PRECISION} |
| 14 | + |
| 15 | +#define VEC4_T ${texel_load_type(DTYPE, "texture3d")} |
| 16 | +#define T ${texel_load_component_type(DTYPE, "texture3d")} |
| 17 | + |
| 18 | +${define_active_storage_type("texture3d")} |
| 19 | + |
| 20 | +#extension GL_EXT_control_flow_attributes : require |
| 21 | + |
| 22 | +layout(std430) buffer; |
| 23 | + |
| 24 | +#include "common.glslh" |
| 25 | +#include "indexing.glslh" |
| 26 | + |
| 27 | +${layout_declare_tensor(B, "w", "t_out", DTYPE, "texture3d")} |
| 28 | +${layout_declare_tensor(B, "r", "t_in", DTYPE, "texture3d")} |
| 29 | + |
| 30 | +${layout_declare_ubo(B, "TextureMetadata", "outp")} |
| 31 | +${layout_declare_ubo(B, "TextureMetadata", "inp")} |
| 32 | +${layout_declare_ubo(B, "ivec4", "pad_per_dim")} |
| 33 | +${layout_declare_ubo(B, "float", "fill_value")} |
| 34 | + |
| 35 | +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; |
| 36 | + |
| 37 | +${layout_declare_spec_const(C, "int", "out_layout", "CONTIG_LAYOUT_INT")} |
| 38 | +const int packed_dim = get_packed_dim(out_layout); |
| 39 | + |
| 40 | +void main() { |
| 41 | + const ivec3 out_pos = ivec3(gl_GlobalInvocationID); |
| 42 | + |
| 43 | + if (out_of_bounds(out_pos, outp)) { |
| 44 | + return; |
| 45 | + } |
| 46 | + |
| 47 | + // Convert the thread position to output tensor indices in element space. |
| 48 | + // out_tidx.data[packed_dim] is the element index of the first component in |
| 49 | + // this texel; the remaining three dims are scalar element indices. |
| 50 | + TensorIndex4D out_tidx = texture_pos_to_tensor4d_idx_simple(outp, out_pos); |
| 51 | + |
| 52 | + // Tail texels may have fewer than 4 valid elements; leave extras as 0. |
| 53 | + const int limit = |
| 54 | + min(4, outp.sizes[packed_dim] - out_tidx.data[packed_dim]); |
| 55 | + |
| 56 | + VEC4_T out_texel = VEC4_T(0); |
| 57 | + |
| 58 | + // Process each of the (up to 4) elements in this output texel independently. |
| 59 | + // For each element: subtract pad offsets to obtain the input element index, |
| 60 | + // then copy from the input if in-bounds or write fill_value if in the padding |
| 61 | + // region. |
| 62 | + [[unroll]] for (int comp = 0; comp < limit; comp++) { |
| 63 | + TensorIndex4D in_tidx = out_tidx; |
| 64 | + in_tidx.data[packed_dim] += comp; |
| 65 | + in_tidx.data[0] -= pad_per_dim[0]; |
| 66 | + in_tidx.data[1] -= pad_per_dim[1]; |
| 67 | + in_tidx.data[2] -= pad_per_dim[2]; |
| 68 | + in_tidx.data[3] -= pad_per_dim[3]; |
| 69 | + |
| 70 | + // Signed underflow (output index < pad) produces a negative value that |
| 71 | + // fails the >= 0 check, correctly identifying the padding region. |
| 72 | + if (in_tidx.data[0] >= 0 && in_tidx.data[0] < inp.sizes[0] && |
| 73 | + in_tidx.data[1] >= 0 && in_tidx.data[1] < inp.sizes[1] && |
| 74 | + in_tidx.data[2] >= 0 && in_tidx.data[2] < inp.sizes[2] && |
| 75 | + in_tidx.data[3] >= 0 && in_tidx.data[3] < inp.sizes[3]) { |
| 76 | + TextureElementIndex elem = |
| 77 | + tensor4d_idx_to_texture_element_idx_simple(inp, in_tidx); |
| 78 | + VEC4_T in_texel = texelFetch(t_in, elem.pos, 0); |
| 79 | + out_texel[comp] = T(in_texel[elem.comp]); |
| 80 | + } else { |
| 81 | + out_texel[comp] = T(fill_value); |
| 82 | + } |
| 83 | + } |
| 84 | + |
| 85 | + imageStore(t_out, out_pos, out_texel); |
| 86 | +} |
0 commit comments