From 1c1f8282712ae10f7f1820ba95a938f47682d774 Mon Sep 17 00:00:00 2001 From: Justin Marshall Date: Mon, 18 May 2026 07:33:20 -0700 Subject: [PATCH] Added nn trainer. --- neo/engine/opengl/nn_mattrain.cpp | 2388 +++++++++++++++++++++++++++++ neo/engine/opengl/opengl_nn.h | 87 ++ 2 files changed, 2475 insertions(+) create mode 100644 neo/engine/opengl/nn_mattrain.cpp create mode 100644 neo/engine/opengl/opengl_nn.h diff --git a/neo/engine/opengl/nn_mattrain.cpp b/neo/engine/opengl/nn_mattrain.cpp new file mode 100644 index 00000000..452c0d95 --- /dev/null +++ b/neo/engine/opengl/nn_mattrain.cpp @@ -0,0 +1,2388 @@ + +/* +=========================================================================== + +IceTech / QD3D12 Neural POM Material Trainer +Copyright (C) 2026 Justin Marshall + +=========================================================================== +*/ + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#pragma comment(lib, "d3d12.lib") +#pragma comment(lib, "d3dcompiler.lib") + +using Microsoft::WRL::ComPtr; + +#include "opengl.h" + +namespace QD3D12NeuralPOM +{ + static const uint32_t kWeightsMagic = 0x4D504E49u; // 'INPM' little-endian + static const uint32_t kWeightsVersion = 1u; + static const int kOutputCount = 13; + + // Deeper Neural POM training/export defaults. The renderer reads + // deltaUVScale from the exported weight header, so this value must match + // the teacher target normalization below. 0.120 gives the network enough + // headroom for the deeper ray target without saturating output0/output1. + static const float kDeltaUVScale = 0.120f; + static const float kTrainingDepthAmplify = 1.50f; + static const float kTrainingNormalAmplify = 1.25f; + static const float kTeacherParallaxBaseScale = 0.026f; + static const float kTeacherParallaxShiftAmplify = 1.50f; + + static const uint32_t kPackedSampleFloat4s = 6u; + static const uint32_t kPackedSampleBytes = kPackedSampleFloat4s * 16u; + + // CPU optimizer minibatching is the primary training-speed fix. + // The original trainer applied Adam to every MLP parameter once per sample; + // at the default 300k samples that means billions of scalar Adam updates. + // Accumulating a larger minibatch applies one averaged Adam step and also + // lets the latent grid use sparse batched Adam below. + static const uint32_t kCpuOptimizerBatch = 256u; + + // Two readback/UAV slots let the GPU generate the next teacher batch while + // the CPU trains on the previous batch instead of blocking every dispatch. + static const uint32_t kTeacherPipelineSlots = 2u; + + static inline float Clamp01(float v) { return v < 0.0f ? 0.0f : (v > 1.0f ? 1.0f : v); } + static inline float ClampF(float v, float lo, float hi) { return v < lo ? lo : (v > hi ? hi : v); } + static inline float LerpF(float a, float b, float t) { return a + (b - a) * t; } + static inline float Fract(float v) { return v - floorf(v); } + + static void Log(const char* fmt, ...) + { + char buffer[4096]; + va_list args; + va_start(args, fmt); + vsnprintf(buffer, sizeof(buffer), fmt, args); + va_end(args); + OutputDebugStringA(buffer); + OutputDebugStringA("\n"); + } + + static void Progress(const QD3D12NeuralPOMTrainDesc& desc, const char* message, uint32_t current, uint32_t total, float loss) + { + if ((desc.flags & QD3D12_NEURAL_POM_TRAIN_FLAG_VERBOSE) != 0u) + Log("%s: %u / %u loss %.6f", message ? message : "neuralPOM", current, total, loss); + if (desc.progress) + desc.progress(message ? message : "neuralPOM", current, total, loss, desc.progressUserData); + } + + static int Fail(QD3D12NeuralPOMTrainStats* stats, const char* fmt, ...) + { + char buffer[1024]; + va_list args; + va_start(args, fmt); + vsnprintf(buffer, sizeof(buffer), fmt, args); + va_end(args); + if (stats) + { + strncpy_s(stats->error, buffer, _TRUNCATE); + } + Log("QD3D12 NeuralPOM trainer failed: %s", buffer); + return 0; + } + + static void CopyPath(char* dst, size_t dstBytes, const std::string& path) + { + if (!dst || dstBytes == 0) + return; + strncpy_s(dst, dstBytes, path.c_str(), _TRUNCATE); + } + + static bool StringIsEmptyOrDash(const char* s) + { + return !s || !s[0] || (s[0] == '-' && s[1] == '\0'); + } + + static std::string WithSuffix(const char* prefix, const char* suffix) + { + std::string out = prefix ? prefix : "neural_material"; + out += suffix; + return out; + } + + static void CreateParentDirectoriesForFile(const char* path) + { + if (!path || !path[0]) + return; + + char tmp[QD3D12_NEURAL_POM_MAX_PATH_CHARS]; + strncpy_s(tmp, path, _TRUNCATE); + + for (char* p = tmp; *p; ++p) + { + if (*p == '/') + *p = '\\'; + } + + for (char* p = tmp; *p; ++p) + { + if (*p == '\\') + { + // Skip drive roots like C:\ + if (p == tmp || (p == tmp + 2 && tmp[1] == ':')) + continue; + char old = *p; + *p = '\0'; + if (tmp[0]) + CreateDirectoryA(tmp, nullptr); + *p = old; + } + } + } + + static bool WriteFileBinary(const char* path, const void* data, size_t bytes) + { + CreateParentDirectoriesForFile(path); + FILE* f = nullptr; + if (fopen_s(&f, path, "wb") != 0 || !f) + return false; + size_t wrote = 0; + if (bytes > 0) + wrote = fwrite(data, 1, bytes, f); + fclose(f); + return wrote == bytes; + } + + static bool WriteFileString(const char* path, const std::string& s) + { + return WriteFileBinary(path, s.c_str(), s.size()); + } + + static bool ReadFileBinary(const char* path, std::vector& out) + { + out.clear(); + if (!path || !path[0]) + return false; + FILE* f = nullptr; + if (fopen_s(&f, path, "rb") != 0 || !f) + return false; + fseek(f, 0, SEEK_END); + long len = ftell(f); + fseek(f, 0, SEEK_SET); + if (len <= 0) + { + fclose(f); + return false; + } + out.resize((size_t)len); + const size_t got = fread(out.data(), 1, out.size(), f); + fclose(f); + return got == out.size(); + } + + struct Vec2 + { + float x, y; + Vec2() : x(0), y(0) {} + Vec2(float X, float Y) : x(X), y(Y) {} + }; + + struct Vec3 + { + float x, y, z; + Vec3() : x(0), y(0), z(0) {} + Vec3(float X, float Y, float Z) : x(X), y(Y), z(Z) {} + }; + + struct Vec4 + { + float x, y, z, w; + Vec4() : x(0), y(0), z(0), w(1) {} + Vec4(float X, float Y, float Z, float W) : x(X), y(Y), z(Z), w(W) {} + }; + + static inline Vec2 operator+(const Vec2& a, const Vec2& b) { return Vec2(a.x + b.x, a.y + b.y); } + static inline Vec2 operator-(const Vec2& a, const Vec2& b) { return Vec2(a.x - b.x, a.y - b.y); } + static inline Vec2 operator*(const Vec2& a, float s) { return Vec2(a.x * s, a.y * s); } + static inline float Dot2(const Vec2& a, const Vec2& b) { return a.x * b.x + a.y * b.y; } + static inline float Len2(const Vec2& a) { return sqrtf(Dot2(a, a)); } + static inline float Dot3(const Vec3& a, const Vec3& b) { return a.x * b.x + a.y * b.y + a.z * b.z; } + static inline Vec3 Normalize3(const Vec3& v, const Vec3& fallback) + { + float l2 = Dot3(v, v); + if (l2 <= 1e-12f) + return fallback; + float inv = 1.0f / sqrtf(l2); + return Vec3(v.x * inv, v.y * inv, v.z * inv); + } + + struct ImageRGBA8 + { + int width; + int height; + std::vector rgba; + + ImageRGBA8() : width(0), height(0) {} + bool Valid() const { return width > 0 && height > 0 && rgba.size() == (size_t)width * (size_t)height * 4u; } + + bool FromExternal(const QD3D12NeuralPOMImageRGBA8& img) + { + width = 0; + height = 0; + rgba.clear(); + if (!img.pixelsRGBA8 || img.width == 0 || img.height == 0) + return false; + const uint32_t srcPitch = img.rowPitchBytes ? img.rowPitchBytes : img.width * 4u; + if (srcPitch < img.width * 4u) + return false; + width = (int)img.width; + height = (int)img.height; + rgba.resize((size_t)width * (size_t)height * 4u); + const uint8_t* src = (const uint8_t*)img.pixelsRGBA8; + for (uint32_t y = 0; y < img.height; ++y) + { + memcpy(&rgba[(size_t)y * (size_t)width * 4u], src + (size_t)y * srcPitch, (size_t)width * 4u); + } + return Valid(); + } + }; + + static uint16_t ReadLE16(const uint8_t* p) { return (uint16_t)(p[0] | (p[1] << 8)); } + + static bool LoadTGAFromMemory(const uint8_t* data, size_t size, ImageRGBA8& out) + { + out = ImageRGBA8(); + if (!data || size < 18) + return false; + + const uint8_t idLen = data[0]; + const uint8_t colorMapType = data[1]; + const uint8_t imageType = data[2]; + const uint16_t width = ReadLE16(data + 12); + const uint16_t height = ReadLE16(data + 14); + const uint8_t bits = data[16]; + const uint8_t desc = data[17]; + + if (colorMapType != 0 || width == 0 || height == 0) + return false; + if (!(imageType == 2 || imageType == 3 || imageType == 10 || imageType == 11)) + return false; + if (!(bits == 8 || bits == 24 || bits == 32)) + return false; + if ((imageType == 3 || imageType == 11) && bits != 8) + return false; + + const size_t pixelBytes = (size_t)bits / 8u; + size_t pos = 18u + (size_t)idLen; + if (pos > size) + return false; + + out.width = (int)width; + out.height = (int)height; + out.rgba.assign((size_t)width * (size_t)height * 4u, 255u); + + const bool rle = (imageType == 10 || imageType == 11); + const bool topOrigin = (desc & 0x20) != 0; + const size_t totalPixels = (size_t)width * (size_t)height; + + size_t pixelIndex = 0; + while (pixelIndex < totalPixels) + { + int runCount = 1; + uint8_t packet[4] = { 0, 0, 0, 255 }; + bool packetIsRLE = false; + + if (rle) + { + if (pos >= size) + return false; + uint8_t header = data[pos++]; + packetIsRLE = (header & 0x80u) != 0; + runCount = (header & 0x7Fu) + 1; + } + + for (int r = 0; r < runCount && pixelIndex < totalPixels; ++r, ++pixelIndex) + { + if (!rle || !packetIsRLE || r == 0) + { + if (pos + pixelBytes > size) + return false; + if (bits == 8) + { + uint8_t l = data[pos++]; + packet[0] = l; packet[1] = l; packet[2] = l; packet[3] = 255; + } + else if (bits == 24) + { + uint8_t b = data[pos++]; + uint8_t g = data[pos++]; + uint8_t rr = data[pos++]; + packet[0] = rr; packet[1] = g; packet[2] = b; packet[3] = 255; + } + else + { + uint8_t b = data[pos++]; + uint8_t g = data[pos++]; + uint8_t rr = data[pos++]; + uint8_t a = data[pos++]; + packet[0] = rr; packet[1] = g; packet[2] = b; packet[3] = a; + } + } + + size_t srcX = pixelIndex % width; + size_t srcY = pixelIndex / width; + size_t dstY = topOrigin ? srcY : ((size_t)height - 1u - srcY); + uint8_t* dst = &out.rgba[(dstY * (size_t)width + srcX) * 4u]; + dst[0] = packet[0]; + dst[1] = packet[1]; + dst[2] = packet[2]; + dst[3] = packet[3]; + } + } + + return out.Valid(); + } + + static bool LoadTGAFile(const char* path, ImageRGBA8& out) + { + std::vector bytes; + if (!ReadFileBinary(path, bytes)) + { + std::string withExt = std::string(path ? path : "") + ".tga"; + if (!ReadFileBinary(withExt.c_str(), bytes)) + return false; + } + return LoadTGAFromMemory(bytes.empty() ? nullptr : bytes.data(), bytes.size(), out); + } + + struct Rng + { + uint32_t s; + explicit Rng(uint32_t seed) : s(seed ? seed : 0x1234567u) {} + uint32_t U32() + { + uint32_t x = s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + s = x ? x : 0xBADC0DEu; + return s; + } + float Float01() { return (float)(U32() & 0x00FFFFFFu) * (1.0f / 16777216.0f); } + float FloatRange(float a, float b) { return a + (b - a) * Float01(); } + }; + + struct AdamVector + { + std::vector m; + std::vector v; + void Resize(size_t n) + { + m.assign(n, 0.0f); + v.assign(n, 0.0f); + } + }; + + struct AdamStepParams + { + float beta1; + float beta2; + float oneMinusBeta1; + float oneMinusBeta2; + float invBiasCorrection1; + float invBiasCorrection2; + float eps; + }; + + static AdamStepParams MakeAdamStepParams(int t) + { + AdamStepParams p; + p.beta1 = 0.9f; + p.beta2 = 0.999f; + p.oneMinusBeta1 = 1.0f - p.beta1; + p.oneMinusBeta2 = 1.0f - p.beta2; + p.eps = 1e-8f; + const float ft = (float)max(1, t); + p.invBiasCorrection1 = 1.0f / max(1.0f - powf(p.beta1, ft), 1e-12f); + p.invBiasCorrection2 = 1.0f / max(1.0f - powf(p.beta2, ft), 1e-12f); + return p; + } + + static inline void AdamUpdateFast(float& p, float& m, float& v, float g, float lr, const AdamStepParams& ap) + { + m = ap.beta1 * m + ap.oneMinusBeta1 * g; + v = ap.beta2 * v + ap.oneMinusBeta2 * g * g; + const float mh = m * ap.invBiasCorrection1; + const float vh = v * ap.invBiasCorrection2; + p -= lr * mh / (sqrtf(vh) + ap.eps); + p = ClampF(p, -8.0f, 8.0f); + } + + static void AdamUpdate(float& p, float& m, float& v, float g, float lr, int t) + { + const AdamStepParams ap = MakeAdamStepParams(t); + AdamUpdateFast(p, m, v, g, lr, ap); + } + + struct MLP + { + int inputCount; + int hidden; + int outputCount; + std::vector w1, b1, w2, b2, w3, b3; + std::vector gw1, gb1, gw2, gb2, gw3, gb3; + AdamVector aw1, ab1, aw2, ab2, aw3, ab3; + + MLP() : inputCount(0), hidden(0), outputCount(0) {} + + void Init(int in, int h, int out, Rng& rng) + { + inputCount = in; + hidden = h; + outputCount = out; + w1.resize((size_t)hidden * inputCount); + b1.assign(hidden, 0.0f); + w2.resize((size_t)hidden * hidden); + b2.assign(hidden, 0.0f); + w3.resize((size_t)outputCount * hidden); + b3.assign(outputCount, 0.0f); + gw1.resize(w1.size()); gb1.resize(b1.size()); + gw2.resize(w2.size()); gb2.resize(b2.size()); + gw3.resize(w3.size()); gb3.resize(b3.size()); + aw1.Resize(w1.size()); ab1.Resize(b1.size()); + aw2.Resize(w2.size()); ab2.Resize(b2.size()); + aw3.Resize(w3.size()); ab3.Resize(b3.size()); + + float s1 = sqrtf(2.0f / (float)(inputCount + hidden)); + float s2 = sqrtf(2.0f / (float)(hidden + hidden)); + float s3 = sqrtf(2.0f / (float)(hidden + outputCount)); + for (size_t i = 0; i < w1.size(); ++i) w1[i] = rng.FloatRange(-s1, s1); + for (size_t i = 0; i < w2.size(); ++i) w2[i] = rng.FloatRange(-s2, s2); + for (size_t i = 0; i < w3.size(); ++i) w3[i] = rng.FloatRange(-s3, s3); + } + + static float Sigmoid(float x) + { + x = ClampF(x, -40.0f, 40.0f); + return 1.0f / (1.0f + expf(-x)); + } + + static float ActivateOut(int index, float z) + { + if (index == 0 || index == 1 || index == 3 || index == 4) + return tanhf(z); + return Sigmoid(z); + } + + static float ActivateOutDeriv(int index, float y) + { + if (index == 0 || index == 1 || index == 3 || index == 4) + return 1.0f - y * y; + return y * (1.0f - y); + } + + void Forward(const std::vector& x, std::vector& h1, std::vector& h2, std::vector& y) const + { + if (h1.size() != (size_t)hidden) h1.resize((size_t)hidden); + if (h2.size() != (size_t)hidden) h2.resize((size_t)hidden); + if (y.size() != (size_t)outputCount) y.resize((size_t)outputCount); + for (int j = 0; j < hidden; ++j) + { + float z = b1[j]; + const float* w = &w1[(size_t)j * inputCount]; + for (int i = 0; i < inputCount; ++i) + z += w[i] * x[i]; + h1[j] = tanhf(z); + } + for (int j = 0; j < hidden; ++j) + { + float z = b2[j]; + const float* w = &w2[(size_t)j * hidden]; + for (int i = 0; i < hidden; ++i) + z += w[i] * h1[i]; + h2[j] = tanhf(z); + } + for (int j = 0; j < outputCount; ++j) + { + float z = b3[j]; + const float* w = &w3[(size_t)j * hidden]; + for (int i = 0; i < hidden; ++i) + z += w[i] * h2[i]; + y[j] = ActivateOut(j, z); + } + } + + void ZeroGrad() + { + std::fill(gw1.begin(), gw1.end(), 0.0f); + std::fill(gb1.begin(), gb1.end(), 0.0f); + std::fill(gw2.begin(), gw2.end(), 0.0f); + std::fill(gb2.begin(), gb2.end(), 0.0f); + std::fill(gw3.begin(), gw3.end(), 0.0f); + std::fill(gb3.begin(), gb3.end(), 0.0f); + } + + void BackwardAccumulate(const std::vector& x, const std::vector& h1, const std::vector& h2, + const std::vector& y, const std::vector& dy, std::vector& dx, + std::vector& dz3, std::vector& dh2, std::vector& dz2, + std::vector& dh1, std::vector& dz1) + { + if (dx.size() != (size_t)inputCount) dx.resize((size_t)inputCount); + if (dz3.size() != (size_t)outputCount) dz3.resize((size_t)outputCount); + if (dh2.size() != (size_t)hidden) dh2.resize((size_t)hidden); + if (dz2.size() != (size_t)hidden) dz2.resize((size_t)hidden); + if (dh1.size() != (size_t)hidden) dh1.resize((size_t)hidden); + if (dz1.size() != (size_t)hidden) dz1.resize((size_t)hidden); + std::fill(dx.begin(), dx.end(), 0.0f); + std::fill(dh2.begin(), dh2.end(), 0.0f); + std::fill(dh1.begin(), dh1.end(), 0.0f); + + for (int o = 0; o < outputCount; ++o) + { + dz3[o] = dy[o] * ActivateOutDeriv(o, y[o]); + gb3[o] += dz3[o]; + const float* w = &w3[(size_t)o * hidden]; + float* gw = &gw3[(size_t)o * hidden]; + for (int j = 0; j < hidden; ++j) + { + gw[j] += dz3[o] * h2[j]; + dh2[j] += w[j] * dz3[o]; + } + } + for (int j = 0; j < hidden; ++j) + { + dz2[j] = dh2[j] * (1.0f - h2[j] * h2[j]); + gb2[j] += dz2[j]; + const float* w = &w2[(size_t)j * hidden]; + float* gw = &gw2[(size_t)j * hidden]; + for (int i = 0; i < hidden; ++i) + { + gw[i] += dz2[j] * h1[i]; + dh1[i] += w[i] * dz2[j]; + } + } + for (int j = 0; j < hidden; ++j) + { + dz1[j] = dh1[j] * (1.0f - h1[j] * h1[j]); + gb1[j] += dz1[j]; + const float* w = &w1[(size_t)j * inputCount]; + float* gw = &gw1[(size_t)j * inputCount]; + for (int i = 0; i < inputCount; ++i) + { + gw[i] += dz1[j] * x[i]; + dx[i] += w[i] * dz1[j]; + } + } + } + + void Backward(const std::vector& x, const std::vector& h1, const std::vector& h2, + const std::vector& y, const std::vector& dy, std::vector& dx) + { + ZeroGrad(); + std::vector dz3, dh2, dz2, dh1, dz1; + BackwardAccumulate(x, h1, h2, y, dy, dx, dz3, dh2, dz2, dh1, dz1); + } + + void ApplyAdamScaled(float lr, int t, float gradScale, float l2) + { + const AdamStepParams ap = MakeAdamStepParams(t); + for (size_t i = 0; i < w1.size(); ++i) AdamUpdateFast(w1[i], aw1.m[i], aw1.v[i], gw1[i] * gradScale + w1[i] * l2, lr, ap); + for (size_t i = 0; i < b1.size(); ++i) AdamUpdateFast(b1[i], ab1.m[i], ab1.v[i], gb1[i] * gradScale, lr, ap); + for (size_t i = 0; i < w2.size(); ++i) AdamUpdateFast(w2[i], aw2.m[i], aw2.v[i], gw2[i] * gradScale + w2[i] * l2, lr, ap); + for (size_t i = 0; i < b2.size(); ++i) AdamUpdateFast(b2[i], ab2.m[i], ab2.v[i], gb2[i] * gradScale, lr, ap); + for (size_t i = 0; i < w3.size(); ++i) AdamUpdateFast(w3[i], aw3.m[i], aw3.v[i], gw3[i] * gradScale + w3[i] * l2, lr, ap); + for (size_t i = 0; i < b3.size(); ++i) AdamUpdateFast(b3[i], ab3.m[i], ab3.v[i], gb3[i] * gradScale, lr, ap); + } + + void ApplyAdam(float lr, int t) + { + ApplyAdamScaled(lr, t, 1.0f, 0.0f); + } + }; + + struct LatentGrid + { + int res; + int channels; + std::vector data; + std::vector adamM; + std::vector adamV; + + // Sparse minibatch gradient storage. A dense clear of the full latent + // texture every optimizer step is expensive; stamps let us clear only + // the texels touched by the current CPU optimizer batch. + std::vector grad; + std::vector gradStamp; + std::vector touched; + uint32_t gradEpoch; + + LatentGrid() : res(0), channels(0), gradEpoch(1u) {} + + void Init(int r, int c, Rng& rng) + { + res = r; + channels = c; + data.resize((size_t)res * res * channels); + adamM.assign(data.size(), 0.0f); + adamV.assign(data.size(), 0.0f); + grad.assign(data.size(), 0.0f); + gradStamp.assign(data.size(), 0u); + touched.clear(); + touched.reserve((size_t)std::max(1024u, kCpuOptimizerBatch * 4u * (uint32_t)channels)); + gradEpoch = 1u; + for (size_t i = 0; i < data.size(); ++i) + data[i] = rng.FloatRange(-0.05f, 0.05f); + } + + int Index(int x, int y, int c) const + { + x %= res; y %= res; + if (x < 0) x += res; + if (y < 0) y += res; + return ((y * res + x) * channels + c); + } + + void Sample(float u, float v, std::vector& out, int idx[4], float w[4]) const + { + if (out.size() != (size_t)channels) out.resize((size_t)channels); + u = Fract(u); v = Fract(v); + if (u < 0.0f) u += 1.0f; + if (v < 0.0f) v += 1.0f; + float fx = u * (float)res - 0.5f; + float fy = v * (float)res - 0.5f; + int x0 = (int)floorf(fx); + int y0 = (int)floorf(fy); + float tx = fx - (float)x0; + float ty = fy - (float)y0; + int x1 = x0 + 1; + int y1 = y0 + 1; + idx[0] = Index(x0, y0, 0); + idx[1] = Index(x1, y0, 0); + idx[2] = Index(x0, y1, 0); + idx[3] = Index(x1, y1, 0); + w[0] = (1.0f - tx) * (1.0f - ty); + w[1] = tx * (1.0f - ty); + w[2] = (1.0f - tx) * ty; + w[3] = tx * ty; + for (int c = 0; c < channels; ++c) + { + out[c] = data[idx[0] + c] * w[0] + data[idx[1] + c] * w[1] + data[idx[2] + c] * w[2] + data[idx[3] + c] * w[3]; + } + } + + void BeginGradBatch() + { + touched.clear(); + ++gradEpoch; + if (gradEpoch == 0u) + { + std::fill(gradStamp.begin(), gradStamp.end(), 0u); + gradEpoch = 1u; + } + } + + void AddGrad(int id, float g) + { + if ((uint32_t)id >= gradStamp.size()) + return; + if (gradStamp[(size_t)id] != gradEpoch) + { + gradStamp[(size_t)id] = gradEpoch; + grad[(size_t)id] = 0.0f; + touched.push_back((uint32_t)id); + } + grad[(size_t)id] += g; + } + + void AccumulateGrad(const int idx[4], const float w[4], const std::vector& dLatent) + { + for (int k = 0; k < 4; ++k) + { + for (int c = 0; c < channels; ++c) + { + const int id = idx[k] + c; + AddGrad(id, dLatent[c] * w[k]); + } + } + } + + void ApplyAdamSparse(float lr, int t, float gradScale, float l2) + { + if (touched.empty()) + return; + const AdamStepParams ap = MakeAdamStepParams(t); + for (size_t n = 0; n < touched.size(); ++n) + { + const uint32_t idu = touched[n]; + const size_t id = (size_t)idu; + const float g = grad[id] * gradScale + data[id] * l2; + AdamUpdateFast(data[id], adamM[id], adamV[id], g, lr, ap); + data[id] = ClampF(data[id], -4.0f, 4.0f); + } + } + + // Compatibility path for single-sample latent Adam. The trainer below + // now uses sparse minibatch Adam, but this remains useful for A/B tests. + void BackwardAdam(const int idx[4], const float w[4], const std::vector& dLatent, float lr, int t) + { + const AdamStepParams ap = MakeAdamStepParams(t); + for (int k = 0; k < 4; ++k) + { + for (int c = 0; c < channels; ++c) + { + int id = idx[k] + c; + float g = dLatent[c] * w[k] + data[id] * 0.000001f; + AdamUpdateFast(data[id], adamM[id], adamV[id], g, lr, ap); + data[id] = ClampF(data[id], -4.0f, 4.0f); + } + } + } + }; + + static uint16_t FloatToHalf(float f) + { + union { float f; uint32_t u; } v; + v.f = f; + uint32_t sign = (v.u >> 16) & 0x8000u; + int exp = (int)((v.u >> 23) & 0xFFu) - 127 + 15; + uint32_t mant = v.u & 0x7FFFFFu; + + if (exp <= 0) + { + if (exp < -10) + return (uint16_t)sign; + mant = (mant | 0x800000u) >> (1 - exp); + return (uint16_t)(sign | ((mant + 0x1000u) >> 13)); + } + if (exp >= 31) + { + return (uint16_t)(sign | 0x7C00u); + } + return (uint16_t)(sign | ((uint32_t)exp << 10) | ((mant + 0x1000u) >> 13)); + } + +#pragma pack(push, 1) + struct WeightsHeader + { + uint32_t magic; + uint32_t version; + uint32_t inputCount; + uint32_t hiddenCount; + uint32_t outputCount; + uint32_t latentResolution; + uint32_t latentChannels; + float deltaUVScale; + }; +#pragma pack(pop) + + static void AppendBytes(std::vector& out, const void* data, size_t bytes) + { + if (!data || bytes == 0) + return; + size_t old = out.size(); + out.resize(old + bytes); + memcpy(out.data() + old, data, bytes); + } + + static void AppendFloatArray(std::vector& out, const std::vector& a) + { + if (!a.empty()) + AppendBytes(out, a.data(), a.size() * sizeof(float)); + } + + static bool SaveWeights(const char* path, const MLP& net, const LatentGrid& latent) + { + WeightsHeader h; + h.magic = kWeightsMagic; + h.version = kWeightsVersion; + h.inputCount = (uint32_t)net.inputCount; + h.hiddenCount = (uint32_t)net.hidden; + h.outputCount = (uint32_t)net.outputCount; + h.latentResolution = (uint32_t)latent.res; + h.latentChannels = (uint32_t)latent.channels; + h.deltaUVScale = kDeltaUVScale; + + std::vector bytes; + AppendBytes(bytes, &h, sizeof(h)); + AppendFloatArray(bytes, net.w1); AppendFloatArray(bytes, net.b1); + AppendFloatArray(bytes, net.w2); AppendFloatArray(bytes, net.b2); + AppendFloatArray(bytes, net.w3); AppendFloatArray(bytes, net.b3); + return WriteFileBinary(path, bytes.empty() ? nullptr : bytes.data(), bytes.size()); + } + + static bool SaveLatentF32(const char* path, const LatentGrid& latent) + { + return WriteFileBinary(path, latent.data.empty() ? nullptr : latent.data.data(), latent.data.size() * sizeof(float)); + } + + static bool SaveLatentRGBA16F(const char* path, const LatentGrid& latent) + { + const int rgbaSlices = (latent.channels + 3) / 4; + const size_t texels = (size_t)latent.res * (size_t)latent.res; + std::vector half(texels * (size_t)rgbaSlices * 4u, 0u); + for (int y = 0; y < latent.res; ++y) + { + for (int x = 0; x < latent.res; ++x) + { + size_t srcBase = ((size_t)y * (size_t)latent.res + (size_t)x) * (size_t)latent.channels; + size_t dstTexel = (size_t)y * (size_t)latent.res + (size_t)x; + for (int c = 0; c < latent.channels; ++c) + { + int slice = c / 4; + int comp = c & 3; + size_t dst = ((size_t)slice * texels + dstTexel) * 4u + (size_t)comp; + half[dst] = FloatToHalf(latent.data[srcBase + c]); + } + } + } + return WriteFileBinary(path, half.empty() ? nullptr : half.data(), half.size() * sizeof(uint16_t)); + } + + static std::string JsonEscape(const char* s) + { + std::string out; + if (!s) + return out; + for (; *s; ++s) + { + char c = *s; + if (c == '\\' || c == '"') + { + out.push_back('\\'); + out.push_back(c); + } + else if (c == '\n') + { + out += "\\n"; + } + else if (c == '\r') + { + out += "\\r"; + } + else + { + out.push_back(c); + } + } + return out; + } + + static std::string BuildManifest( + const char* outPrefix, + const char* albedoSource, + const char* normalSource, + const char* specularSource, + const MLP& net, + const LatentGrid& latent, + uint32_t samples, + uint32_t samplesPerGpuBatch, + float normalStrength, + float ySign, + float finalLoss) + { + const int rgbaSlices = (latent.channels + 3) / 4; + const size_t sliceBytes = (size_t)latent.res * (size_t)latent.res * 4u * sizeof(uint16_t); + char buf[8192]; + snprintf(buf, sizeof(buf), + "{\n" + " \"format\": \"IceTechNeuralPOMMaterialV1\",\n" + " \"trainingEntryPoint\": \"QD3D12_NeuralPOMTrainMaterialD3D12\",\n" + " \"teacherBackend\": \"D3D12Compute\",\n" + " \"albedoSource\": \"%s\",\n" + " \"normalSource\": \"%s\",\n" + " \"specularSource\": \"%s\",\n" + " \"samples\": %u,\n" + " \"samplesPerGpuBatch\": %u,\n" + " \"inputCount\": %d,\n" + " \"hiddenCount\": %d,\n" + " \"outputCount\": %d,\n" + " \"latentResolution\": %d,\n" + " \"latentChannels\": %d,\n" + " \"latentRGBA16FSlices\": %d,\n" + " \"latentRGBA16FSliceBytes\": %u,\n" + " \"deltaUVScale\": %.9g,\n" + " \"trainingDepthAmplify\": %.9g,\n" + " \"trainingNormalAmplify\": %.9g,\n" + " \"teacherParallaxBaseScale\": %.9g,\n" + " \"teacherParallaxShiftAmplify\": %.9g,\n" + " \"normalStrength\": %.9g,\n" + " \"normalMapYSign\": %.9g,\n" + " \"finalMeanLoss\": %.9g,\n" + " \"weightsFile\": \"%s_weights.bin\",\n" + " \"latentF32File\": \"%s_latent_f32.bin\",\n" + " \"latentRGBA16FFile\": \"%s_latent_rgba16f.bin\",\n" + " \"outputs\": [\n" + " \"deltaUV.x = output0 * deltaUVScale\",\n" + " \"deltaUV.y = output1 * deltaUVScale\",\n" + " \"hitDepth\",\n" + " \"normalTS.x\",\n" + " \"normalTS.y\",\n" + " \"confidence\",\n" + " \"selfShadowVisibility\",\n" + " \"albedo.r\", \"albedo.g\", \"albedo.b\",\n" + " \"specular.r\", \"specular.g\", \"specular.b\"\n" + " ]\n" + "}\n", + JsonEscape(albedoSource).c_str(), + JsonEscape(normalSource).c_str(), + JsonEscape(specularSource ? specularSource : "").c_str(), + samples, + samplesPerGpuBatch, + net.inputCount, + net.hidden, + net.outputCount, + latent.res, + latent.channels, + rgbaSlices, + (unsigned int)sliceBytes, + kDeltaUVScale, + kTrainingDepthAmplify, + kTrainingNormalAmplify, + kTeacherParallaxBaseScale, + kTeacherParallaxShiftAmplify, + normalStrength, + ySign, + finalLoss, + outPrefix, + outPrefix, + outPrefix); + return std::string(buf); + } + + static UINT64 AlignUp(UINT64 v, UINT64 a) + { + return (v + (a - 1)) & ~(a - 1); + } + + static D3D12_RESOURCE_DESC BufferDesc(UINT64 size, D3D12_RESOURCE_FLAGS flags = D3D12_RESOURCE_FLAG_NONE) + { + D3D12_RESOURCE_DESC d = {}; + d.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER; + d.Alignment = 0; + d.Width = size; + d.Height = 1; + d.DepthOrArraySize = 1; + d.MipLevels = 1; + d.Format = DXGI_FORMAT_UNKNOWN; + d.SampleDesc.Count = 1; + d.SampleDesc.Quality = 0; + d.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR; + d.Flags = flags; + return d; + } + + static bool CreateCommittedBuffer( + ID3D12Device* device, + UINT64 size, + D3D12_HEAP_TYPE heapType, + D3D12_RESOURCE_STATES initialState, + D3D12_RESOURCE_FLAGS flags, + ComPtr& out) + { + out.Reset(); + D3D12_HEAP_PROPERTIES hp = {}; + hp.Type = heapType; + D3D12_RESOURCE_DESC rd = BufferDesc(size, flags); + HRESULT hr = device->CreateCommittedResource( + &hp, + D3D12_HEAP_FLAG_NONE, + &rd, + initialState, + nullptr, + IID_PPV_ARGS(&out)); + return SUCCEEDED(hr); + } + + struct GpuContext + { + ComPtr device; + ComPtr queue; + ComPtr allocator; + ComPtr list; + ComPtr fence; + HANDLE fenceEvent; + UINT64 nextFence; + + GpuContext() : fenceEvent(nullptr), nextFence(1) {} + ~GpuContext() + { + if (fenceEvent) + CloseHandle(fenceEvent); + } + + bool Init(ID3D12Device* inDevice, ID3D12CommandQueue* inQueue) + { + if (!inDevice || !inQueue) + return false; + device = inDevice; + queue = inQueue; + if (FAILED(device->CreateCommandAllocator(D3D12_COMMAND_LIST_TYPE_DIRECT, IID_PPV_ARGS(&allocator)))) + return false; + if (FAILED(device->CreateCommandList(0, D3D12_COMMAND_LIST_TYPE_DIRECT, allocator.Get(), nullptr, IID_PPV_ARGS(&list)))) + return false; + list->Close(); + if (FAILED(device->CreateFence(0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(&fence)))) + return false; + fenceEvent = CreateEventA(nullptr, FALSE, FALSE, nullptr); + return fenceEvent != nullptr; + } + + bool Begin() + { + if (FAILED(allocator->Reset())) + return false; + if (FAILED(list->Reset(allocator.Get(), nullptr))) + return false; + return true; + } + + bool EndSignal(UINT64* outFenceValue) + { + if (FAILED(list->Close())) + return false; + ID3D12CommandList* lists[] = { list.Get() }; + queue->ExecuteCommandLists(1, lists); + const UINT64 value = nextFence++; + if (FAILED(queue->Signal(fence.Get(), value))) + return false; + if (outFenceValue) + *outFenceValue = value; + return true; + } + + bool WaitFence(UINT64 value) + { + if (value == 0) + return true; + if (fence->GetCompletedValue() < value) + { + if (FAILED(fence->SetEventOnCompletion(value, fenceEvent))) + return false; + WaitForSingleObject(fenceEvent, INFINITE); + } + return true; + } + + bool EndWait() + { + UINT64 value = 0; + return EndSignal(&value) && WaitFence(value); + } + + void Transition(ID3D12Resource* res, D3D12_RESOURCE_STATES before, D3D12_RESOURCE_STATES after) + { + if (!res || before == after) + return; + D3D12_RESOURCE_BARRIER b = {}; + b.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION; + b.Transition.pResource = res; + b.Transition.StateBefore = before; + b.Transition.StateAfter = after; + b.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES; + list->ResourceBarrier(1, &b); + } + + void UavBarrier(ID3D12Resource* res) + { + D3D12_RESOURCE_BARRIER b = {}; + b.Type = D3D12_RESOURCE_BARRIER_TYPE_UAV; + b.UAV.pResource = res; + list->ResourceBarrier(1, &b); + } + }; + + struct GpuTexture + { + ComPtr resource; + UINT width; + UINT height; + }; + + static bool UploadImageAsTexture(GpuContext& gpu, const ImageRGBA8& img, GpuTexture& out, std::vector>& keepAlive) + { + out = GpuTexture{}; + if (!img.Valid()) + return false; + + D3D12_RESOURCE_DESC rd = {}; + rd.Dimension = D3D12_RESOURCE_DIMENSION_TEXTURE2D; + rd.Width = (UINT)img.width; + rd.Height = (UINT)img.height; + rd.DepthOrArraySize = 1; + rd.MipLevels = 1; + rd.Format = DXGI_FORMAT_R8G8B8A8_UNORM; + rd.SampleDesc.Count = 1; + rd.Layout = D3D12_TEXTURE_LAYOUT_UNKNOWN; + rd.Flags = D3D12_RESOURCE_FLAG_NONE; + + D3D12_HEAP_PROPERTIES hpDef = {}; + hpDef.Type = D3D12_HEAP_TYPE_DEFAULT; + if (FAILED(gpu.device->CreateCommittedResource(&hpDef, D3D12_HEAP_FLAG_NONE, &rd, D3D12_RESOURCE_STATE_COPY_DEST, nullptr, IID_PPV_ARGS(&out.resource)))) + return false; + + D3D12_PLACED_SUBRESOURCE_FOOTPRINT footprint = {}; + UINT rows = 0; + UINT64 rowSize = 0; + UINT64 uploadSize = 0; + gpu.device->GetCopyableFootprints(&rd, 0, 1, 0, &footprint, &rows, &rowSize, &uploadSize); + + ComPtr upload; + if (!CreateCommittedBuffer(gpu.device.Get(), uploadSize, D3D12_HEAP_TYPE_UPLOAD, D3D12_RESOURCE_STATE_GENERIC_READ, D3D12_RESOURCE_FLAG_NONE, upload)) + return false; + + uint8_t* dstBase = nullptr; + if (FAILED(upload->Map(0, nullptr, (void**)&dstBase)) || !dstBase) + return false; + memset(dstBase, 0, (size_t)uploadSize); + const uint8_t* srcBase = img.rgba.data(); + const UINT srcPitch = (UINT)img.width * 4u; + for (UINT y = 0; y < (UINT)img.height; ++y) + { + memcpy(dstBase + footprint.Offset + (size_t)y * footprint.Footprint.RowPitch, + srcBase + (size_t)y * srcPitch, + srcPitch); + } + upload->Unmap(0, nullptr); + + D3D12_TEXTURE_COPY_LOCATION src = {}; + src.pResource = upload.Get(); + src.Type = D3D12_TEXTURE_COPY_TYPE_PLACED_FOOTPRINT; + src.PlacedFootprint = footprint; + + D3D12_TEXTURE_COPY_LOCATION dst = {}; + dst.pResource = out.resource.Get(); + dst.Type = D3D12_TEXTURE_COPY_TYPE_SUBRESOURCE_INDEX; + dst.SubresourceIndex = 0; + + gpu.list->CopyTextureRegion(&dst, 0, 0, 0, &src, nullptr); + gpu.Transition(out.resource.Get(), D3D12_RESOURCE_STATE_COPY_DEST, D3D12_RESOURCE_STATE_NON_PIXEL_SHADER_RESOURCE); + + out.width = (UINT)img.width; + out.height = (UINT)img.height; + keepAlive.push_back(upload); + return true; + } + + static const char* gGpuTeacherHLSL = R"NPOMHLSL( +Texture2D gAlbedo : register(t0); +Texture2D gNormal : register(t1); +Texture2D gSpecular : register(t2); +SamplerState gSamp : register(s0); +RWStructuredBuffer gOut : register(u0); + +cbuffer TeacherCB : register(b0) +{ + uint gSampleCount; + uint gSeed; + uint gHasSpecular; + uint gPad0; + float gNormalStrength; + float gNormalMapYSign; + float gDepthAmplify; + float gNormalAmplify; + float gDeltaUVScale; + float gParallaxBaseScale; + float gParallaxShiftAmplify; + float gPad1; +}; + +static const float QD3D12_POM_DISTANCE_NEAR = 384.0; +static const float QD3D12_POM_DISTANCE_FAR = 1800.0; +static const float PI2 = 6.28318530718; + +float TeacherDeltaUVScale() +{ + return max(abs(gDeltaUVScale), 1.0e-6); +} + +float TeacherDepthAmplify() +{ + return clamp(gDepthAmplify, 0.25, 3.0); +} + +float TeacherNormalAmplify() +{ + return clamp(gNormalAmplify, 0.25, 3.0); +} + +float TeacherParallaxBaseScale() +{ + return max(gParallaxBaseScale, 0.0); +} + +float TeacherShiftAmplify() +{ + return clamp(gParallaxShiftAmplify, 1.0, 3.0); +} + +float ApplyTeacherDepthAmplify(float h) +{ + return saturate(0.5 + (saturate(h) - 0.5) * TeacherDepthAmplify()); +} + +float3 ApplyTeacherNormalAmplify(float3 n) +{ + n.xy *= TeacherNormalAmplify(); + float lenSq = dot(n, n); + return (lenSq > 1e-8) ? (n * rsqrt(lenSq)) : float3(0.0, 0.0, 1.0); +} + +uint PcgHash(uint input) +{ + uint state = input * 747796405u + 2891336453u; + uint word = ((state >> ((state >> 28u) + 4u)) ^ state) * 277803737u; + return (word >> 22u) ^ word; +} + +float Rand(inout uint s) +{ + s = PcgHash(s + 0x9E3779B9u); + return (float)(s & 0x00FFFFFFu) * (1.0 / 16777216.0); +} + +float RandRange(inout uint s, float a, float b) +{ + return a + (b - a) * Rand(s); +} + +float3 RandomHemisphereTS(inout uint rng, float minZ, float maxZ) +{ + float z = RandRange(rng, minZ, maxZ); + float a = RandRange(rng, 0.0, PI2); + float r = sqrt(max(0.0, 1.0 - z * z)); + return float3(cos(a) * r, sin(a) * r, z); +} + +float2 WrapUv(float2 uv) +{ + return frac(uv); +} + +float4 SampleAlbedo(float2 uv) +{ + return gAlbedo.SampleLevel(gSamp, WrapUv(uv), 0.0); +} + +float4 SampleNormal(float2 uv) +{ + return gNormal.SampleLevel(gSamp, WrapUv(uv), 0.0); +} + +float4 SampleSpecular(float2 uv) +{ + return gSpecular.SampleLevel(gSamp, WrapUv(uv), 0.0); +} + +float AuthoredPomAlphaWeight(float alphaValue) +{ + return (alphaValue > 0.0001 && alphaValue < 0.999) ? 1.0 : 0.0; +} + +float Luma(float3 c) +{ + return dot(saturate(c), float3(0.299, 0.587, 0.114)); +} + +float2 NormalMapTexelSize() +{ + uint w = 1; + uint h = 1; + gNormal.GetDimensions(w, h); + return rcp(max(float2((float)w, (float)h), float2(1.0, 1.0))); +} + +float GetPomHeightFromSamples(float2 uv) +{ + float4 nm = SampleNormal(uv); + float authored = AuthoredPomAlphaWeight(nm.a); + if (authored > 0.5) + return ApplyTeacherDepthAmplify(nm.a); + + float3 decoded = nm.xyz * 2.0 - 1.0; + decoded.y *= gNormalMapYSign; + + float slope = saturate(length(decoded.xy)); + float normalHeight = saturate(0.5 + (pow(slope, 0.80) - 0.35) * 0.42); + + float normalDetailGate = saturate((slope - 0.040) * 8.0); + if (normalDetailGate > 0.001) + { + float2 texel = NormalMapTexelSize(); + float lumC = Luma(SampleAlbedo(uv).rgb); + float lumL = Luma(SampleAlbedo(uv - float2(texel.x, 0.0)).rgb); + float lumR = Luma(SampleAlbedo(uv + float2(texel.x, 0.0)).rgb); + float lumU = Luma(SampleAlbedo(uv - float2(0.0, texel.y)).rgb); + float lumD = Luma(SampleAlbedo(uv + float2(0.0, texel.y)).rgb); + + float lumAvg = (lumL + lumR + lumU + lumD) * 0.25; + float localContrast = abs(lumC - lumAvg); + float diffuseWeight = saturate((localContrast - 0.018) * 16.0) * normalDetailGate; + float diffuseHeight = saturate(0.5 + (lumC - lumAvg) * 1.85); + normalHeight = lerp(normalHeight, diffuseHeight, diffuseWeight * 0.28); + } + + return ApplyTeacherDepthAmplify(normalHeight); +} + +void GetPomHeightStats(float2 uv, out float minHeight, out float maxHeight, out float authoredWeight) +{ + float2 texel = NormalMapTexelSize(); + float4 nmC = SampleNormal(uv); + float4 nmL = SampleNormal(uv - float2(texel.x, 0.0)); + float4 nmR = SampleNormal(uv + float2(texel.x, 0.0)); + float4 nmU = SampleNormal(uv - float2(0.0, texel.y)); + float4 nmD = SampleNormal(uv + float2(0.0, texel.y)); + + authoredWeight = max( + AuthoredPomAlphaWeight(nmC.a), + max(max(AuthoredPomAlphaWeight(nmL.a), AuthoredPomAlphaWeight(nmR.a)), + max(AuthoredPomAlphaWeight(nmU.a), AuthoredPomAlphaWeight(nmD.a)))); + + float hC = GetPomHeightFromSamples(uv); + float hL = GetPomHeightFromSamples(uv - float2(texel.x, 0.0)); + float hR = GetPomHeightFromSamples(uv + float2(texel.x, 0.0)); + float hU = GetPomHeightFromSamples(uv - float2(0.0, texel.y)); + float hD = GetPomHeightFromSamples(uv + float2(0.0, texel.y)); + + minHeight = min(hC, min(min(hL, hR), min(hU, hD))); + maxHeight = max(hC, max(max(hL, hR), max(hU, hD))); +} + +float GetPomDepth(float2 uv) +{ + return saturate(1.0 - GetPomHeightFromSamples(uv)); +} + +float GetPomConfidence(float2 uv) +{ + float minHeight = 0.5; + float maxHeight = 0.5; + float authoredWeight = 0.0; + GetPomHeightStats(uv, minHeight, maxHeight, authoredWeight); + + float heightRange = maxHeight - minHeight; + float threshold = lerp(0.055, 0.014, authoredWeight); + float confidence = saturate((heightRange - threshold) / max(0.18 - threshold, 0.001)); + return confidence * confidence * (3.0 - 2.0 * confidence); +} + +float3 DecodeNormalTS(float4 nm) +{ + float3 n = nm.xyz * 2.0 - 1.0; + n.y *= gNormalMapYSign; + n.xy *= max(gNormalStrength, 0.0); + float lenSq = dot(n, n); + return (lenSq > 1e-8) ? (n * rsqrt(lenSq)) : float3(0.0, 0.0, 1.0); +} + +void EvaluatePomTeacher( + float2 baseUv, + float3 viewTS, + float3 lightTS, + float distanceFade, + out float2 outUv, + out float hitDepth, + out float3 normalTS, + out float confidence, + out float visibility, + out float3 albedo, + out float3 specular) +{ + outUv = baseUv; + hitDepth = GetPomDepth(baseUv); + confidence = GetPomConfidence(baseUv); + visibility = 1.0; + + viewTS = (dot(viewTS, viewTS) > 1e-8) ? normalize(viewTS) : float3(0.0, 0.0, 1.0); + lightTS = (dot(lightTS, lightTS) > 1e-8) ? normalize(lightTS) : float3(0.0, 0.0, 1.0); + + if (confidence > 0.03) + { + float parallaxScale = clamp(gNormalStrength, 0.0, 4.0) * TeacherParallaxBaseScale() * confidence * TeacherDepthAmplify(); + float ndotv = saturate(viewTS.z); + if (ndotv > 0.035 && abs(parallaxScale) >= 1e-6) + { + float vz = max(ndotv, 0.22); + float grazingFade = smoothstep(0.08, 0.22, ndotv); + float scale = parallaxScale * distanceFade * grazingFade; + if (scale > 0.00045) + { + float layerCountF = lerp(10.0, 30.0, saturate(1.0 - ndotv)); + layerCountF = lerp(8.0, layerCountF, saturate(distanceFade)); + uint layerCount = (uint)clamp(layerCountF + 0.5, 8.0, 32.0); + + float2 parallaxVector = (viewTS.xy / vz) * scale; + float parallaxLen = length(parallaxVector); + float maxParallaxShift = min(lerp(0.024, 0.058, confidence) * TeacherShiftAmplify(), TeacherDeltaUVScale() * 0.92); + if (parallaxLen > maxParallaxShift && parallaxLen > 1e-6) + parallaxVector *= maxParallaxShift / parallaxLen; + + float invLayerCount = rcp((float)layerCount); + float2 deltaUv = parallaxVector * invLayerCount; + + float2 uv = baseUv; + float2 prevUv = uv; + float currentLayerDepth = 0.0; + float prevLayerDepth = 0.0; + float currentDepth = GetPomDepth(uv); + float prevDepth = currentDepth; + + [loop] + for (uint layer = 0u; layer < 32u; ++layer) + { + if (layer >= layerCount || currentLayerDepth >= currentDepth) + break; + + prevUv = uv; + prevLayerDepth = currentLayerDepth; + prevDepth = currentDepth; + uv -= deltaUv; + currentLayerDepth += invLayerCount; + currentDepth = GetPomDepth(uv); + } + + float afterDepth = currentDepth - currentLayerDepth; + float beforeDepth = prevDepth - prevLayerDepth; + float denom = afterDepth - beforeDepth; + float weight = (abs(denom) > 1e-5) ? saturate(afterDepth / denom) : 0.0; + float2 refinedUv = lerp(uv, prevUv, weight); + float refinedLayerDepth = lerp(currentLayerDepth, prevLayerDepth, weight); + float blend = saturate(distanceFade * grazingFade); + outUv = lerp(baseUv, refinedUv, blend); + hitDepth = saturate(lerp(hitDepth, refinedLayerDepth, blend)); + } + } + } + + outUv = frac(outUv); + normalTS = ApplyTeacherNormalAmplify(DecodeNormalTS(SampleNormal(outUv))); + albedo = saturate(SampleAlbedo(outUv).rgb); + specular = (gHasSpecular != 0u) ? saturate(SampleSpecular(outUv).rgb) : float3(0.0, 0.0, 0.0); + + if (confidence > 0.03 && lightTS.z > 0.04) + { + float vz = max(lightTS.z, 0.16); + float scale = clamp(gNormalStrength, 0.0, 4.0) * TeacherParallaxBaseScale() * confidence * TeacherDepthAmplify(); + float2 shadowVector = (lightTS.xy / vz) * scale; + float len = length(shadowVector); + float maxShift = min(lerp(0.018, 0.050, confidence) * TeacherShiftAmplify(), TeacherDeltaUVScale() * 0.84); + if (len > maxShift && len > 1e-6) + shadowVector *= maxShift / len; + + float shadow = 1.0; + const int steps = 12; + [unroll] + for (int i = 1; i <= steps; ++i) + { + float t = (float)i / (float)steps; + float2 suv = outUv + shadowVector * t; + float rayDepth = saturate(hitDepth - t * 0.55); + float sampleDepth = GetPomDepth(suv); + float blocker = smoothstep(0.015, 0.080, rayDepth - sampleDepth); + shadow = min(shadow, 1.0 - blocker * (1.0 - t * 0.35)); + } + visibility = saturate(lerp(1.0, shadow, confidence)); + } +} + +[numthreads(128, 1, 1)] +void CSMain(uint3 dispatchThreadId : SV_DispatchThreadID) +{ + uint id = dispatchThreadId.x; + if (id >= gSampleCount) + return; + + uint rng = PcgHash(gSeed ^ (id * 747796405u + 2891336453u)); + float u0 = Rand(rng); + float v0 = Rand(rng); + float2 uv = float2(u0, v0); + + float zMin = (Rand(rng) < 0.55) ? 0.08 : 0.18; + float zMax = (Rand(rng) < 0.55) ? 0.55 : 1.0; + float3 viewTS = RandomHemisphereTS(rng, zMin, zMax); + float3 lightTS = RandomHemisphereTS(rng, 0.10, 1.0); + + float viewDistance = RandRange(rng, 48.0, 2200.0); + float distanceFade = 1.0 - smoothstep(QD3D12_POM_DISTANCE_NEAR, QD3D12_POM_DISTANCE_FAR, viewDistance); + + float2 hitUv; + float hitDepth; + float3 normalTS; + float confidence; + float visibility; + float3 albedo; + float3 specular; + EvaluatePomTeacher(uv, viewTS, lightTS, distanceFade, hitUv, hitDepth, normalTS, confidence, visibility, albedo, specular); + + float2 deltaUv = hitUv - uv; + // Choose the shortest wrapped delta so the network does not learn a 0.99 jump at texture seams. + deltaUv = deltaUv - round(deltaUv); + + float invDeltaScale = rcp(TeacherDeltaUVScale()); + float target0 = clamp(deltaUv.x * invDeltaScale, -1.0, 1.0); + float target1 = clamp(deltaUv.y * invDeltaScale, -1.0, 1.0); + float target2 = saturate(hitDepth); + float target3 = clamp(normalTS.x, -1.0, 1.0); + float target4 = clamp(normalTS.y, -1.0, 1.0); + float target5 = saturate(confidence); + float target6 = saturate(visibility); + + uint base = id * 6u; + gOut[base + 0u] = float4(uv.x, uv.y, viewTS.x, viewTS.y); + gOut[base + 1u] = float4(viewTS.z, lightTS.x, lightTS.y, lightTS.z); + gOut[base + 2u] = float4(distanceFade, target0, target1, target2); + gOut[base + 3u] = float4(target3, target4, target5, target6); + gOut[base + 4u] = float4(albedo.r, albedo.g, albedo.b, specular.r); + gOut[base + 5u] = float4(specular.g, specular.b, 0.0, 0.0); +} +)NPOMHLSL"; + + struct GpuTeacherCB + { + uint32_t sampleCount; + uint32_t seed; + uint32_t hasSpecular; + uint32_t pad0; + float normalStrength; + float normalMapYSign; + float depthAmplify; + float normalAmplify; + float deltaUVScale; + float parallaxBaseScale; + float parallaxShiftAmplify; + float pad1; + }; + + struct GpuTeacher + { + GpuContext gpu; + GpuTexture albedoTex; + GpuTexture normalTex; + GpuTexture specularTex; + bool hasSpecular; + uint32_t maxBatch; + + struct BatchSlot + { + ComPtr sampleBuffer; + ComPtr readbackBuffer; + UINT64 fenceValue; + uint32_t count; + bool mapped; + + BatchSlot() : fenceValue(0), count(0), mapped(false) {} + }; + + ComPtr rootSignature; + ComPtr pso; + ComPtr heap; + UINT descriptorStride; + ComPtr constantBuffer; + BatchSlot slots[kTeacherPipelineSlots]; + uint8_t* cbMapped; + + GpuTeacher() : hasSpecular(false), maxBatch(0), descriptorStride(0), cbMapped(nullptr) {} + ~GpuTeacher() + { + UnmapReadback(); + if (constantBuffer && cbMapped) + constantBuffer->Unmap(0, nullptr); + } + + bool CompilePipeline(std::string& error) + { + ComPtr shader; + ComPtr errors; + UINT flags = D3DCOMPILE_ENABLE_STRICTNESS; +#if defined(_DEBUG) + flags |= D3DCOMPILE_DEBUG | D3DCOMPILE_SKIP_OPTIMIZATION; +#endif + HRESULT hr = D3DCompile( + gGpuTeacherHLSL, + strlen(gGpuTeacherHLSL), + "QD3D12NeuralPOMTeacherCS", + nullptr, + nullptr, + "CSMain", + "cs_5_0", + flags, + 0, + &shader, + &errors); + if (FAILED(hr)) + { + if (errors) + error.assign((const char*)errors->GetBufferPointer(), errors->GetBufferSize()); + else + error = "D3DCompile failed for Neural POM teacher compute shader"; + return false; + } + + D3D12_DESCRIPTOR_RANGE ranges[2] = {}; + ranges[0].RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_SRV; + ranges[0].NumDescriptors = 3; + ranges[0].BaseShaderRegister = 0; + ranges[0].RegisterSpace = 0; + ranges[0].OffsetInDescriptorsFromTableStart = 0; + + ranges[1].RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_UAV; + ranges[1].NumDescriptors = 1; + ranges[1].BaseShaderRegister = 0; + ranges[1].RegisterSpace = 0; + ranges[1].OffsetInDescriptorsFromTableStart = 3; + + D3D12_ROOT_PARAMETER params[2] = {}; + params[0].ParameterType = D3D12_ROOT_PARAMETER_TYPE_CBV; + params[0].Descriptor.ShaderRegister = 0; + params[0].Descriptor.RegisterSpace = 0; + params[0].ShaderVisibility = D3D12_SHADER_VISIBILITY_ALL; + + params[1].ParameterType = D3D12_ROOT_PARAMETER_TYPE_DESCRIPTOR_TABLE; + params[1].DescriptorTable.NumDescriptorRanges = 2; + params[1].DescriptorTable.pDescriptorRanges = ranges; + params[1].ShaderVisibility = D3D12_SHADER_VISIBILITY_ALL; + + D3D12_STATIC_SAMPLER_DESC sampler = {}; + sampler.Filter = D3D12_FILTER_MIN_MAG_MIP_LINEAR; + sampler.AddressU = D3D12_TEXTURE_ADDRESS_MODE_WRAP; + sampler.AddressV = D3D12_TEXTURE_ADDRESS_MODE_WRAP; + sampler.AddressW = D3D12_TEXTURE_ADDRESS_MODE_WRAP; + sampler.MipLODBias = 0.0f; + sampler.MaxAnisotropy = 1; + sampler.ComparisonFunc = D3D12_COMPARISON_FUNC_ALWAYS; + sampler.BorderColor = D3D12_STATIC_BORDER_COLOR_OPAQUE_WHITE; + sampler.MinLOD = 0.0f; + sampler.MaxLOD = D3D12_FLOAT32_MAX; + sampler.ShaderRegister = 0; + sampler.RegisterSpace = 0; + sampler.ShaderVisibility = D3D12_SHADER_VISIBILITY_ALL; + + D3D12_ROOT_SIGNATURE_DESC rs = {}; + rs.NumParameters = 2; + rs.pParameters = params; + rs.NumStaticSamplers = 1; + rs.pStaticSamplers = &sampler; + rs.Flags = D3D12_ROOT_SIGNATURE_FLAG_NONE; + + ComPtr rsBlob; + ComPtr rsError; + hr = D3D12SerializeRootSignature(&rs, D3D_ROOT_SIGNATURE_VERSION_1, &rsBlob, &rsError); + if (FAILED(hr)) + { + if (rsError) + error.assign((const char*)rsError->GetBufferPointer(), rsError->GetBufferSize()); + else + error = "D3D12SerializeRootSignature failed"; + return false; + } + + hr = gpu.device->CreateRootSignature(0, rsBlob->GetBufferPointer(), rsBlob->GetBufferSize(), IID_PPV_ARGS(&rootSignature)); + if (FAILED(hr)) + { + error = "CreateRootSignature failed"; + return false; + } + + D3D12_COMPUTE_PIPELINE_STATE_DESC ps = {}; + ps.pRootSignature = rootSignature.Get(); + ps.CS.pShaderBytecode = shader->GetBufferPointer(); + ps.CS.BytecodeLength = shader->GetBufferSize(); + hr = gpu.device->CreateComputePipelineState(&ps, IID_PPV_ARGS(&pso)); + if (FAILED(hr)) + { + error = "CreateComputePipelineState failed"; + return false; + } + return true; + } + + D3D12_CPU_DESCRIPTOR_HANDLE CpuHandle(UINT index) const + { + D3D12_CPU_DESCRIPTOR_HANDLE h = heap->GetCPUDescriptorHandleForHeapStart(); + h.ptr += SIZE_T(index) * SIZE_T(descriptorStride); + return h; + } + + D3D12_GPU_DESCRIPTOR_HANDLE GpuHandle(UINT index) const + { + D3D12_GPU_DESCRIPTOR_HANDLE h = heap->GetGPUDescriptorHandleForHeapStart(); + h.ptr += UINT64(index) * UINT64(descriptorStride); + return h; + } + + bool CreateDescriptors(std::string& error) + { + D3D12_DESCRIPTOR_HEAP_DESC hd = {}; + hd.NumDescriptors = kTeacherPipelineSlots * 4u; + hd.Type = D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV; + hd.Flags = D3D12_DESCRIPTOR_HEAP_FLAG_SHADER_VISIBLE; + if (FAILED(gpu.device->CreateDescriptorHeap(&hd, IID_PPV_ARGS(&heap)))) + { + error = "CreateDescriptorHeap failed"; + return false; + } + descriptorStride = gpu.device->GetDescriptorHandleIncrementSize(D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV); + + auto CreateSRV = [&](ID3D12Resource* res, UINT index) + { + D3D12_SHADER_RESOURCE_VIEW_DESC sd = {}; + sd.ViewDimension = D3D12_SRV_DIMENSION_TEXTURE2D; + sd.Format = DXGI_FORMAT_R8G8B8A8_UNORM; + sd.Shader4ComponentMapping = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING; + sd.Texture2D.MipLevels = 1; + gpu.device->CreateShaderResourceView(res, &sd, CpuHandle(index)); + }; + + D3D12_UNORDERED_ACCESS_VIEW_DESC uav = {}; + uav.ViewDimension = D3D12_UAV_DIMENSION_BUFFER; + uav.Format = DXGI_FORMAT_UNKNOWN; + uav.Buffer.FirstElement = 0; + uav.Buffer.NumElements = maxBatch * kPackedSampleFloat4s; + uav.Buffer.StructureByteStride = 16; + uav.Buffer.CounterOffsetInBytes = 0; + uav.Buffer.Flags = D3D12_BUFFER_UAV_FLAG_NONE; + + for (uint32_t slotIndex = 0; slotIndex < kTeacherPipelineSlots; ++slotIndex) + { + const UINT base = slotIndex * 4u; + CreateSRV(albedoTex.resource.Get(), base + 0u); + CreateSRV(normalTex.resource.Get(), base + 1u); + CreateSRV(specularTex.resource.Get(), base + 2u); + gpu.device->CreateUnorderedAccessView(slots[slotIndex].sampleBuffer.Get(), nullptr, &uav, CpuHandle(base + 3u)); + } + return true; + } + + bool Init( + ID3D12Device* device, + ID3D12CommandQueue* queue, + const ImageRGBA8& albedo, + const ImageRGBA8& normal, + const ImageRGBA8* specular, + uint32_t inMaxBatch, + std::string& error) + { + maxBatch = std::max(1u, inMaxBatch); + hasSpecular = specular && specular->Valid(); + + if (!gpu.Init(device, queue)) + { + error = "failed to initialize trainer GPU command context"; + return false; + } + + if (!CompilePipeline(error)) + return false; + + if (!gpu.Begin()) + { + error = "failed to begin trainer upload command list"; + return false; + } + + std::vector> uploadKeepAlive; + if (!UploadImageAsTexture(gpu, albedo, albedoTex, uploadKeepAlive)) + { + error = "failed to upload albedo texture to D3D12"; + return false; + } + if (!UploadImageAsTexture(gpu, normal, normalTex, uploadKeepAlive)) + { + error = "failed to upload normal texture to D3D12"; + return false; + } + if (hasSpecular) + { + if (!UploadImageAsTexture(gpu, *specular, specularTex, uploadKeepAlive)) + { + error = "failed to upload specular texture to D3D12"; + return false; + } + } + else + { + ImageRGBA8 black; + black.width = 1; + black.height = 1; + black.rgba = { 0, 0, 0, 255 }; + if (!UploadImageAsTexture(gpu, black, specularTex, uploadKeepAlive)) + { + error = "failed to upload default black specular texture to D3D12"; + return false; + } + } + + if (!gpu.EndWait()) + { + error = "failed to submit trainer texture uploads"; + return false; + } + + const UINT64 cbBytes = AlignUp(sizeof(GpuTeacherCB), 256); + if (!CreateCommittedBuffer(gpu.device.Get(), cbBytes, D3D12_HEAP_TYPE_UPLOAD, D3D12_RESOURCE_STATE_GENERIC_READ, D3D12_RESOURCE_FLAG_NONE, constantBuffer)) + { + error = "failed to create teacher constant buffer"; + return false; + } + D3D12_RANGE readRange = {}; + if (FAILED(constantBuffer->Map(0, &readRange, (void**)&cbMapped)) || !cbMapped) + { + error = "failed to map teacher constant buffer"; + return false; + } + + const UINT64 sampleBytes = (UINT64)maxBatch * (UINT64)kPackedSampleBytes; + for (uint32_t slotIndex = 0; slotIndex < kTeacherPipelineSlots; ++slotIndex) + { + BatchSlot& slot = slots[slotIndex]; + if (!CreateCommittedBuffer(gpu.device.Get(), sampleBytes, D3D12_HEAP_TYPE_DEFAULT, D3D12_RESOURCE_STATE_UNORDERED_ACCESS, D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS, slot.sampleBuffer)) + { + error = "failed to create teacher sample UAV buffer"; + return false; + } + if (!CreateCommittedBuffer(gpu.device.Get(), sampleBytes, D3D12_HEAP_TYPE_READBACK, D3D12_RESOURCE_STATE_COPY_DEST, D3D12_RESOURCE_FLAG_NONE, slot.readbackBuffer)) + { + error = "failed to create teacher readback buffer"; + return false; + } + } + + if (!CreateDescriptors(error)) + return false; + + return true; + } + + bool GenerateBatchAsync(uint32_t slotIndex, uint32_t count, uint32_t seed, float normalStrength, float normalYSign, std::string& error) + { + if (slotIndex >= kTeacherPipelineSlots) + { + error = "invalid teacher batch slot"; + return false; + } + if (count == 0 || count > maxBatch) + { + error = "invalid teacher batch count"; + return false; + } + + BatchSlot& slot = slots[slotIndex]; + UnmapReadback(slotIndex); + + GpuTeacherCB cb = {}; + cb.sampleCount = count; + cb.seed = seed; + cb.hasSpecular = hasSpecular ? 1u : 0u; + cb.normalStrength = normalStrength; + cb.normalMapYSign = normalYSign; + cb.depthAmplify = kTrainingDepthAmplify; + cb.normalAmplify = kTrainingNormalAmplify; + cb.deltaUVScale = kDeltaUVScale; + cb.parallaxBaseScale = kTeacherParallaxBaseScale; + cb.parallaxShiftAmplify = kTeacherParallaxShiftAmplify; + memcpy(cbMapped, &cb, sizeof(cb)); + + if (!gpu.Begin()) + { + error = "failed to begin teacher dispatch command list"; + return false; + } + + ID3D12DescriptorHeap* heaps[] = { heap.Get() }; + gpu.list->SetDescriptorHeaps(1, heaps); + gpu.list->SetComputeRootSignature(rootSignature.Get()); + gpu.list->SetPipelineState(pso.Get()); + gpu.list->SetComputeRootConstantBufferView(0, constantBuffer->GetGPUVirtualAddress()); + gpu.list->SetComputeRootDescriptorTable(1, GpuHandle(slotIndex * 4u)); + + const UINT groups = (count + 127u) / 128u; + gpu.list->Dispatch(groups, 1, 1); + gpu.UavBarrier(slot.sampleBuffer.Get()); + gpu.Transition(slot.sampleBuffer.Get(), D3D12_RESOURCE_STATE_UNORDERED_ACCESS, D3D12_RESOURCE_STATE_COPY_SOURCE); + gpu.list->CopyBufferRegion(slot.readbackBuffer.Get(), 0, slot.sampleBuffer.Get(), 0, (UINT64)count * (UINT64)kPackedSampleBytes); + gpu.Transition(slot.sampleBuffer.Get(), D3D12_RESOURCE_STATE_COPY_SOURCE, D3D12_RESOURCE_STATE_UNORDERED_ACCESS); + + UINT64 fenceValue = 0; + if (!gpu.EndSignal(&fenceValue)) + { + error = "failed to submit teacher dispatch"; + return false; + } + + slot.fenceValue = fenceValue; + slot.count = count; + return true; + } + + bool WaitBatch(uint32_t slotIndex, const float** outRecords, uint32_t* outCount, std::string& error) + { + if (!outRecords) + return false; + *outRecords = nullptr; + if (outCount) + *outCount = 0; + if (slotIndex >= kTeacherPipelineSlots) + { + error = "invalid teacher batch slot"; + return false; + } + + BatchSlot& slot = slots[slotIndex]; + if (!gpu.WaitFence(slot.fenceValue)) + { + error = "failed to wait for teacher dispatch"; + return false; + } + + void* mapped = nullptr; + D3D12_RANGE readRange; + readRange.Begin = 0; + readRange.End = (SIZE_T)slot.count * (SIZE_T)kPackedSampleBytes; + if (FAILED(slot.readbackBuffer->Map(0, &readRange, &mapped)) || !mapped) + { + error = "failed to map teacher readback buffer"; + return false; + } + slot.mapped = true; + *outRecords = (const float*)mapped; + if (outCount) + *outCount = slot.count; + return true; + } + + bool GenerateBatch(uint32_t count, uint32_t seed, float normalStrength, float normalYSign, const float** outRecords, std::string& error) + { + if (!GenerateBatchAsync(0u, count, seed, normalStrength, normalYSign, error)) + return false; + uint32_t actualCount = 0; + return WaitBatch(0u, outRecords, &actualCount, error); + } + + void UnmapReadback(uint32_t slotIndex) + { + if (slotIndex >= kTeacherPipelineSlots) + return; + BatchSlot& slot = slots[slotIndex]; + if (slot.readbackBuffer && slot.mapped) + { + D3D12_RANGE noWrite = { 0, 0 }; + slot.readbackBuffer->Unmap(0, &noWrite); + slot.mapped = false; + } + } + + void UnmapReadback() + { + for (uint32_t slotIndex = 0; slotIndex < kTeacherPipelineSlots; ++slotIndex) + UnmapReadback(slotIndex); + } + }; + + static float GetPackedTarget(const float* packedSample, int outputIndex) + { + return packedSample[9 + outputIndex]; + } + + struct TrainScratch + { + std::vector latentSample; + std::vector input; + std::vector h1; + std::vector h2; + std::vector y; + std::vector dy; + std::vector dx; + std::vector dz3; + std::vector dh2; + std::vector dz2; + std::vector dh1; + std::vector dz1; + std::vector dLatent; + int latentIdx[4]; + float latentW[4]; + + void Init(int inputCount, int hidden, int outputCount, int latentChannels) + { + latentSample.resize(latentChannels); + input.resize(inputCount); + h1.resize(hidden); + h2.resize(hidden); + y.resize(outputCount); + dy.resize(outputCount); + dx.resize(inputCount); + dz3.resize(outputCount); + dh2.resize(hidden); + dz2.resize(hidden); + dh1.resize(hidden); + dz1.resize(hidden); + dLatent.resize(latentChannels); + memset(latentIdx, 0, sizeof(latentIdx)); + memset(latentW, 0, sizeof(latentW)); + } + }; + + static float TrainingLearningRate(const QD3D12NeuralPOMTrainDesc& opt, uint32_t sampleStep) + { + float lr = 0.0012f; + if (sampleStep > opt.samples * 2u / 3u) lr = 0.00035f; + else if (sampleStep > opt.samples / 3u) lr = 0.0007f; + return lr; + } + + static void TrainPackedGpuSample( + const float* packedSample, + const QD3D12NeuralPOMTrainDesc& opt, + MLP& net, + LatentGrid& latent, + int step, + float& emaLoss, + TrainScratch& scratch) + { + const float uvx = packedSample[0]; + const float uvy = packedSample[1]; + const float viewX = packedSample[2]; + const float viewY = packedSample[3]; + const float viewZ = packedSample[4]; + const float lightX = packedSample[5]; + const float lightY = packedSample[6]; + const float lightZ = packedSample[7]; + const float distanceFade = packedSample[8]; + + latent.Sample(uvx, uvy, scratch.latentSample, scratch.latentIdx, scratch.latentW); + + if (scratch.input.size() != (size_t)net.inputCount) scratch.input.resize((size_t)net.inputCount); + for (int c = 0; c < latent.channels; ++c) + scratch.input[c] = scratch.latentSample[c]; + const int base = latent.channels; + scratch.input[base + 0] = viewX; + scratch.input[base + 1] = viewY; + scratch.input[base + 2] = viewZ; + scratch.input[base + 3] = lightX; + scratch.input[base + 4] = lightY; + scratch.input[base + 5] = lightZ; + scratch.input[base + 6] = ClampF(opt.normalStrength / 4.0f, 0.0f, 1.0f); + scratch.input[base + 7] = distanceFade; + scratch.input[base + 8] = 0.0f; // mip/ray-cone placeholder for the runtime path. + + net.Forward(scratch.input, scratch.h1, scratch.h2, scratch.y); + + static const float lossW[kOutputCount] = + { + 10.0f, 10.0f, // deeper delta UV needs tighter fit + 2.5f, // depth + 3.5f, 3.5f, // stronger learned normal xy + 0.8f, // confidence + 1.5f, // visibility + 0.25f, 0.25f, 0.25f, + 0.20f, 0.20f, 0.20f + }; + + if (scratch.dy.size() != (size_t)kOutputCount) scratch.dy.resize((size_t)kOutputCount); + float loss = 0.0f; + float grazingWeight = 1.0f / max(viewZ, 0.16f); + grazingWeight = ClampF(grazingWeight, 1.0f, 4.0f); + for (int i = 0; i < kOutputCount; ++i) + { + float w = lossW[i]; + if (i <= 4) + w *= grazingWeight; + const float target = GetPackedTarget(packedSample, i); + const float d = scratch.y[i] - target; + loss += w * d * d; + scratch.dy[i] = 2.0f * w * d; + } + emaLoss = (step == 1) ? loss : (emaLoss * 0.995f + loss * 0.005f); + + net.BackwardAccumulate( + scratch.input, + scratch.h1, + scratch.h2, + scratch.y, + scratch.dy, + scratch.dx, + scratch.dz3, + scratch.dh2, + scratch.dz2, + scratch.dh1, + scratch.dz1); + + for (int c = 0; c < latent.channels; ++c) + scratch.dLatent[c] = scratch.dx[c]; + + // Accumulate sparse latent gradients and apply Adam once per CPU optimizer + // minibatch. This removes per-sample Adam/powf overhead and is much more + // cache-friendly for the 128x128xC latent texture. + latent.AccumulateGrad(scratch.latentIdx, scratch.latentW, scratch.dLatent); + } + + static void ApplyDefaults(QD3D12NeuralPOMTrainDesc& d) + { + if (d.samples == 0) d.samples = 300000; + if (d.samplesPerGpuBatch == 0) d.samplesPerGpuBatch = 32768; + if (d.latentResolution == 0) d.latentResolution = 128; + if (d.latentChannels == 0) d.latentChannels = 8; + if (d.hiddenCount == 0) d.hiddenCount = 48; + if (d.randomSeed == 0) d.randomSeed = 0xC001D00Du; + if (d.normalStrength == 0.0f) d.normalStrength = 1.0f; + if (d.normalMapYSign == 0.0f) d.normalMapYSign = 1.0f; + + d.samples = std::max(1000u, d.samples); + d.samplesPerGpuBatch = std::max(256u, std::min(d.samplesPerGpuBatch, 65536u)); + d.latentResolution = std::max(16u, std::min(d.latentResolution, 512u)); + d.latentChannels = std::max(4u, std::min(d.latentChannels, 16u)); + d.hiddenCount = std::max(16u, std::min(d.hiddenCount, 128u)); + d.normalStrength = ClampF(d.normalStrength, 0.0f, 4.0f); + d.normalMapYSign = (d.normalMapYSign < 0.0f) ? -1.0f : 1.0f; + } + + static bool LoadImagesFromDesc(const QD3D12NeuralPOMTrainDesc& desc, ImageRGBA8& albedo, ImageRGBA8& normal, ImageRGBA8& specular, bool& hasSpecular, QD3D12NeuralPOMTrainStats* stats) + { + hasSpecular = false; + + if (desc.albedoImage) + { + if (!albedo.FromExternal(*desc.albedoImage)) + return Fail(stats, "invalid albedoImage RGBA8 input") != 0; + } + else + { + if (!LoadTGAFile(desc.albedoPath, albedo)) + return Fail(stats, "failed to load albedo TGA '%s'", desc.albedoPath ? desc.albedoPath : "") != 0; + } + + if (desc.normalImage) + { + if (!normal.FromExternal(*desc.normalImage)) + return Fail(stats, "invalid normalImage RGBA8 input") != 0; + } + else + { + if (!LoadTGAFile(desc.normalPath, normal)) + return Fail(stats, "failed to load normal TGA '%s'", desc.normalPath ? desc.normalPath : "") != 0; + } + + if (desc.specularImage) + { + hasSpecular = specular.FromExternal(*desc.specularImage); + if (!hasSpecular) + return Fail(stats, "invalid specularImage RGBA8 input") != 0; + } + else if (!StringIsEmptyOrDash(desc.specularPath)) + { + hasSpecular = LoadTGAFile(desc.specularPath, specular); + if (!hasSpecular) + Log("QD3D12 NeuralPOM: specular TGA '%s' not found; training zero specular", desc.specularPath); + } + + if (!albedo.Valid() || !normal.Valid()) + return Fail(stats, "invalid source images") != 0; + return true; + } + + static int TrainImpl(const QD3D12NeuralPOMTrainDesc* inDesc, QD3D12NeuralPOMTrainStats* stats) + { + if (stats) + { + memset(stats, 0, sizeof(*stats)); + stats->size = sizeof(*stats); + } + + if (!inDesc) + return Fail(stats, "null QD3D12NeuralPOMTrainDesc"); + + QD3D12NeuralPOMTrainDesc desc = *inDesc; + ApplyDefaults(desc); + + if (!desc.outputPrefix || !desc.outputPrefix[0]) + return Fail(stats, "missing outputPrefix"); + + ImageRGBA8 albedo; + ImageRGBA8 normal; + ImageRGBA8 specular; + bool hasSpecular = false; + if (!LoadImagesFromDesc(desc, albedo, normal, specular, hasSpecular, stats)) + return 0; + + ID3D12Device* device = QD3D12_GetDevice(); + ID3D12CommandQueue* queue = QD3D12_GetQueue(); + if (!device || !queue) + return Fail(stats, "QD3D12 device/queue are not initialized; call after shim D3D12 init"); + + // Drain the shim's open frame/queue before borrowing the shared queue for long trainer work. + QD3D12_WaitForGPU_External(); + + std::string gpuError; + GpuTeacher teacher; + if (!teacher.Init(device, queue, albedo, normal, hasSpecular ? &specular : nullptr, desc.samplesPerGpuBatch, gpuError)) + return Fail(stats, "D3D12 teacher init failed: %s", gpuError.c_str()); + + Rng rng(desc.randomSeed); + MLP net; + LatentGrid latent; + const int inputCount = (int)desc.latentChannels + 9; + net.Init(inputCount, (int)desc.hiddenCount, kOutputCount, rng); + latent.Init((int)desc.latentResolution, (int)desc.latentChannels, rng); + + float emaLoss = 0.0f; + const uint32_t total = desc.samples; + const uint32_t reportEvery = std::max(1u, total / 20u); + uint32_t trained = 0; + uint32_t gpuBatches = 0; + uint32_t netAccumulatedSamples = 0; + uint32_t netOptimizerSteps = 0; + TrainScratch scratch; + scratch.Init(net.inputCount, net.hidden, net.outputCount, latent.channels); + net.ZeroGrad(); + latent.BeginGradBatch(); + + Progress(desc, "neuralPOM D3D12 trainer started", 0, total, 0.0f); + + uint32_t scheduled = 0; + auto ScheduleTeacherBatch = [&](uint32_t slotIndex) -> bool + { + if (scheduled >= total) + return false; + const uint32_t batchCount = std::min(desc.samplesPerGpuBatch, total - scheduled); + const uint32_t seed = desc.randomSeed ^ (scheduled * 1664525u + gpuBatches * 1013904223u); + if (!teacher.GenerateBatchAsync(slotIndex, batchCount, seed, desc.normalStrength, desc.normalMapYSign, gpuError)) + return false; + scheduled += batchCount; + ++gpuBatches; + return true; + }; + + uint32_t readSlot = 0u; + if (!ScheduleTeacherBatch(readSlot)) + return Fail(stats, "D3D12 teacher dispatch failed: %s", gpuError.c_str()); + + while (trained < total) + { + const float* packed = nullptr; + uint32_t batchCount = 0; + if (!teacher.WaitBatch(readSlot, &packed, &batchCount, gpuError)) + return Fail(stats, "D3D12 teacher dispatch failed: %s", gpuError.c_str()); + + const uint32_t nextSlot = (readSlot + 1u) % kTeacherPipelineSlots; + const bool shouldScheduleNext = scheduled < total; + if (shouldScheduleNext && !ScheduleTeacherBatch(nextSlot)) + return Fail(stats, "D3D12 teacher dispatch failed: %s", gpuError.c_str()); + + for (uint32_t i = 0; i < batchCount; ++i) + { + ++trained; + const float* sample = packed + (size_t)i * (size_t)(kPackedSampleBytes / sizeof(float)); + TrainPackedGpuSample(sample, desc, net, latent, (int)trained, emaLoss, scratch); + ++netAccumulatedSamples; + + if (netAccumulatedSamples >= kCpuOptimizerBatch || trained == total) + { + ++netOptimizerSteps; + const float lr = TrainingLearningRate(desc, trained); + const float gradScale = 1.0f / (float)std::max(1u, netAccumulatedSamples); + net.ApplyAdamScaled(lr, (int)netOptimizerSteps, gradScale, 1e-7f); + latent.ApplyAdamSparse(lr * 0.65f, (int)netOptimizerSteps, gradScale, 1e-6f); + net.ZeroGrad(); + latent.BeginGradBatch(); + netAccumulatedSamples = 0; + } + + if (trained == 1u || trained % reportEvery == 0u || trained == total) + Progress(desc, "neuralPOM D3D12 trainer", trained, total, emaLoss); + } + + teacher.UnmapReadback(readSlot); + readSlot = nextSlot; + } + + const std::string manifestPath = WithSuffix(desc.outputPrefix, "_manifest.json"); + const std::string weightsPath = WithSuffix(desc.outputPrefix, "_weights.bin"); + const std::string latentF32Path = WithSuffix(desc.outputPrefix, "_latent_f32.bin"); + const std::string latentHalfPath = WithSuffix(desc.outputPrefix, "_latent_rgba16f.bin"); + + bool ok = true; + ok = SaveWeights(weightsPath.c_str(), net, latent) && ok; + ok = SaveLatentF32(latentF32Path.c_str(), latent) && ok; + ok = SaveLatentRGBA16F(latentHalfPath.c_str(), latent) && ok; + + const char* albedoSource = desc.albedoImage ? "" : desc.albedoPath; + const char* normalSource = desc.normalImage ? "" : desc.normalPath; + const char* specSource = desc.specularImage ? "" : (hasSpecular ? desc.specularPath : ""); + const std::string manifest = BuildManifest( + desc.outputPrefix, + albedoSource, + normalSource, + specSource, + net, + latent, + desc.samples, + desc.samplesPerGpuBatch, + desc.normalStrength, + desc.normalMapYSign, + emaLoss); + ok = WriteFileString(manifestPath.c_str(), manifest) && ok; + + if (!ok) + return Fail(stats, "training finished but one or more output files could not be written"); + + if (stats) + { + stats->samplesTrained = trained; + stats->gpuBatches = gpuBatches; + stats->finalMeanLoss = emaLoss; + CopyPath(stats->manifestPath, sizeof(stats->manifestPath), manifestPath); + CopyPath(stats->weightsPath, sizeof(stats->weightsPath), weightsPath); + CopyPath(stats->latentF32Path, sizeof(stats->latentF32Path), latentF32Path); + CopyPath(stats->latentRGBA16FPath, sizeof(stats->latentRGBA16FPath), latentHalfPath); + } + + Progress(desc, "neuralPOM D3D12 trainer finished", total, total, emaLoss); + QD3D12_WaitForGPU_External(); + return 1; + } + +} // namespace QD3D12NeuralPOM + +QD3D12_NEURAL_POM_API int QD3D12_NeuralPOMTrainMaterialD3D12( + const QD3D12NeuralPOMTrainDesc* desc, + QD3D12NeuralPOMTrainStats* stats) +{ + return QD3D12NeuralPOM::TrainImpl(desc, stats); +} + +QD3D12_NEURAL_POM_API int QD3D12_NeuralPOMTrainMaterialFromFilesD3D12( + const char* albedoTgaPath, + const char* normalTgaPath, + const char* specularTgaPathOrNull, + const char* outputPrefix, + uint32_t samples, + uint32_t latentResolution, + uint32_t hiddenCount, + float normalStrength, + float normalMapYSign, + QD3D12NeuralPOMTrainStats* stats) +{ + QD3D12NeuralPOMTrainDesc d = {}; + d.size = sizeof(d); + d.albedoPath = albedoTgaPath; + d.normalPath = normalTgaPath; + d.specularPath = specularTgaPathOrNull; + strcpy(d.outputPrefix, outputPrefix); + d.samples = samples; + d.samplesPerGpuBatch = 32768; + d.latentResolution = latentResolution; + d.latentChannels = 8; + d.hiddenCount = hiddenCount; + d.randomSeed = 0xC001D00Du; + d.flags = QD3D12_NEURAL_POM_TRAIN_FLAG_VERBOSE; + d.normalStrength = normalStrength; + d.normalMapYSign = normalMapYSign; + return QD3D12_NeuralPOMTrainMaterialD3D12(&d, stats); +} diff --git a/neo/engine/opengl/opengl_nn.h b/neo/engine/opengl/opengl_nn.h new file mode 100644 index 00000000..cced9570 --- /dev/null +++ b/neo/engine/opengl/opengl_nn.h @@ -0,0 +1,87 @@ +#pragma once +// ----------------------------------------------------------------------------- +// Shim entry points supplied by gl_d3d12shim.cpp. +// These have C++ linkage in the shim source, so do not wrap them in extern "C". +// ----------------------------------------------------------------------------- +extern ID3D12Device* QD3D12_GetDevice(void); +extern ID3D12CommandQueue* QD3D12_GetQueue(void); +extern void QD3D12_WaitForGPU_External(void); + +#ifndef QD3D12_NEURAL_POM_API +#define QD3D12_NEURAL_POM_API +#endif + +#ifndef QD3D12_NEURAL_POM_MAX_PATH_CHARS +#define QD3D12_NEURAL_POM_MAX_PATH_CHARS 520 +#endif + +#define QD3D12_NEURAL_POM_TRAIN_FLAG_NONE 0u +#define QD3D12_NEURAL_POM_TRAIN_FLAG_VERBOSE 1u +#define QD3D12_NEURAL_POM_TRAIN_FLAG_RESERVED_FORCE_CPU 2u +#define QD3D12_NEURAL_POM_TRAIN_FLAG_RESERVED_ALLOW_FALLBACK 4u + +typedef void(__cdecl* QD3D12NeuralPOMProgressCallback)( + const char* message, + uint32_t currentSample, + uint32_t totalSamples, + float meanLoss, + void* userData); + +struct QD3D12NeuralPOMImageRGBA8 +{ + uint32_t width; + uint32_t height; + uint32_t rowPitchBytes; + const void* pixelsRGBA8; +}; + +struct QD3D12NeuralPOMTrainDesc +{ + uint32_t size; + + // File path path. Used when the matching image pointer below is null. + const char* albedoPath; + const char* normalPath; + const char* specularPath; // null, empty, or "-" means no authored specular. + char outputPrefix[512]; + + // Optional direct image data path. This is the recommended path when calling + // from gl_d3d12shim.cpp because TextureResource::sysmem is already RGBA8. + const QD3D12NeuralPOMImageRGBA8* albedoImage; + const QD3D12NeuralPOMImageRGBA8* normalImage; + const QD3D12NeuralPOMImageRGBA8* specularImage; + + uint32_t samples; // default 300000 + uint32_t samplesPerGpuBatch; // default 8192 + uint32_t latentResolution; // default 128 + uint32_t latentChannels; // default 8 + uint32_t hiddenCount; // default 48 + uint32_t randomSeed; // default 0xC001D00D + uint32_t flags; + + float normalStrength; // default 1.0, clamped 0..4 + float normalMapYSign; // default +1, use -1 if your normal map green channel is inverted + + QD3D12NeuralPOMProgressCallback progress; + void* progressUserData; +}; + +struct QD3D12NeuralPOMTrainStats +{ + uint32_t size; + uint32_t samplesTrained; + uint32_t gpuBatches; + float finalMeanLoss; + char manifestPath[QD3D12_NEURAL_POM_MAX_PATH_CHARS]; + char weightsPath[QD3D12_NEURAL_POM_MAX_PATH_CHARS]; + char latentF32Path[QD3D12_NEURAL_POM_MAX_PATH_CHARS]; + char latentRGBA16FPath[QD3D12_NEURAL_POM_MAX_PATH_CHARS]; + char error[1024]; +}; + +QD3D12_NEURAL_POM_API int QD3D12_NeuralPOMTrainMaterialD3D12( + const QD3D12NeuralPOMTrainDesc* desc, + QD3D12NeuralPOMTrainStats* stats); + +void APIENTRY glNeuralPOMMaterialQD3D12(GLuint texture, GLsizei weightsBytes, const GLvoid* weightsData, GLsizei latentBytes, const GLvoid* latentRGBA16FData); +void APIENTRY glBindNeuralPOMTextureQD3D12(GLuint texture); \ No newline at end of file