Source browser

examples/shaders/lenet_conv.cpp

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

1// Apply a valid convolution with optional zero padding and ReLU.
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 output feature value.
6layout(std430, binding = 0) readonly buffer Inputs { // Read the previous layer in NCHW order.
7    float input_values[]; // Store complete channel planes inside each image.
8}; // Finish the input resource declaration.
9layout(std430, binding = 1) readonly buffer Parameters { // Read flattened model weights and biases.
10    float parameters[]; // Match the offline model parameter order.
11}; // Finish the parameter resource declaration.
12layout(std430, binding = 2) // Use the third tightly packed shader storage binding.
13    writeonly buffer Outputs { // Store each next-layer feature in its own output position.
14    float output_values[]; // Keep the next layer in the same NCHW order.
15}; // Finish the output resource declaration.
16uniform uint batch_count; // Receive the number of independent images.
17uniform uint input_width; // Receive the width of each input channel plane.
18uniform uint input_height; // Receive the height of each input channel plane.
19uniform uint input_channels; // Receive the number of input feature planes.
20uniform uint output_channels; // Receive the number of learned convolution filters.
21uniform uint kernel_size; // Receive the side length of each square filter.
22uniform uint padding; // Receive the number of implicit zero pixels on each edge.
23uniform uint weight_offset; // Receive the first convolution weight in the flat model.
24uniform uint bias_offset; // Receive the first convolution bias in the flat model.
25void main() { // Compute the output feature value assigned to this invocation.
26    uint width = // Compute the output width after kernel extent and zero padding.
27        input_width + 2u * padding - kernel_size + 1u; // Find the valid padded output width.
28    uint height = // Compute the output height after kernel extent and zero padding.
29        input_height + 2u * padding - kernel_size + 1u; // Find the valid padded output height.
30    uint plane = width * height, // Count values per output plane before identifying the image.
31         index = gl_GlobalInvocationID.x; // Find this output within the flattened layer.
32    if (index >= batch_count * output_channels * plane) // Ignore excess invocations beyond the final convolution output.
33        return; // Ignore unused invocations after the layer.
34    uint row = index / (output_channels * plane); // Find the independent image in this batch.
35    uint channel = // Select the learned output feature plane.
36        (index / plane) % output_channels; // Find the learned filter producing this feature.
37    uint x = index % width, // Decode the horizontal and vertical output coordinates.
38         y = (index % plane) / width; // Recover the pixel coordinates in the output plane.
39    float value = // Begin convolution accumulation with the learned channel bias.
40        parameters[bias_offset + channel]; // Start the feature value with its learned bias.
41    for (uint source = 0u; source < input_channels; // Visit every input feature plane connected to this output.
42         ++source) { // Accumulate contributions from every input plane.
43        for (uint ky = 0u; ky < kernel_size; ++ky) { // Visit each row of the convolution kernel.
44            for (uint kx = 0u; kx < kernel_size; // Visit every horizontal position in the learned kernel.
45                 ++kx) { // Visit each column of the convolution kernel.
46                int sx = int(x + kx) - // Move the kernel position into signed padded input coordinates.
47                         int(padding); // Place the kernel column inside the unpadded input.
48                int sy = // Move the vertical kernel position into signed padded input coordinates.
49                    int(y + ky) - int(padding); // Place the kernel row inside the unpadded input.
50                if (sx < 0 || sy < 0 || sx >= int(input_width) || sy >= int(input_height)) // Treat positions outside the original feature plane as zero.
51                    continue; // Use zero outside the input image.
52                uint sample_index = (row * input_channels + source) * input_width * input_height + // Locate this image and source channel in NCHW storage.
53                                    uint(sy) * input_width + // Move to the selected row of the source feature plane.
54                                    uint(sx); // Select the NCHW input feature.
55                uint coefficient = // Locate this learned kernel coefficient.
56                    weight_offset + // Begin at the first weight of this convolution layer.
57                    ((channel * input_channels + source) * kernel_size + ky) * kernel_size + // Select the output channel, input channel and kernel row.
58                    kx; // Select the matching learned kernel coefficient.
59                value += parameters[coefficient] * // Multiply the learned coefficient by the selected source feature.
60                         input_values[sample_index]; // Accumulate this convolution contribution.
61            } // Finish this kernel row.
62        } // Finish this source plane's kernel.
63    } // Finish the sum over input feature planes.
64    output_values[index] = // Store the completed feature with the chosen activation.
65        max(value, 0.0); // Apply ReLU once after the complete learned projection.
66} // Finish this convolution invocation.