yum-mirror/slang

Making it easier to work with shaders

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

venkataram-nvReport AD checkpoint contexts (#5058)b808aa4df

master
1.4 KiB49 linesraw
1//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
2//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
3//TEST(compute):COMPARE_COMPUTE_EX:-cpu -compute -output-using-type -shaderobj
4//TEST:SIMPLE(filecheck=CHECK): -target hlsl -profile cs_5_0 -entry computeMain -line-directive-mode none
5//TEST:SIMPLE(filecheck=CHK):-target glsl -stage compute -entry computeMain -report-checkpoint-intermediates
6
7//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
8RWStructuredBuffer<float> outputBuffer;
9
10typedef DifferentialPair<float> dpfloat;
11typedef float.Differential dfloat;
12
13[BackwardDifferentiable]
14[PreferCheckpoint]
15float g(float x)
16{
17    return log(x);
18}
19
20//CHK: note: checkpointing context of 4 bytes associated with function: 'f'
21[BackwardDifferentiable]
22float f(int p, float x)
23{
24    float y = 1.0;
25    // Test that phi parameter can be restored.
26    if (p == 0)
27        //CHK: note: 4 bytes (float) used to checkpoint the following item:
28        y = g(x);
29
30    return y * y;
31}
32
33// Check that there are no calls to primal_g in bwd_f.
34
35// CHECK: void s_bwd_f_{{[0-9]+}}
36// CHECK-NOT: {{[_a-zA-Z0-9]+}} = s_primal_ctx_g_{{[0-9]+}}
37// CHECK: return
38
39
40[numthreads(1, 1, 1)]
41void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
42{
43    dpfloat dpa = dpfloat(2.0, 0.0);
44
45    __bwd_diff(f)(0, dpa, 1.0f);
46    outputBuffer[0] = dpa.d; // Expect: 1
47}
48
49//CHK-NOT: note