id-software/DoomRTX
1
0
Fork 0
mirror of https://github.com/jmarshall23/DoomRTX.git synced 2026-08-12 08:11:10 +02:00
Code Issues Packages Projects Releases Wiki Activity
Files
main
DoomRTX/neo/engine/opengl/nn_mattrain.cpp
T
Justin Marshall 1c1f828271 Added nn trainer.
2026-05-18 07:33:20 -07:00

2389 lines
72 KiB
C++
Raw Permalink Blame History

/*
===========================================================================
IceTech / QD3D12 Neural POM Material Trainer
Copyright (C) 2026 Justin Marshall
===========================================================================
*/
#include <windows.h>
#include <d3d12.h>
#include <d3dcompiler.h>
#include <wrl/client.h>
#include <stdint.h>
#include <stdio.h>
#include <stdarg.h>
#include <math.h>
#include <float.h>
#include <string.h>
#include <algorithm>
#include <string>
#include <vector>
#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<uint8_t>& 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<uint8_t> 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<uint8_t> 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<float> m;
std::vector<float> 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<float> w1, b1, w2, b2, w3, b3;
std::vector<float> 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<float>& x, std::vector<float>& h1, std::vector<float>& h2, std::vector<float>& 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<float>& x, const std::vector<float>& h1, const std::vector<float>& h2,
const std::vector<float>& y, const std::vector<float>& dy, std::vector<float>& dx,
std::vector<float>& dz3, std::vector<float>& dh2, std::vector<float>& dz2,
std::vector<float>& dh1, std::vector<float>& 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<float>& x, const std::vector<float>& h1, const std::vector<float>& h2,
const std::vector<float>& y, const std::vector<float>& dy, std::vector<float>& dx)
{
ZeroGrad();
std::vector<float> 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<float> data;
std::vector<float> adamM;
std::vector<float> 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<float> grad;
std::vector<uint32_t> gradStamp;
std::vector<uint32_t> 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<uint32_t>(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<float>& 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<float>& 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<float>& 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<uint8_t>& 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<uint8_t>& out, const std::vector<float>& 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<uint8_t> 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<uint16_t> 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<ID3D12Resource>& 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<ID3D12Device> device;
ComPtr<ID3D12CommandQueue> queue;
ComPtr<ID3D12CommandAllocator> allocator;
ComPtr<ID3D12GraphicsCommandList> list;
ComPtr<ID3D12Fence> 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<ID3D12Resource> resource;
UINT width;
UINT height;
};
static bool UploadImageAsTexture(GpuContext& gpu, const ImageRGBA8& img, GpuTexture& out, std::vector<ComPtr<ID3D12Resource>>& 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<ID3D12Resource> 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<float4> gAlbedo : register(t0);
Texture2D<float4> gNormal : register(t1);
Texture2D<float4> gSpecular : register(t2);
SamplerState gSamp : register(s0);
RWStructuredBuffer<float4> 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<ID3D12Resource> sampleBuffer;
ComPtr<ID3D12Resource> readbackBuffer;
UINT64 fenceValue;
uint32_t count;
bool mapped;
BatchSlot() : fenceValue(0), count(0), mapped(false) {}
};
ComPtr<ID3D12RootSignature> rootSignature;
ComPtr<ID3D12PipelineState> pso;
ComPtr<ID3D12DescriptorHeap> heap;
UINT descriptorStride;
ComPtr<ID3D12Resource> 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<ID3DBlob> shader;
ComPtr<ID3DBlob> 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<ID3DBlob> rsBlob;
ComPtr<ID3DBlob> 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<uint32_t>(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<ComPtr<ID3D12Resource>> 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<float> latentSample;
std::vector<float> input;
std::vector<float> h1;
std::vector<float> h2;
std::vector<float> y;
std::vector<float> dy;
std::vector<float> dx;
std::vector<float> dz3;
std::vector<float> dh2;
std::vector<float> dz2;
std::vector<float> dh1;
std::vector<float> dz1;
std::vector<float> 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<uint32_t>(1000u, d.samples);
d.samplesPerGpuBatch = std::max<uint32_t>(256u, std::min<uint32_t>(d.samplesPerGpuBatch, 65536u));
d.latentResolution = std::max<uint32_t>(16u, std::min<uint32_t>(d.latentResolution, 512u));
d.latentChannels = std::max<uint32_t>(4u, std::min<uint32_t>(d.latentChannels, 16u));
d.hiddenCount = std::max<uint32_t>(16u, std::min<uint32_t>(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 : "<null>") != 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 : "<null>") != 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<uint32_t>(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<uint32_t>(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<uint32_t>(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 ? "<shim RGBA8 albedo image>" : desc.albedoPath;
const char* normalSource = desc.normalImage ? "<shim RGBA8 normal image>" : desc.normalPath;
const char* specSource = desc.specularImage ? "<shim RGBA8 specular image>" : (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);
}
Reference in New Issue View Git Blame Copy Permalink