yum-mirror/slang

Making it easier to work with shaders

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

Harsh Aggarwal (NVIDIA)Enable CUDA Test Enablement - Batch 1: Autodiff Tests (1-16) (#8139)c3df36043

master
branchesmasterall refs →
1.4 KiB49 linesraw
1//TEST:COMPARE_COMPUTE(filecheck-buffer=CHECK):-output-using-type
2//TEST:COMPARE_COMPUTE(filecheck-buffer=CHECK):-output-using-type -cuda
3public struct ReadOnlyIndex
4{
5    private int _idx;
6    __init(int i) { _idx = i; }
7    public property int idx { get { return _idx; } }
8}
9struct GradientBuffer
10{
11    RWStructuredBuffer<float> primal;
12    StructuredBuffer<float> grad;
13
14    [Differentiable]
15    void write(int idx, float v) { primal[idx] = detach(v); }
16
17    [BackwardDerivativeOf(write)]
18    void write_bwd(int idx, inout DifferentialPair<float> d) { d = diffPair(d.p, grad[idx]); }
19
20    [Differentiable]
21    void store(ReadOnlyIndex idx, float v) { write(idx.idx, v); }
22}
23[Differentiable]
24void test(GradientBuffer buf, ReadOnlyIndex b, float x)
25{
26    buf.store(b, x);
27}
28public float repro(RWStructuredBuffer<float> primal, StructuredBuffer<float> grad)
29{
30    DifferentialPair<float> result = diffPair(1.0f);
31    GradientBuffer buf = { primal, grad };
32    bwd_diff(test)(buf, ReadOnlyIndex(5), result);
33    return result.d;
34}
35
36//TEST_INPUT: set output = out ubuffer(data=[0 0 0 0], stride=4)
37RWStructuredBuffer<float> output;
38
39//TEST_INPUT: set gPrimal = ubuffer(data=[0.0 1.0 2.0 3.0 4.0 5.0 6.0 7.0], stride=4)
40RWStructuredBuffer<float> gPrimal;
41//TEST_INPUT: set gGrad = ubuffer(data=[0.0 1.0 2.0 3.0 4.0 5.0 6.0 7.0], stride=4)
42StructuredBuffer<float> gGrad;
43
44[numthreads(1,1,1)]
45void computeMain()
46{
47    // CHECK: 5.0
48    output[0] = repro(gPrimal, gGrad);
49}