1
+ #![ no_std]
2
+
3
+ use spirv_std:: glam:: { ivec2, vec2, vec4, IVec2 , UVec2 , Vec2 , Vec3 , Vec4 , Vec4Swizzles } ;
4
+ use spirv_std:: image:: { ImageWithMethods , sample_with} ;
5
+ use spirv_std:: { num_traits:: Float , spirv, Image } ;
6
+
7
+ #[ repr( C , align( 16 ) ) ]
8
+ #[ derive( Copy , Clone ) ]
9
+ pub struct Light {
10
+ pub position : Vec4 , // 16 bytes, aligned to 16
11
+ pub color_radius : Vec4 , // 16 bytes - store color.xyz in xyz, radius in w
12
+ }
13
+
14
+ #[ repr( C , align( 16 ) ) ]
15
+ #[ derive( Copy , Clone ) ]
16
+ pub struct UBO {
17
+ pub lights : [ Light ; 6 ] , // 6 * 32 = 192 bytes
18
+ pub view_pos : Vec4 , // 16 bytes, total = 208 bytes
19
+ pub _padding1 : Vec4 , // 16 bytes padding to align next member to 224
20
+ pub debug_display_target : Vec4 , // Use Vec4 with debug_display_target in x component
21
+ }
22
+
23
+ #[ spirv( vertex) ]
24
+ pub fn main_vs (
25
+ #[ spirv( vertex_index) ] vert_index : i32 ,
26
+ #[ spirv( position) ] out_position : & mut Vec4 ,
27
+ out_uv : & mut Vec2 ,
28
+ ) {
29
+ let uv = vec2 ( ( ( vert_index << 1 ) & 2 ) as f32 , ( vert_index & 2 ) as f32 ) ;
30
+ * out_uv = uv;
31
+ * out_position = vec4 ( uv. x * 2.0 - 1.0 , uv. y * 2.0 - 1.0 , 0.0 , 1.0 ) ;
32
+ }
33
+
34
+ const NUM_LIGHTS : usize = 6 ;
35
+
36
+ fn resolve ( tex : & Image ! ( 2 D , format=rgba8, sampled, multisampled) , uv : IVec2 , num_samples : u32 ) -> Vec4 {
37
+ let mut result = Vec4 :: ZERO ;
38
+ for i in 0 ..num_samples {
39
+ let val: Vec4 = tex. fetch_with ( uv, sample_with:: sample_index ( i as i32 ) ) ;
40
+ result += val;
41
+ }
42
+ // Average resolved samples
43
+ result / ( num_samples as f32 )
44
+ }
45
+
46
+ fn calculate_lighting ( pos : Vec3 , normal : Vec3 , albedo : Vec4 , ubo : & UBO ) -> Vec3 {
47
+ let mut result = Vec3 :: ZERO ;
48
+
49
+ for i in 0 ..NUM_LIGHTS {
50
+ // Vector to light
51
+ let l = ubo. lights [ i] . position . xyz ( ) - pos;
52
+ // Distance from light to fragment position
53
+ let dist = l. length ( ) ;
54
+
55
+ // Viewer to fragment
56
+ let v = ( ubo. view_pos . xyz ( ) - pos) . normalize ( ) ;
57
+
58
+ // Light to fragment
59
+ let l = l. normalize ( ) ;
60
+
61
+ // Attenuation
62
+ let atten = ubo. lights [ i] . color_radius . w / ( dist. powf ( 2.0 ) + 1.0 ) ;
63
+
64
+ // Diffuse part
65
+ let n = normal. normalize ( ) ;
66
+ let n_dot_l = n. dot ( l) . max ( 0.0 ) ;
67
+ let diff = ubo. lights [ i] . color_radius . xyz ( ) * albedo. xyz ( ) * n_dot_l * atten;
68
+
69
+ // Specular part
70
+ let r = ( -l) . reflect ( n) ;
71
+ let n_dot_r = r. dot ( v) . max ( 0.0 ) ;
72
+ let spec = ubo. lights [ i] . color_radius . xyz ( ) * albedo. w * n_dot_r. powf ( 8.0 ) * atten;
73
+
74
+ result += diff + spec;
75
+ }
76
+ result
77
+ }
78
+
79
+ #[ spirv( fragment) ]
80
+ pub fn main_fs (
81
+ in_uv : Vec2 ,
82
+ #[ spirv( descriptor_set = 0 , binding = 1 ) ] sampler_position : & Image ! ( 2 D , format=rgba16f, sampled, multisampled) ,
83
+ #[ spirv( descriptor_set = 0 , binding = 2 ) ] sampler_normal : & Image ! ( 2 D , format=rgba16f, sampled, multisampled) ,
84
+ #[ spirv( descriptor_set = 0 , binding = 3 ) ] sampler_albedo : & Image ! ( 2 D , format=rgba8, sampled, multisampled) ,
85
+ #[ spirv( uniform, descriptor_set = 0 , binding = 4 ) ] ubo : & UBO ,
86
+ #[ spirv( spec_constant( id = 0 , default = 8 ) ) ] num_samples : u32 ,
87
+ out_frag_color : & mut Vec4 ,
88
+ ) {
89
+ let att_dim: UVec2 = sampler_position. query_size ( ) ;
90
+ let uv = ivec2 ( ( in_uv. x * att_dim. x as f32 ) as i32 , ( in_uv. y * att_dim. y as f32 ) as i32 ) ;
91
+
92
+ // Debug display
93
+ if ubo. debug_display_target . x as i32 > 0 {
94
+ let val: Vec4 = match ubo. debug_display_target . x as i32 {
95
+ 1 => sampler_position. fetch_with ( uv, sample_with:: sample_index ( 0 ) ) ,
96
+ 2 => sampler_normal. fetch_with ( uv, sample_with:: sample_index ( 0 ) ) ,
97
+ 3 => sampler_albedo. fetch_with ( uv, sample_with:: sample_index ( 0 ) ) ,
98
+ 4 => {
99
+ let alb: Vec4 = sampler_albedo. fetch_with ( uv, sample_with:: sample_index ( 0 ) ) ;
100
+ vec4 ( alb. w , alb. w , alb. w , 1.0 )
101
+ } ,
102
+ _ => Vec4 :: ZERO ,
103
+ } ;
104
+ * out_frag_color = vec4 ( val. x , val. y , val. z , 1.0 ) ;
105
+ return ;
106
+ }
107
+
108
+ const AMBIENT : f32 = 0.15 ;
109
+
110
+ // Ambient part
111
+ let alb = resolve ( sampler_albedo, uv, num_samples) ;
112
+ let mut frag_color = Vec3 :: ZERO ;
113
+
114
+ // Calculate lighting for every MSAA sample
115
+ for i in 0 ..num_samples {
116
+ let pos: Vec4 = sampler_position. fetch_with ( uv, sample_with:: sample_index ( i as i32 ) ) ;
117
+ let normal: Vec4 = sampler_normal. fetch_with ( uv, sample_with:: sample_index ( i as i32 ) ) ;
118
+ let albedo: Vec4 = sampler_albedo. fetch_with ( uv, sample_with:: sample_index ( i as i32 ) ) ;
119
+ frag_color += calculate_lighting ( pos. xyz ( ) , normal. xyz ( ) , albedo, ubo) ;
120
+ }
121
+
122
+ frag_color = ( alb. xyz ( ) * AMBIENT ) + frag_color / ( num_samples as f32 ) ;
123
+
124
+ * out_frag_color = vec4 ( frag_color. x , frag_color. y , frag_color. z , 1.0 ) ;
125
+ }
0 commit comments