|
| 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 | +#define VEC4_T ${texel_load_type(DTYPE, STORAGE)} |
| 13 | + |
| 14 | +${define_active_storage_type(STORAGE)} |
| 15 | + |
| 16 | +#extension GL_EXT_control_flow_attributes : require |
| 17 | + |
| 18 | +layout(std430) buffer; |
| 19 | + |
| 20 | +${layout_declare_tensor(B, "w", "tout", DTYPE, STORAGE)} |
| 21 | +${layout_declare_tensor(B, "r", "tin", DTYPE, STORAGE)} |
| 22 | + |
| 23 | +${layout_declare_ubo(B, "ivec3", "tin_limits")} |
| 24 | +${layout_declare_ubo(B, "ivec4", "tin_sizes")} |
| 25 | + |
| 26 | +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; |
| 27 | + |
| 28 | +layout(constant_id = 3) const int packed_dim = 0; |
| 29 | +layout(constant_id = 4) const int reduce_dim = 0; |
| 30 | +layout(constant_id = 5) const int group_dim = 1; |
| 31 | + |
| 32 | +// A more verbose name would be NWORKERS_PER_GROUP. This describes the number of |
| 33 | +// threads that will co-operate to compute one reduction output. There may be |
| 34 | +// multiple groups computing distinct reduction outputs within one work group. |
| 35 | +#define NWORKERS 4 |
| 36 | + |
| 37 | +// Sets an upper limit on the total size of a work group based on how many |
| 38 | +// elements are allocated in the shared memory array below. Each thread in the |
| 39 | +// work group will write into its assigned element in the shared array. |
| 40 | +#define MAX_NTHREADS 16 |
| 41 | + |
| 42 | + |
| 43 | +shared vec4 shared_vecs[MAX_NTHREADS]; |
| 44 | + |
| 45 | +#include "indexing_utils.h" |
| 46 | + |
| 47 | +int tid_to_smi(const ivec2 tid) { |
| 48 | + return tid.x + tid.y * NWORKERS; |
| 49 | +} |
| 50 | + |
| 51 | +/* |
| 52 | + * The functions below compute reduction along a single dimension for a tensor. |
| 53 | + * The shader template generalize reduction by abstracting the initial value of |
| 54 | + * the accumulator, the calculation used to update the accumulator with new |
| 55 | + * values, and a postprocessing calculation that can be used to modify the |
| 56 | + * accumulator before writing to output. |
| 57 | + * |
| 58 | + * This shader also utilize shared memory to have multiple threads help compute |
| 59 | + * the max and sum reduction operations. A total of NGROUPS x NWORKERS threads |
| 60 | + * are expected to be launched. Each group works on a unique reduction "row", and |
| 61 | + * within a group NWORKERS threads co-operate to compute the max and sum of one |
| 62 | + * "row". Each worker in the group is responsible for computing a partial output |
| 63 | + * of the "row" and uploading it to shared memory; the overall reduction output |
| 64 | + * can then be determined by aggregating the partial outputs stored in shared |
| 65 | + * memory. |
| 66 | + * |
| 67 | + * As a caveat, this shader does not currently support cases where `batch` > 1 |
| 68 | + * and the reduce dim happens to also be the batch concatenation dim. To support |
| 69 | + * this, there will need to be additional logic to set the starting value of |
| 70 | + * `scan_pos[reduce_dim]`. Since this is not expected to be a common use-case, |
| 71 | + * supporting this case is left as an exercise for when it is required. |
| 72 | + */ |
| 73 | + |
| 74 | +// Initializing the accumulator accepts the first value in the reduction row, |
| 75 | +// since some reduction operations (i.e. amax, amin) prefer to initialize with |
| 76 | +// a data point instead of a static value. |
| 77 | +#define INIT_ACCUM(first_val) ${INIT_ACCUM} |
| 78 | +#define UPDATE_ACCUM(accum, new_val) ${UPDATE_ACCUM} |
| 79 | +// Useful for operators such as mean which want to perform a final calculation |
| 80 | +// with the accumulator. |
| 81 | +#define POSTPROCESS(accum) ${POSTPROCESS} |
| 82 | + |
| 83 | +/* |
| 84 | + * Computes reduction where the reduction dim is orthogonal to the packed dim. |
| 85 | + * This case is simpler because each element of a texel belongs to a separate |
| 86 | + * reduction "group", meaning we don't have to perform reduction along a texel. |
| 87 | + */ |
| 88 | +void reduce_nonpacked_dim(const ivec2 tid, ivec3 scan_pos) { |
| 89 | + // shared memory index of this thread |
| 90 | + const int smi = tid_to_smi(tid); |
| 91 | + |
| 92 | + scan_pos[reduce_dim] = 0; |
| 93 | + vec4 accum = INIT_ACCUM(load_texel(tin, scan_pos)); |
| 94 | + |
| 95 | + scan_pos[reduce_dim] = tid.x; |
| 96 | + // Partially accumulate over elements i, i + NWORKERS, i + 2*NWORKERS, ... of |
| 97 | + // the reduction row |
| 98 | + for (int i = tid.x; i < tin_sizes[reduce_dim]; |
| 99 | + i += NWORKERS, scan_pos[reduce_dim] += NWORKERS) { |
| 100 | + accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos)); |
| 101 | + } |
| 102 | + // Write partial output to shared memory and synchronize work group |
| 103 | + shared_vecs[smi] = accum; |
| 104 | + barrier(); |
| 105 | + |
| 106 | + // Since the reduction row is reduced to only one element, only the "main" |
| 107 | + // thread in the group needs aggregate the partial outputs |
| 108 | + if (tid.x == 0) { |
| 109 | + // Iterate over the partial outputs to obtain the overall output |
| 110 | + int group_i = tid.y * NWORKERS; |
| 111 | + accum = shared_vecs[group_i++]; |
| 112 | + for (int i = 1; i < NWORKERS; ++i, group_i++) { |
| 113 | + accum = UPDATE_ACCUM(accum, shared_vecs[group_i]); |
| 114 | + } |
| 115 | + |
| 116 | + // Determine if there are any padding elements in the final texel of the |
| 117 | + // packed dimension |
| 118 | + const int nspill = mod4(tin_sizes[packed_dim]); |
| 119 | + // Detect if this thread is working on the final texels of the packed |
| 120 | + // dimension, which may have padding elements |
| 121 | + const bool is_last_texel = |
| 122 | + scan_pos[packed_dim] == (tin_limits[packed_dim] - 1); |
| 123 | + |
| 124 | + // Explicitly set padding elements to 0 |
| 125 | + if (is_last_texel && nspill > 0) { |
| 126 | + [[unroll]] for (int i = nspill; i < 4; ++i) { |
| 127 | + accum[i] = 0; |
| 128 | + } |
| 129 | + } |
| 130 | + scan_pos[reduce_dim] = tid.x; |
| 131 | + write_texel(tout, scan_pos, POSTPROCESS(accum)); |
| 132 | + } |
| 133 | +} |
| 134 | + |
| 135 | +/* |
| 136 | + * Compute reduction where the reduction dim is also the packed dim. This case is |
| 137 | + * complex because the reduction needs to occur over the individual texels. |
| 138 | + * Therefore, in this algorithm each element of the accumulator texels are |
| 139 | + * themselves partial outputs. Special care has to be taken to ignore padding |
| 140 | + * elements in texels (which occur when the size of the packed dim is not a |
| 141 | + * multiple of 4) so that they do not influence the output of reduction. |
| 142 | + */ |
| 143 | +void reduce_packed_dim(const ivec2 tid, ivec3 scan_pos) { |
| 144 | + // shared memory index of this thread |
| 145 | + const int smi = tid_to_smi(tid); |
| 146 | + |
| 147 | + // Number of non-padding elements in the last texel in the reduction row |
| 148 | + const int nspill = mod4(tin_sizes[packed_dim]); |
| 149 | + // Only reduce up to the last "complete" texel. The last texel will need to be |
| 150 | + // handled specially if it has padding elements. |
| 151 | + const int reduce_len = tin_sizes[packed_dim] - nspill; |
| 152 | + |
| 153 | + scan_pos[reduce_dim] = 0; |
| 154 | + vec4 accum = INIT_ACCUM(vec4(load_texel(tin, scan_pos).x)); |
| 155 | + |
| 156 | + // Partially accumulate over elements i, i + NWORKERS, i + 2*NWORKERS, ... of |
| 157 | + // the reduction row |
| 158 | + scan_pos[reduce_dim] = tid.x; |
| 159 | + for (int i = tid.x * 4; i < reduce_len; |
| 160 | + i += NWORKERS * 4, scan_pos[reduce_dim] += NWORKERS) { |
| 161 | + accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos)); |
| 162 | + } |
| 163 | + // For the last texel in the dim, if there are padding elements then each |
| 164 | + // element of the texel needs to be processed individually such that the |
| 165 | + // padding elements are ignored |
| 166 | + if (scan_pos[reduce_dim] == tin_limits[reduce_dim] - 1 && nspill > 0) { |
| 167 | + const vec4 intex = load_texel(tin, scan_pos); |
| 168 | + for (int i = 0; i < nspill; ++i) { |
| 169 | + accum.x = UPDATE_ACCUM(accum.x, intex[i]); |
| 170 | + } |
| 171 | + } |
| 172 | + // Write partial output to shared memory and synchronize work group |
| 173 | + shared_vecs[smi] = accum; |
| 174 | + barrier(); |
| 175 | + |
| 176 | + // Since the reduction row is reduced to only one element, only the "main" |
| 177 | + // thread in the group needs aggregate the partial outputs |
| 178 | + if (tid.x == 0) { |
| 179 | + // Iterate over the partial maximums to obtain the overall maximum |
| 180 | + int group_i = tid.y * NWORKERS; |
| 181 | + accum = shared_vecs[group_i++]; |
| 182 | + for (int i = 1; i < NWORKERS; ++i, group_i++) { |
| 183 | + accum = UPDATE_ACCUM(accum, shared_vecs[group_i]); |
| 184 | + } |
| 185 | + // Each element of the texel is itself a partial maximum; iterate over the |
| 186 | + // texel to find the actual maximum |
| 187 | + float accum_final = accum.x; |
| 188 | + [[unroll]] for (int i = 1; i < 4; ++i) { |
| 189 | + accum_final = UPDATE_ACCUM(accum[i], accum_final); |
| 190 | + } |
| 191 | + |
| 192 | + scan_pos[reduce_dim] = tid.x; |
| 193 | + write_texel(tout, scan_pos, POSTPROCESS(vec4(accum_final, 0, 0, 0))); |
| 194 | + } |
| 195 | +} |
| 196 | + |
| 197 | +void main() { |
| 198 | + ivec3 scan_pos = ivec3(gl_GlobalInvocationID); |
| 199 | + scan_pos[reduce_dim] = 0; |
| 200 | + |
| 201 | + const ivec2 tid = ivec2( |
| 202 | + gl_LocalInvocationID[reduce_dim], |
| 203 | + gl_LocalInvocationID[group_dim]); |
| 204 | + |
| 205 | + if (any(greaterThanEqual(scan_pos, tin_limits))) { |
| 206 | + return; |
| 207 | + } |
| 208 | + |
| 209 | + if (reduce_dim != packed_dim) { |
| 210 | + reduce_nonpacked_dim(tid, scan_pos); |
| 211 | + } else { |
| 212 | + reduce_packed_dim(tid, scan_pos); |
| 213 | + } |
| 214 | +} |
0 commit comments