yum-mirror/slang

Making it easier to work with shaders

git clone https://git.yummers.dev/yum-mirror/slang

Harsh Aggarwal (NVIDIA)Fix 7723 - Add autodiff tests (#7919)d10732742

master
branchesmasterall refs →
2.1 KiB84 linesraw
1
2//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
3//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
4//TEST(compute):COMPARE_COMPUTE_EX:-cuda -compute -shaderobj -output-using-type
5
6//TEST_INPUT:ubuffer(data=[0], stride=4):out,name=outputBuffer
7RWStructuredBuffer<float> outputBuffer;
8
9static const uint32_t N_LATENT_DIMS = 4;
10static const uint32_t kDecoderInputCount = 6;
11struct LatentTexture
12{
13    static const uint32_t kLatentDimsCount = N_LATENT_DIMS;
14    static const uint32_t kLatentTextureCount = N_LATENT_DIMS / 4;
15
16    [BackwardDifferentiable]
17    void getCodeStochastic(float2 uv, out float code[kLatentDimsCount])
18    {
19        return getCode(uint2(1,2), code);
20    }
21
22    void getCode(uint2 texel, out float code[kLatentDimsCount])
23    {
24        for (uint i = 0; i < kLatentTextureCount; ++i)
25        {
26            for (uint j = 0; j < 4; ++j)
27            {
28                code[i * 4 + j] = j;
29            }
30        }
31    }
32    [BackwardDerivativeOf(getCode)]
33    void bwd_getCode(uint2 texel, float d_out[kLatentDimsCount])
34    {
35        outputBuffer[0] = d_out[0];
36    }
37}
38
39static LatentTexture gLatents;
40
41[BackwardDifferentiable]
42void test(float arr[10], out float result[3])
43{
44    float sum = 0;
45    [ForceUnroll]
46    for (int i = 0; i < LatentTexture.kLatentDimsCount + kDecoderInputCount; i++)
47        sum += arr[i];
48    result[0] = sum;
49    result[1] = sum;
50    result[2] = sum;
51}
52
53[BackwardDifferentiable]
54float evalDecoder()
55{
56    // Latent code.
57    float latentCode[LatentTexture.kLatentDimsCount];
58    gLatents.getCodeStochastic(float2(1,2), latentCode);
59
60    // Model input.
61    float input[kDecoderInputCount + LatentTexture.kLatentDimsCount];
62    input[0] = 0;
63    input[1] = 1;
64    input[2] = 2;
65    input[3] = 3;
66    input[4] = 4;
67    input[5] = 5;
68    [ForceUnroll]
69    for (int i = 0; i < LatentTexture.kLatentDimsCount; i++)
70    {
71        input[kDecoderInputCount + i] = latentCode[i];
72    }
73
74    float res[3];
75    test(input, res);
76    return res[0] + res[1] + res[2];
77}
78
79
80[numthreads(1, 1, 1)]
81void computeMain(uint3 dispatchThreadID: SV_DispatchThreadID)
82{
83    __bwd_diff(evalDecoder)(1.0);
84}