Source browser
scripts/generate_transformer.py
Read the complete file below or download its original bytes. Line numbers provide stable links to a declaration or processing step.
File fingerprint
1#!/usr/bin/env python3
2"""Generate standalone native GLSL inference and analytic-gradient examples.
3
4Both emitted shaders contain the same four-token transformer function and are
5independently compilable. Use --check to compare them with their shared template.
6"""
7import argparse # Read command-line options and show help.
8from pathlib import Path # Build paths and read or write files.
9
10ROOT = Path(__file__).resolve().parents[1] # Find the checkout containing the maintained source files.
11# Shared forward-mode derivative functions.
12BODY = r"""
13
14// Each vec2 carries a value and one analytic parameter derivative.
15vec2 ad_mul(vec2 a, vec2 b) { // Apply the product rule to values and their partial derivatives.
16 return vec2(a.x * b.x, a.y * b.x + a.x * b.y); // Keep the value product and both derivative contributions.
17} // End the product-rule helper.
18vec2 ad_div(vec2 a, vec2 b) { // Apply the quotient rule to one partial derivative.
19 return vec2(a.x / b.x, (a.y * b.x - a.x * b.y) / (b.x * b.x)); // Divide values and propagate the numerator and denominator derivatives.
20} // End the quotient-rule helper.
21vec2 ad_exp(vec2 a) { // Differentiate the exponential using its computed value.
22 float value = exp(a.x); // Compute the scalar value once before propagating its derivative.
23 return vec2(value, value * a.y); // Multiply the input derivative by the exponential value.
24} // End the exponential helper.
25vec2 ad_log(vec2 a) { // Differentiate the natural logarithm.
26 return vec2(log(a.x), a.y / a.x); // Divide the input derivative by its value.
27} // End the logarithm helper.
28vec2 ad_invsqrt(vec2 a) { // Differentiate inverse square root for the positive normalization input.
29 float value = inversesqrt(a.x); // Compute the scalar value once before propagating its derivative.
30 return vec2(value, -0.5 * value * value * value * a.y); // Use the derivative of a power with exponent minus one half.
31} // End the inverse-square-root helper.
32vec2 ad_relu(vec2 a) { // Keep the derivative only where the activation is positive.
33 return a.x > 0.0 ? a : vec2(0.0); // Use a zero value and derivative on the inactive branch.
34} // End the ReLU helper.
35vec2 weight(uint index, uint seed) { // Read a parameter and seed its selected partial derivative.
36 return vec2(parameters[index], index == seed ? 1.0 : 0.0); // Seed one for this parameter and zero for every other parameter.
37} // End the parameter derivative seed helper.
38
39// Parameter ranges: position0:16,Q16:32,K32:48,V48:64,O64:80,
40// RMSscale80:84,FFN1weight84:100,bias100:104,FFN2weight104:120,
41// bias120:124,classifierweight124:132,bias132:134.
42// The returned vec4 stores two logits followed by their two derivatives.
43vec4 transformer(uint row, uint seed) { // Evaluate two classifier logits and one parameter derivative.
44 vec2 x[16]; // Hold four tokens after adding their learned positional features.
45 vec2 q[16]; // Hold the four query vectors and their derivatives.
46 vec2 k[16]; // Hold the four key vectors and their derivatives.
47 vec2 v[16]; // Hold the four value vectors and their derivatives.
48 vec2 attended[16]; // Hold attention output after its projection and residual.
49 vec2 normalized[16]; // Hold each normalized feature and its derivative.
50 vec2 features[16]; // Hold the feed-forward output after its residual.
51 for (uint i = 0u; i < 16u; ++i) // Add learned positions to all sixteen input features.
52 x[i] = vec2(inputs[row * 16u + i], 0.0) + weight(i, seed); // Add the learned position to this input value.
53 for (uint token = 0u; token < 4u; ++token) // Project each of the four tokens into query, key and value features.
54 for (uint c = 0u; c < 4u; ++c) { // Accumulate query, key and value matrix products.
55 vec2 qvalue = vec2(0.0); // Start the query dot product at zero.
56 vec2 kvalue = vec2(0.0); // Start the key dot product at zero.
57 vec2 vvalue = vec2(0.0); // Start the value dot product at zero.
58 for (uint j = 0u; j < 4u; ++j) { // Accumulate query, key and value matrix products.
59 qvalue += ad_mul(x[token * 4u + j], weight(16u + j * 4u + c, seed)); // Accumulate the query projection and its derivative.
60 kvalue += ad_mul(x[token * 4u + j], weight(32u + j * 4u + c, seed)); // Accumulate the key projection and its derivative.
61 vvalue += ad_mul(x[token * 4u + j], weight(48u + j * 4u + c, seed)); // Accumulate the value projection and its derivative.
62 } // The query, key and value dot products are complete.
63 q[token * 4u + c] = qvalue; // Store the projected query feature.
64 k[token * 4u + c] = kvalue; // Store the projected key feature.
65 v[token * 4u + c] = vvalue; // Store the projected value feature.
66 } // All projected token features are now stored.
67 for (uint token = 0u; token < 4u; ++token) { // Evaluate this token's attention and feed-forward features.
68 vec2 scores[4]; // Keep one query-key score for each source token.
69 vec2 probabilities[4]; // Keep exponentiated attention scores before normalization.
70 float maximum = -1000000000.0; // Find a score shift that prevents exponential overflow.
71 for (uint source = 0u; source < 4u; ++source) { // Create an attention score for every source token.
72 scores[source] = vec2(0.0); // Start this scaled query-key dot product at zero.
73 for (uint c = 0u; c < 4u; ++c) // Sum the query-key product across four feature channels.
74 scores[source] += ad_mul(q[token * 4u + c], k[source * 4u + c]) * 0.5; // Accumulate the scaled query-key dot product.
75 maximum = max(maximum, scores[source].x); // Retain the largest attention score for stable softmax.
76 } // Every source token now has an attention score.
77 vec2 total = vec2(0.0); // Accumulate the denominator of the softmax.
78 for (uint source = 0u; source < 4u; ++source) { // Exponentiate shifted scores for the softmax denominator.
79 probabilities[source] = ad_exp(scores[source] - vec2(maximum, 0.0)); // Keep the shifted exponential and its derivative.
80 total += probabilities[source]; // Sum the exponentials for probability normalization.
81 } // The softmax denominator now includes every source score.
82 vec2 context[4]; // Combine source value vectors using their attention probabilities.
83 for (uint c = 0u; c < 4u; ++c) { // Build each channel of the attention value mixture.
84 context[c] = vec2(0.0); // Start this output feature mixture at zero.
85 for (uint source = 0u; source < 4u; ++source) // Mix source value vectors with their normalized attention probabilities.
86 context[c] += ad_mul(ad_div(probabilities[source], total), v[source * 4u + c]); // Weight source values by their normalized attention probabilities.
87 } // The attention mixture now contains four feature values.
88 for (uint c = 0u; c < 4u; ++c) { // Project attention output and preserve the input residual.
89 attended[token * 4u + c] = x[token * 4u + c]; // Start with the input token as a residual connection.
90 for (uint j = 0u; j < 4u; ++j) // Apply the learned attention output matrix to this feature.
91 attended[token * 4u + c] += ad_mul(context[j], weight(64u + j * 4u + c, seed)); // Add the learned attention output projection.
92 } // The attention projection and residual are complete.
93 vec2 square = vec2(0.00001, 0.0); // Start the mean squared feature value with a small positive epsilon.
94 for (uint c = 0u; c < 4u; ++c) // Accumulate the mean squared feature value for RMS normalization.
95 square += ad_mul(attended[token * 4u + c], attended[token * 4u + c]) * 0.25; // Add one quarter of this squared feature to the RMS estimate.
96 vec2 inverse = ad_invsqrt(square); // Compute the inverse RMS value and its derivative.
97 for (uint c = 0u; c < 4u; ++c) // Apply inverse RMS and the learned channel scale.
98 normalized[token * 4u + c] = // Scale the normalized feature with its learned RMS parameter.
99 ad_mul(ad_mul(attended[token * 4u + c], inverse), weight(80u + c, seed)); // Propagate both the inverse RMS and learned scale derivatives.
100 vec2 hidden[4]; // Hold four feed-forward hidden activations.
101 for (uint c = 0u; c < 4u; ++c) { // Evaluate each feed-forward hidden feature.
102 hidden[c] = weight(100u + c, seed); // Start the hidden neuron with its learned bias.
103 for (uint j = 0u; j < 4u; ++j) // Apply the first feed-forward matrix to normalized features.
104 hidden[c] += ad_mul(normalized[token * 4u + j], weight(84u + j * 4u + c, seed)); // Accumulate normalized input through the first feed-forward matrix.
105 hidden[c] = ad_relu(hidden[c]); // Apply the ReLU value and derivative together.
106 } // The four hidden features now include ReLU derivatives.
107 for (uint c = 0u; c < 4u; ++c) { // Start each feed-forward output with its residual and bias.
108 features[token * 4u + c] = normalized[token * 4u + c] + weight(120u + c, seed); // Start the residual output with its bias and normalized input.
109 for (uint j = 0u; j < 4u; ++j) // Project the hidden activations through the second feed-forward matrix.
110 features[token * 4u + c] += ad_mul(hidden[j], weight(104u + j * 4u + c, seed)); // Add the second feed-forward matrix projection.
111 } // The second feed-forward projection and residual are complete.
112 } // This token now has its final learned features.
113 vec2 pooled[4]; // Average token features before classification.
114 for (uint c = 0u; c < 4u; ++c) { // Mean-pool each feature across the four tokens.
115 pooled[c] = vec2(0.0); // Start this pooled feature at zero.
116 for (uint token = 0u; token < 4u; ++token) // Average this feature across the four token outputs.
117 pooled[c] += features[token * 4u + c] * 0.25; // Add one quarter of this token feature.
118 } // All token features have been mean-pooled.
119 vec2 logits[2]; // Hold the two class scores and their derivatives.
120 for (uint c = 0u; c < 2u; ++c) { // Build each of the two classifier scores.
121 logits[c] = weight(132u + c, seed); // Start the classifier score with its learned bias.
122 for (uint j = 0u; j < 4u; ++j) // Apply the classifier matrix to the pooled feature values.
123 logits[c] += ad_mul(pooled[j], weight(124u + j * 2u + c, seed)); // Accumulate pooled features through the classifier matrix.
124 } // Both class scores and their partials are complete.
125 return vec4(logits[0].x, logits[1].x, logits[0].y, logits[1].y); // Return both logits followed by their two partial derivatives.
126} // The shared transformer forward pass is complete.
127
128// Stable two-class cross entropy. The maximum shift cancels analytically.
129vec2 cross_entropy(vec4 logits, uint label) { // Evaluate the selected class loss and its derivative.
130 float shift = max(logits.x, logits.y); // Use the larger logit as a safe exponential shift.
131 vec2 first = vec2(logits.x, logits.z); // Keep the first class value or its shifted exponential.
132 vec2 second = vec2(logits.y, logits.w); // Keep the second class value or its shifted exponential.
133 vec2 total = ad_exp(first - vec2(shift, 0.0)) + ad_exp(second - vec2(shift, 0.0)); // Accumulate the denominator of the softmax.
134 return ad_log(total) + vec2(shift, 0.0) - (label == 0u ? first : second); // Subtract the selected class score from the stable log-sum-exp.
135} // End the stable cross-entropy helper.
136"""
137
138# Common model input declarations.
139PREAMBLE = r"""
140// Generated by scripts/generate_transformer.py. Edit its shared function body.
141// Four-token attention, RMS normalization, feed-forward residual and classifier.
142#ifdef __cplusplus // Include the C++ runtime only for CPU compilation.
143#include "shdr.h" // Use the portable math types and kernel registration API.
144#endif // Keep GPU compilation free of C++ header contents.
145layout(local_size_x = 64) in; // Process independent sample or parameter indices in groups of 64.
146layout(std430, binding = 0) readonly buffer Inputs { // Declare the readonly model resource named by this block.
147 float inputs[]; // Read sixteen feature values for each four-token sample.
148}; // Close the input feature storage block.
149layout(std430, binding = 1) readonly buffer Parameters { // Declare the readonly model resource named by this block.
150 float parameters[]; // Read all 134 trainable parameter values.
151}; // Close the model parameter storage block.
152layout(std430, binding = 2) readonly buffer Labels { // Declare the readonly model resource named by this block.
153 uint labels[]; // Read the binary target class for each sequence.
154}; // Close the target label storage block.
155uniform uint batch_count; // Limit every dispatch to the declared number of sequences.
156"""
157
158SUFFIXES = { # Give inference and differentiation separate entry points.
159 'transformer_infer.cpp': r"""
160
161void main() { // Write the result belonging to this compute invocation.
162 uint row = gl_GlobalInvocationID.x; // Locate this sample in the batch.
163 if (row >= batch_count) // Ignore invocations outside the allocated batch.
164 return; // Stop padding invocations before processing a sample.
165 vec4 logits = transformer(row, 4294967295u); // Hold the two class scores and their derivatives.
166 float shift = max(logits.x, logits.y); // Use the larger logit as a safe exponential shift.
167 float first = exp(logits.x - shift); // Keep the first class value or its shifted exponential.
168 float second = exp(logits.y - shift); // Keep the second class value or its shifted exponential.
169 float total = first + second; // Accumulate the denominator of the softmax.
170 predictions[row] = vec4(first / total, // Publish both probabilities, the loss and the predicted class.
171 second / total, // Publish the probability of class one.
172 cross_entropy(logits, labels[row]).x, // Publish the selected class loss.
173 logits.y > logits.x ? 1.0 : 0.0); // Select the class with the larger score.
174} // Finish this prediction invocation.
175""",
176 'transformer_gradient.cpp': r"""
177
178void main() { // Write the result belonging to this compute invocation.
179 uint index = gl_GlobalInvocationID.x; // Locate this sample-and-parameter gradient pair.
180 uint row = index / 134u; // Locate this sample in the batch.
181 uint parameter = index % 134u; // Select one of the 134 trainable parameters.
182 if (row >= batch_count) // Ignore invocations outside the allocated batch.
183 return; // Stop padding invocations before processing a sample.
184 // One invocation evaluates one exact analytic partial derivative.
185 gradients[index] = cross_entropy(transformer(row, parameter), labels[row]).y; // Publish this exact analytic sample-and-parameter derivative.
186} // Finish this parameter gradient invocation.
187""",
188} # Both shaders use the same forward and loss functions.
189
190
191def main(): # Read options and run this command-line tool.
192 parser = argparse.ArgumentParser(description=__doc__) # Describe the command and its supported options.
193 parser.add_argument("--check", action="store_true") # Compare generated files in check mode.
194 arguments = parser.parse_args() # Read whether generated files should be checked or written.
195 stale = [] # Collect generated files that need regeneration.
196 for name, suffix in SUFFIXES.items(): # Generate both inference and gradient shader files.
197 if name.endswith("infer.cpp"): # Give inference one output vector per sample.
198 port = r"""
199layout(std430, binding = 3) writeonly buffer Predictions { // Declare the inference output.
200 vec4 predictions[]; // Store class probabilities, loss and the selected class.
201}; // Close the prediction output block.
202"""
203 else: # Give differentiation one output scalar per sample and parameter.
204 port = r"""
205layout(std430, binding = 3) writeonly buffer Gradients { // Declare the analytic derivative output.
206 float gradients[]; // Store one partial derivative for each sample and parameter.
207}; // Close the gradient output block.
208"""
209 source = PREAMBLE + port + BODY + suffix # Join declarations, the shared forward body and the selected entry point.
210 destination = ROOT / "examples" / "shaders" / name # Locate the generated shader file.
211 if arguments.check: # Compare generated bytes with the existing source files.
212 if not destination.exists() or destination.read_text() != source: # Detect a missing generated file or a byte difference.
213 stale.append(name) # Remember this shader needs regeneration.
214 else: # Write the generated files in output mode.
215 destination.write_text(source) # Update this shader from the shared transformer template.
216 if stale: # Report generated files that differ from their source template.
217 raise SystemExit("generated transformer shaders are stale: " + ", ".join(stale)) # List the shader files that need regeneration.
218
219
220if __name__ == "__main__": # Run only when this file is invoked as a command.
221 main() # Read options and run the requested checks.