Source browser

examples/shaders/ann_infer.cpp

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

1// A fully connected 2 -> ReLU(8) -> sigmoid(1) classifier, one invocation per row.
2#ifdef __cplusplus // Expose host declarations only when a C++ compiler reads this shader.
3#include "shdr.h" // Include shdr.h declarations.
4#endif // Finish this conditional compilation branch.
5layout(local_size_x = 64) in; // Group 64 independent invocations along x.
6layout(std430, binding = 0) readonly buffer Inputs { // Read the model's sample-major input features through buffer binding 0.
7    float inputs[]; // Store the model's sample-major input features.
8}; // Finish this declaration.
9layout(std430, binding = 1) readonly buffer Parameters { // Read the complete flat learned parameter vector through buffer binding 1.
10    float parameters[]; // Store the complete flat learned parameter vector.
11}; // Finish this declaration.
12layout(std430, binding = 2) writeonly buffer Predictions { // Write one predicted class probability per sample through buffer binding 2.
13    float predictions[]; // Store one predicted class probability per sample.
14}; // Finish this declaration.
15uniform uint batch_count; // Set the number of independent samples in this model batch.
16float logistic(float value) { // Evaluate the sigmoid with a sign-selected, numerically stable exponential.
17    if (value >= 0.0) // Use the negative exponent when the input logit is nonnegative.
18        return 1.0 / (1.0 + exp(-value)); // Convert a nonnegative logit into its sigmoid probability.
19    float e = exp(value); // Use a nonpositive exponent for the negative logit branch.
20    return e / (1.0 + e); // Convert a negative logit into its sigmoid probability.
21} // End this block.
22void main() { // Evaluate the fully connected binary classifier for one input row.
23    uint row = gl_GlobalInvocationID.x; // Select the independent input sample assigned to this invocation.
24    if (row >= batch_count) // Ignore invocations outside the sample batch.
25        return; // Stop this invocation before performing another resource access.
26    float x = inputs[2u * row], y = inputs[2u * row + 1u], logit = parameters[32]; // Read the two input features and initialize the output with its bias.
27    for (uint neuron = 0u; neuron < 8u; ++neuron) { // Evaluate each of the eight hidden neurons.
28        float activation = max(0.0, // Apply ReLU to this hidden neuron's weighted input and bias.
29                               parameters[2u * neuron] * x + parameters[2u * neuron + 1u] * y + // Multiply the two input features by this neuron's learned weights.
30                                   parameters[16u + neuron]); // Include this hidden neuron's bias before the ReLU.
31        logit += activation * parameters[24u + neuron]; // Add this hidden activation through its learned output weight.
32    } // End this block.
33    predictions[row] = logistic(logit); // Convert the completed output logit into a binary class probability.
34} // End this block.