From f28f67d988158d6c46f7ffe967152f98d32a37b2 Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Mon, 30 Jun 2025 14:32:50 -0700
Subject: Add MLP training examples. (#7550)

* Add MLP training examples.

* Formatting fix.

* Fix.

* Improve documentation on coopvector.

* Improve doc.

* Update doc.

* Fix typo.

* Cleanup shader.

* Cleanup.

* Fix test.

* Fix type check recursion.

* Fix.

* Fix.

* Fix override check.
---
 examples/mlp-training/kernels.slang | 41 +++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)
 create mode 100644 examples/mlp-training/kernels.slang

(limited to 'examples/mlp-training/kernels.slang')

diff --git a/examples/mlp-training/kernels.slang b/examples/mlp-training/kernels.slang
new file mode 100644
index 000000000..5be076879
--- /dev/null
+++ b/examples/mlp-training/kernels.slang
@@ -0,0 +1,41 @@
+module kernels;
+
+import common;
+import mlp_sw;
+import network;
+import adam;
+
+[numthreads(256, 1, 1)]
+[require(spvGroupNonUniformBallot, spvGroupNonUniformArithmetic)]
+void learnGradient(
+    uint32_t tid : SV_DispatchThreadID,
+    uniform MyNetwork* network, 
+    uniform Atomic<uint32_t>* lossBuffer,
+    uniform float2* inputs,
+    uniform uint32_t count)
+{
+    if (tid >= count)
+        return;
+
+    var input = (half2)inputs[tid];
+    bwd_diff(loss)(network, input.x, input.y, 1.0h);
+    let thisLoss = (float)loss(network, input.x, input.y);
+    let maxLoss = WaveActiveMax(thisLoss);
+    if (WaveIsFirstLane())
+    {
+        lossBuffer.max(bit_cast<uint32_t>(maxLoss));
+    }
+}
+
+[numthreads(256, 1, 1)]
+void adjustParameters(uint32_t tid : SV_DispatchThreadID, uniform AdamState* states, uniform NFloat* params, uniform NFloat* gradients, uniform uint32_t count)
+{
+    if (tid >= count)
+        return;
+    if (isnan(gradients[tid]))
+    {
+        gradients[tid] = 0.0h;
+        return;
+    }
+    AdamOptimizer::step(states[tid], params[tid], gradients[tid]);
+}
\ No newline at end of file
-- 
cgit v1.2.3