yum-mirror/slang

Making it easier to work with shaders

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

Yong HeAllow partial specialization of existential arguments. (#6487)dd9d24d29

master
3.0 KiB93 linesraw
1//TEST:SIMPLE(filecheck=CHECK): -target spirv
2
3// This is a test that checks that we can apply partial specialization to a function
4// that takes existential parameters.
5
6// CHECK: OpRayQueryProceedKHR
7// CHECK: OpImageWrite
8
9public interface IRandom {
10    [mutating] uint32_t next_uint();
11    [mutating] float next_float();
12};
13
14public struct TEA : IRandom {
15    int val;
16    public __init(uint32_t v0, uint32_t v1) {}
17    [mutating] public uint32_t next_uint() {
18        return val++;
19    }
20    [mutating] public float next_float() {
21        return val++;
22    }
23};
24
25public interface IScene {
26    property RaytracingAccelerationStructure as;
27};
28
29// The Scene type contains a resource field, if dynamic dispatch code were generated
30// for this type, we will get a compile error.
31struct Scene : IScene {
32    RaytracingAccelerationStructure as;
33};
34
35public interface IIntegrator {
36    
37    // This function takes two existential parameters, `scene` and `rng`.
38    // if we call this function with `rng` being dynamic, and `scene` being static,
39    // we should still be able to specialize the `sample` function with the statically known
40    // type of `scene`.
41    public float3 sample(IScene scene, RayDesc ray, IRandom rng);
42};
43namespace integrator {
44    public struct NoShading : IIntegrator {
45        public float3 sample(IScene scene, RayDesc _ray, IRandom rng) {
46            return float3( 0.0f, 0.0f, 0.0f );
47        }
48    };
49
50    struct Test {
51
52        float4 sample(IScene scene, RayDesc _ray, IRandom rng, int pixel_aabb_uv, uint3 id) {
53            float4 grad = { 0.0f, 0.0f, 0.0f, 1.0f };
54
55            RayDesc ray = _ray; uint32_t depth = 0;
56            RayQuery<0> rayQuery;
57            while (rayQuery.Proceed()) {
58                float rand = rng.next_float();
59
60                IIntegrator integrator = integrator::NoShading();
61
62                // Here `rng` is mutating in the loop, so its type is dynamic and we need
63                // to generate dynamic dispatch code around it.
64                // But this shouldn't result in `scene` being dynamic as well.
65                // We should still be able to specialize `integrator.sample` with the statically
66                // known type of `scene`.
67                // If this doesn't happen, then the compiler will try to synthesize dynamic dispatch
68                // logic for `scene` and fail to compile.
69                let color_in = integrator.sample(scene, {}, rng);
70                let color_out = integrator.sample(scene, {}, rng);
71                grad.xyz += color_out;
72            }
73            return grad;
74        }
75    };
76};
77
78[[vk::binding(0, 0)]] RWTexture2D<float4> output;
79
80[vk::constant_id(0)] const int WGS_X = 1;
81[vk::constant_id(1)] const int WGS_Y = 1;
82[shader("compute"), numthreads(WGS_X, WGS_Y, 1)]
83void main(
84    uint3 id : SV_DispatchThreadID
85)
86{
87    IRandom rng = TEA(id.y * id.x, 1);
88    Scene scene = { RaytracingAccelerationStructure(0) };
89
90    integrator::Test integrator = integrator::Test();
91    float4 grad = integrator.sample(scene, {}, rng, {}, id);
92    output[id.xy] += float4(grad.xyz, grad.w);
93}