Source browser
examples/transformer_support.hpp
Read the complete file below or download its original bytes. Line numbers provide stable links to a declaration or processing step.
File fingerprint
1/** @file transformer_support.hpp
2 * @brief Host wiring for a four-token transformer with analytic SGD training.
3 * The same native shader functions perform CPU and GPU inference and evaluate
4 * every parameter derivative. Gradient dispatches read the current parameters.
5 * Subsequent SGD dispatches apply their averaged updates.
6 */
7#ifndef SHDR_EXAMPLES_TRANSFORMER_SUPPORT_HPP // Include this example helper only once.
8#define SHDR_EXAMPLES_TRANSFORMER_SUPPORT_HPP // Remember that this example helper is available.
9#include "shdr.h" // Use the same header for CPU and GPU resource wiring.
10#include <cmath> // Use ordinary scalar math in host designers and oracles.
11#include <random> // Create repeatable random fixtures and initial parameters.
12
13namespace shdr_examples { // Keep these example helpers in their own namespace.
14
15/** @brief A tiny single-head sequence classifier with all 134 parameters learned.
16 * Architecture: four 4-component tokens, learned positions, scaled attention,
17 * output residual, RMSNorm, a 4-wide ReLU feed-forward residual, mean pooling,
18 * and a two-class softmax classifier. This compact implementation demonstrates
19 * training mechanics and explicit CPU/GPU resource wiring.
20 */
21class transformer_model { // Bundle all transformer resources and linked stages.
22 static unsigned checked_batch(unsigned count) { // Validate the batch before allocating model resources.
23 if (!count || count > 65536u) // Handle this validation or boundary case explicitly.
24 throw shdr::error("transformer batch count must be between 1 and 65536"); // Report the invalid design or failed check clearly.
25 return count; // Return the validated sequence count.
26 } // Finish the checked_batch operation.
27 unsigned batch_count_; // Remember how many token sequences share this model.
28 std::unique_ptr<shdr::kernel> infer_, gradient_, update_; // Own the three compiled transformer stages.
29 shdr::pipeline inference_, training_; // Keep inference and training as separate linked pipelines.
30
31 public: // Expose the resources and processing helpers.
32 enum { token_count = 4, width = 4, parameter_count = 134 }; // Fix the four-token, four-feature model and its parameter count.
33 shdr::buffer<float> inputs, parameters, gradients; // Keep input sequences, trainable values and per-sample gradients.
34 shdr::buffer<unsigned> labels; // Keep each sequence class in a portable unsigned buffer.
35 shdr::buffer<shdr::glsl::vec4> predictions; // Keep probabilities, loss and predicted class for each sequence.
36
37 /** @brief Allocate and link inference, analytic gradients and averaged SGD.
38 * @param batch_count Positive number of four-token training sequences.
39 * @param learning_rate SGD step size applied to the batch-mean gradient.
40 */
41 explicit transformer_model(unsigned batch_count = 8, float learning_rate = 0.15f) // Connect the transformer inference and training stages.
42 : batch_count_(checked_batch(batch_count)), // Validate the batch size before allocating model resources.
43 infer_(shdr::make_kernel("transformer_infer_kernel")), // Create the next compiled inference, derivative or update stage.
44 gradient_(shdr::make_kernel("transformer_gradient_kernel")), // Create the next compiled inference, derivative or update stage.
45 update_(shdr::make_kernel("nn_sgd_kernel")), inputs(batch_count_ * token_count * width), // Create the next compiled inference, derivative or update stage.
46 parameters(parameter_count), gradients(batch_count_ * parameter_count), // Allocate every learned transformer value.
47 labels(batch_count_), predictions(batch_count_) { // Allocate one target class for each sequence.
48 if (!std::isfinite(learning_rate) || learning_rate <= 0.0f) // Reject a nonfinite or out-of-range value.
49 throw shdr::error("transformer learning rate must be positive and finite"); // Report the invalid design or failed check clearly.
50 infer_->bind("inputs", inputs) // Connect the next model stage to its shared resources.
51 .bind("parameters", parameters) // Connect the next named resource to the same shader instance.
52 .bind("labels", labels) // Connect the next named resource to the same shader instance.
53 .bind("predictions", predictions) // Connect the next named resource to the same shader instance.
54 .set<unsigned>("batch_count", batch_count_); // Set the shader dimensions or scalar processing parameter.
55 gradient_->bind("inputs", inputs) // Connect the next model stage to its shared resources.
56 .bind("parameters", parameters) // Connect the next named resource to the same shader instance.
57 .bind("labels", labels) // Connect the next named resource to the same shader instance.
58 .bind("gradients", gradients) // Connect the next named resource to the same shader instance.
59 .set<unsigned>("batch_count", batch_count_); // Set the shader dimensions or scalar processing parameter.
60 update_->bind("parameters", parameters) // Connect the next model stage to its shared resources.
61 .bind("gradients", gradients) // Connect the next named resource to the same shader instance.
62 .set<unsigned>("batch_count", batch_count_) // Set the shader dimensions or scalar processing parameter.
63 .set<unsigned>("parameter_count", parameter_count) // Set the shader dimensions or scalar processing parameter.
64 .set<float>("learning_rate", learning_rate) // Set the shader dimensions or scalar processing parameter.
65 .set<float>("weight_decay", 0.0f); // Set the shader dimensions or scalar processing parameter.
66 inference_ = shdr::dispatch(*infer_, shdr::extent((batch_count_ + 63u) / 64u)); // Dispatch one inference invocation per sequence.
67 training_ = shdr::dispatch(*gradient_, // Dispatch one analytic derivative per sequence and parameter.
68 shdr::extent((batch_count_ * parameter_count + 63u) / 64u)) >> // Cover every sequence and parameter in the derivative dispatch.
69 shdr::dispatch(*update_, shdr::extent((parameter_count + 63u) / 64u)); // Reduce batch gradients and update each parameter exactly once.
70 reset_inputs(); // Create the initial repeatable training dataset.
71 reset_parameters(); // Create the initial repeatable trainable parameters.
72 } // Finish the transformer_model operation.
73 /** @brief Recreate a balanced, repeatable two-class sequence dataset.
74 * The class controls a feature pattern shared across tokens. small seeded
75 * noise and token variations prevent identical sequences within each class.
76 */
77 void reset_inputs(unsigned seed = 7127u) { // Build the repeatable two-class token dataset.
78 std::mt19937 generator(seed); // Use a fixed seed so every backend receives identical data.
79 std::uniform_real_distribution<float> noise(-0.12f, 0.12f); // Choose a bounded range for reproducible noise or weights.
80 const float pattern[4] = {1.0f, -0.7f, 0.4f, -0.2f}; // Define the feature pattern distinguishing the two classes.
81 float* values = inputs.data(); // Acquire the host storage that this initialization will overwrite.
82 unsigned* targets = labels.data(); // Acquire the host storage that this initialization will overwrite.
83 for (unsigned row = 0; row < batch_count_; ++row) { // Visit each token sequence in the batch.
84 targets[row] = row % 2u; // Alternate the two classes for a balanced dataset.
85 const float sign = targets[row] ? 1.0f : -1.0f; // Choose the class-dependent sign of the feature pattern.
86 for (unsigned token = 0; token < token_count; ++token) // Visit each of the four tokens.
87 for (unsigned c = 0; c < width; ++c) // Visit every value needed by this calculation.
88 values[row * 16u + token * 4u + c] = // Write this feature into its sequence, token and component slot.
89 sign * pattern[c] + noise(generator) + float(token) * 0.02f; // Combine the class pattern with small noise and token variation.
90 } // Finish this channel, sample or feature pass.
91 } // Finish the reset_inputs operation.
92 /** @brief Restore deterministic trainable parameters, including RMSNorm scales. */
93 void reset_parameters(unsigned seed = 59315u) { // Restore repeatable trainable model parameters.
94 std::mt19937 generator(seed); // Use a fixed seed so every backend receives identical data.
95 std::uniform_real_distribution<float> distribution(-0.15f, 0.15f); // Choose a bounded range for reproducible noise or weights.
96 float* values = parameters.data(); // Acquire the host storage that this initialization will overwrite.
97 for (unsigned i = 0; i < parameter_count; ++i) // Visit every trainable parameter.
98 values[i] = distribution(generator); // Initialize this trainable value from the fixed random seed.
99 for (unsigned i = 0; i < 16u; ++i) // Visit every value needed by this calculation.
100 values[i] *= 0.1f; // Keep positional embeddings small at initialization.
101 for (unsigned i = 80u; i < 84u; ++i) // Visit every value needed by this calculation.
102 values[i] = 1.0f; // Initialize RMS normalization scales to unity.
103 for (unsigned i = 100u; i < 104u; ++i) // Visit every value needed by this calculation.
104 values[i] = 0.5f; // Keep the initial feed-forward hidden units active.
105 for (unsigned i = 120u; i < 124u; ++i) // Visit every value needed by this calculation.
106 values[i] = 0.0f; // Start the output feed-forward bias at zero.
107 } // Finish the reset_parameters operation.
108 /** @brief Return the sample count shared by inference and gradient dispatches. */
109 unsigned batch_count() const { // Return the number of sequences in this model batch.
110 return batch_count_; // Remember how many token sequences share this model.
111 } // Finish the batch_count operation.
112 /** @brief Return a two-stage analytic gradient followed by averaged SGD update. */
113 const shdr::pipeline& training_pipeline() const { // Expose analytic derivatives followed by averaged SGD.
114 return training_; // Expose analytic gradients followed by averaged SGD.
115 } // Finish the training_pipeline operation.
116 /** @brief Return the classifier inference dispatch. */
117 const shdr::pipeline& prediction_pipeline() const { // Expose the classifier inference dispatch.
118 return inference_; // Expose the classifier inference dispatch.
119 } // Finish the prediction_pipeline operation.
120 /** @brief Return the dispatch that computes analytic per-sample derivatives. */
121 shdr::pipeline gradient_pipeline() const { // Build the analytic derivative dispatch.
122 return shdr::pipeline( // Expose the per-sample analytic gradient stage.
123 shdr::dispatch(*gradient_, shdr::extent((batch_count_ * parameter_count + 63u) / 64u))); // Cover the sample-major gradient resource with 64-wide groups.
124 } // Finish the gradient_pipeline operation.
125 /** @brief Read batch-mean cross entropy after inference has produced predictions. */
126 float loss() const { // Read the mean cross entropy for the current predictions.
127 const shdr::glsl::vec4* values = predictions.data(); // Read the resource after any required GPU result synchronization.
128 double sum = 0.0; // Start the scalar accumulator at zero.
129 for (unsigned row = 0; row < batch_count_; ++row) // Visit each token sequence in the batch.
130 sum += values[row].z; // Add this sequence loss to the batch total.
131 return static_cast<float>(sum / batch_count_); // Read the mean cross entropy for the current predictions.
132 } // Finish the loss operation.
133 /** @brief Read the number of correctly classified training sequences. */
134 unsigned correct() const { // Count correctly classified sequences.
135 const shdr::glsl::vec4* values = predictions.data(); // Read the resource after any required GPU result synchronization.
136 const unsigned* targets = labels.data(); // Read the resource after any required GPU result synchronization.
137 unsigned count = 0; // Begin counting correct sequence classifications.
138 for (unsigned row = 0; row < batch_count_; ++row) // Visit each token sequence in the batch.
139 if (static_cast<unsigned>(values[row].w) == targets[row]) // Handle this validation or boundary case explicitly.
140 ++count; // Count this correctly classified sequence.
141 return count; // Return the validated sequence count.
142 } // Finish the correct operation.
143}; // Finish the transformer_model resource helper.
144} // namespace shdr_examples
145#endif // Finish the conditional compilation block.