Source browser

examples/shaders/lenet_pool.cpp

Read the complete file below or download its original bytes. Line numbers provide stable links to a declaration or processing step.

1// Downsample independent channel planes with two-by-two average pooling.
2#ifdef __cplusplus // Include the runtime only when building the CPU adapter.
3#include "shdr.h" // Share resource and math types with the application.
4#endif // Finish header inclusion for the generated C++ build.
5layout(local_size_x = 64) in; // Give each invocation one pooled feature value.
6layout(std430, binding = 0) readonly buffer Inputs { // Read the previous layer in NCHW order.
7    float input_values[]; // Store the feature planes produced by convolution.
8}; // Finish the input resource declaration.
9layout(std430, // Use tightly packed shader storage for the following port.
10       binding = 1) writeonly buffer Outputs { // Write a distinct downsampled feature layer.
11    float output_values[]; // Preserve the image and channel ordering.
12}; // Finish the output resource declaration.
13uniform uint batch_count; // Receive the number of independent images.
14uniform uint channels; // Receive the number of feature planes per image.
15uniform uint input_width; // Receive the even input-plane width.
16uniform uint input_height; // Receive the even input-plane height.
17void main() { // Average the source square assigned to this output feature.
18    uint width = input_width / 2u, height = input_height / 2u; // Halve the two spatial dimensions.
19    uint plane = width * height, // Count values per output plane before identifying the image.
20         index = gl_GlobalInvocationID.x; // Find this output in the flattened layer.
21    if (index >= batch_count * channels * plane) // Ignore excess invocations beyond the pooled output.
22        return; // Ignore unused invocations after the output layer.
23    uint x = index % width, y = (index % plane) / width; // Recover this pooled pixel's coordinates.
24    uint source = (index / plane) * input_width * input_height + 2u * y * input_width + // Locate the upper-left source pixel of this two-by-two pooling tile.
25                  2u * x; // Select the top-left source pixel in its channel plane.
26    float value = // Begin pooling with the two pixels in the upper source row.
27        input_values[source] + input_values[source + 1u]; // Add the upper two source pixels.
28    value += input_values[source + input_width] + // Add the two pixels in the second row of the pooling tile.
29             input_values[source + input_width + 1u]; // Add the lower two source pixels.
30    output_values[index] = value * 0.25; // Divide the complete sum by the four-pixel area.
31} // Finish this average-pooling invocation.