1use ash::vk;
4use std::ffi::CString;
5
6use super::VulkanDevice;
7use crate::{HalError, TextureFormat};
8
9pub use ash::vk::Pipeline as RawPipeline;
10
11pub struct VulkanShader {
12 pub(crate) module: vk::ShaderModule,
13}
14
15unsafe impl Send for VulkanShader {}
16unsafe impl Sync for VulkanShader {}
17
18pub struct VulkanComputePipeline {
19 pub(crate) pipeline: vk::Pipeline,
20 pub(crate) layout: vk::PipelineLayout,
21 pub(crate) descriptor_set_layout: vk::DescriptorSetLayout,
22 pub(crate) descriptor_pool: vk::DescriptorPool,
23 pub(crate) descriptor_set: vk::DescriptorSet,
24}
25
26unsafe impl Send for VulkanComputePipeline {}
27unsafe impl Sync for VulkanComputePipeline {}
28
29impl VulkanDevice {
30 pub fn create_shader_from_spirv(&self, spirv_words: &[u32]) -> Result<VulkanShader, HalError> {
31 const SPIRV_MAGIC: u32 = 0x07230203;
32 if spirv_words.len() < 5 || spirv_words[0] != SPIRV_MAGIC {
33 return Err(HalError::ShaderError(
34 "Invalid SPIR-V bytecode: missing or invalid magic header".to_string(),
35 ));
36 }
37
38 let create_info = vk::ShaderModuleCreateInfo::default().code(spirv_words);
39 let module = unsafe {
40 self.device
41 .create_shader_module(&create_info, None)
42 .map_err(|e| HalError::ShaderError(e.to_string()))?
43 };
44 Ok(VulkanShader { module })
45 }
46
47 pub fn create_compute_pipeline(
48 &self,
49 shader: &VulkanShader,
50 buffer_count: u32,
51 ) -> Result<VulkanComputePipeline, HalError> {
52 let mut bindings = Vec::with_capacity(buffer_count as usize);
53 for i in 0..buffer_count {
54 bindings.push(
55 vk::DescriptorSetLayoutBinding::default()
56 .binding(i)
57 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
58 .descriptor_count(1)
59 .stage_flags(vk::ShaderStageFlags::COMPUTE),
60 );
61 }
62
63 let layout_info = vk::DescriptorSetLayoutCreateInfo::default().bindings(&bindings);
64 let descriptor_set_layout = unsafe {
65 self.device
66 .create_descriptor_set_layout(&layout_info, None)
67 .map_err(|e| HalError::PipelineError(e.to_string()))?
68 };
69
70 let set_layouts = [descriptor_set_layout];
71 let pipeline_layout_info =
72 vk::PipelineLayoutCreateInfo::default().set_layouts(&set_layouts);
73
74 let layout = unsafe {
75 self.device
76 .create_pipeline_layout(&pipeline_layout_info, None)
77 .map_err(|e| HalError::PipelineError(e.to_string()))?
78 };
79
80 let entry_point = CString::new("main").unwrap();
81 let stage_info = vk::PipelineShaderStageCreateInfo::default()
82 .stage(vk::ShaderStageFlags::COMPUTE)
83 .module(shader.module)
84 .name(&entry_point);
85
86 let pipeline_info = vk::ComputePipelineCreateInfo::default()
87 .stage(stage_info)
88 .layout(layout);
89
90 let pipelines = unsafe {
91 self.device
92 .create_compute_pipelines(self.pipeline_cache, &[pipeline_info], None)
93 .map_err(|(_, e)| HalError::PipelineError(e.to_string()))?
94 };
95
96 let pool_sizes = [vk::DescriptorPoolSize::default()
97 .ty(vk::DescriptorType::STORAGE_BUFFER)
98 .descriptor_count(buffer_count.max(1))];
99
100 let pool_info = vk::DescriptorPoolCreateInfo::default()
101 .pool_sizes(&pool_sizes)
102 .max_sets(1);
103
104 let descriptor_pool = unsafe {
105 self.device
106 .create_descriptor_pool(&pool_info, None)
107 .map_err(|e| HalError::PipelineError(e.to_string()))?
108 };
109
110 let alloc_info = vk::DescriptorSetAllocateInfo::default()
111 .descriptor_pool(descriptor_pool)
112 .set_layouts(&set_layouts);
113
114 let descriptor_sets = unsafe {
115 self.device
116 .allocate_descriptor_sets(&alloc_info)
117 .map_err(|e| HalError::PipelineError(e.to_string()))?
118 };
119
120 Ok(VulkanComputePipeline {
121 pipeline: pipelines[0],
122 layout,
123 descriptor_set_layout,
124 descriptor_pool,
125 descriptor_set: descriptor_sets[0],
126 })
127 }
128
129 pub fn create_bindless_compute_pipeline(
130 &self,
131 shader: &VulkanShader,
132 ) -> Result<vk::Pipeline, HalError> {
133 let heap = self.bindless_heap.as_ref().ok_or_else(|| {
134 HalError::BackendError("Bindless descriptor heap is not enabled".to_string())
135 })?;
136
137 let entry_point = CString::new("main").unwrap();
138 let stage_info = vk::PipelineShaderStageCreateInfo::default()
139 .stage(vk::ShaderStageFlags::COMPUTE)
140 .module(shader.module)
141 .name(&entry_point);
142
143 let pipeline_info = vk::ComputePipelineCreateInfo::default()
144 .stage(stage_info)
145 .layout(heap.pipeline_layout);
146
147 let pipelines = unsafe {
148 self.device
149 .create_compute_pipelines(self.pipeline_cache, &[pipeline_info], None)
150 .map_err(|(_, e)| HalError::PipelineError(e.to_string()))?
151 };
152
153 Ok(pipelines[0])
154 }
155
156 pub fn destroy_raw_pipeline(&self, pipeline: vk::Pipeline) {
157 unsafe {
158 self.device.destroy_pipeline(pipeline, None);
159 }
160 }
161
162 pub fn create_graphics_pipeline(
163 &self,
164 shader: &VulkanShader,
165 vs_entry: &str,
166 fs_entry: &str,
167 color_format: TextureFormat,
168 ) -> Result<vk::Pipeline, HalError> {
169 self.create_graphics_pipeline_with_depth(
170 shader,
171 vs_entry,
172 fs_entry,
173 color_format,
174 None,
175 false,
176 false,
177 )
178 }
179
180 pub fn create_graphics_pipeline_with_depth(
181 &self,
182 shader: &VulkanShader,
183 vs_entry: &str,
184 fs_entry: &str,
185 color_format: TextureFormat,
186 depth_format: Option<TextureFormat>,
187 depth_test: bool,
188 depth_write: bool,
189 ) -> Result<vk::Pipeline, HalError> {
190 self.create_graphics_pipeline_with_depth_and_samples(
191 shader,
192 vs_entry,
193 fs_entry,
194 color_format,
195 depth_format,
196 depth_test,
197 depth_write,
198 1,
199 )
200 }
201
202 pub fn create_graphics_pipeline_with_depth_and_samples(
203 &self,
204 shader: &VulkanShader,
205 vs_entry: &str,
206 fs_entry: &str,
207 color_format: TextureFormat,
208 depth_format: Option<TextureFormat>,
209 depth_test: bool,
210 depth_write: bool,
211 sample_count: u32,
212 ) -> Result<vk::Pipeline, HalError> {
213 let desc = crate::GraphicsPipelineDesc {
214 label: None,
215 vs_entry,
216 fs_entry,
217 color_format,
218 depth_format,
219 depth_test,
220 depth_write,
221 sample_count,
222 primitive: crate::PrimitiveState::default(),
223 blend: None,
224 };
225 self.create_graphics_pipeline_with_desc(shader, &desc)
226 }
227
228 pub fn create_graphics_pipeline_with_desc(
229 &self,
230 shader: &VulkanShader,
231 desc: &crate::GraphicsPipelineDesc,
232 ) -> Result<vk::Pipeline, HalError> {
233 let heap = self.bindless_heap.as_ref().ok_or_else(|| {
234 HalError::BackendError("Bindless descriptor heap is not enabled".to_string())
235 })?;
236
237 let vk_format = desc.color_format.to_vk();
238
239 let vs_name = CString::new(desc.vs_entry).unwrap();
240 let fs_name = CString::new(desc.fs_entry).unwrap();
241
242 let stages = [
243 vk::PipelineShaderStageCreateInfo::default()
244 .stage(vk::ShaderStageFlags::VERTEX)
245 .module(shader.module)
246 .name(&vs_name),
247 vk::PipelineShaderStageCreateInfo::default()
248 .stage(vk::ShaderStageFlags::FRAGMENT)
249 .module(shader.module)
250 .name(&fs_name),
251 ];
252
253 let vertex_input_info = vk::PipelineVertexInputStateCreateInfo::default();
254
255 let input_assembly = vk::PipelineInputAssemblyStateCreateInfo::default()
256 .topology(desc.primitive.topology.to_vk())
257 .primitive_restart_enable(false);
258
259 let dynamic_states = [vk::DynamicState::VIEWPORT, vk::DynamicState::SCISSOR];
260 let dynamic_state_info =
261 vk::PipelineDynamicStateCreateInfo::default().dynamic_states(&dynamic_states);
262
263 let viewport_state = vk::PipelineViewportStateCreateInfo::default()
264 .viewport_count(1)
265 .scissor_count(1);
266
267 let rasterizer = vk::PipelineRasterizationStateCreateInfo::default()
268 .depth_clamp_enable(false)
269 .rasterizer_discard_enable(false)
270 .polygon_mode(vk::PolygonMode::FILL)
271 .line_width(1.0)
272 .cull_mode(desc.primitive.cull_mode.to_vk())
273 .front_face(desc.primitive.front_face.to_vk())
274 .depth_bias_enable(false);
275
276 let vk_samples = match desc.sample_count {
277 2 => vk::SampleCountFlags::TYPE_2,
278 4 => vk::SampleCountFlags::TYPE_4,
279 8 => vk::SampleCountFlags::TYPE_8,
280 16 => vk::SampleCountFlags::TYPE_16,
281 32 => vk::SampleCountFlags::TYPE_32,
282 64 => vk::SampleCountFlags::TYPE_64,
283 _ => vk::SampleCountFlags::TYPE_1,
284 };
285
286 let multisampling = vk::PipelineMultisampleStateCreateInfo::default()
287 .sample_shading_enable(false)
288 .rasterization_samples(vk_samples);
289
290 let color_blend_attachment = if let Some(b) = desc.blend {
291 [vk::PipelineColorBlendAttachmentState::default()
292 .color_write_mask(vk::ColorComponentFlags::RGBA)
293 .blend_enable(true)
294 .src_color_blend_factor(b.color.src_factor.to_vk())
295 .dst_color_blend_factor(b.color.dst_factor.to_vk())
296 .color_blend_op(b.color.operation.to_vk())
297 .src_alpha_blend_factor(b.alpha.src_factor.to_vk())
298 .dst_alpha_blend_factor(b.alpha.dst_factor.to_vk())
299 .alpha_blend_op(b.alpha.operation.to_vk())]
300 } else {
301 [vk::PipelineColorBlendAttachmentState::default()
302 .color_write_mask(vk::ColorComponentFlags::RGBA)
303 .blend_enable(false)]
304 };
305
306 let color_blending = vk::PipelineColorBlendStateCreateInfo::default()
307 .logic_op_enable(false)
308 .attachments(&color_blend_attachment);
309
310 let color_formats = [vk_format];
311 let mut rendering_info =
312 vk::PipelineRenderingCreateInfo::default().color_attachment_formats(&color_formats);
313
314 let vk_depth_format = match desc.depth_format {
315 Some(TextureFormat::Depth32Float) => Some(vk::Format::D32_SFLOAT),
316 Some(f) => {
317 return Err(HalError::BackendError(format!(
318 "Unsupported depth format: {f:?}"
319 )))
320 }
321 None => None,
322 };
323
324 if let Some(df) = vk_depth_format {
325 rendering_info = rendering_info.depth_attachment_format(df);
326 }
327
328 let depth_stencil = vk::PipelineDepthStencilStateCreateInfo::default()
329 .depth_test_enable(desc.depth_test)
330 .depth_write_enable(desc.depth_write)
331 .depth_compare_op(if desc.depth_test {
332 vk::CompareOp::LESS_OR_EQUAL
333 } else {
334 vk::CompareOp::ALWAYS
335 })
336 .depth_bounds_test_enable(false)
337 .stencil_test_enable(false);
338
339 let mut pipeline_info = vk::GraphicsPipelineCreateInfo::default()
340 .stages(&stages)
341 .vertex_input_state(&vertex_input_info)
342 .input_assembly_state(&input_assembly)
343 .viewport_state(&viewport_state)
344 .rasterization_state(&rasterizer)
345 .multisample_state(&multisampling)
346 .color_blend_state(&color_blending)
347 .dynamic_state(&dynamic_state_info)
348 .layout(heap.pipeline_layout)
349 .push_next(&mut rendering_info);
350
351 if desc.depth_format.is_some() || desc.depth_test {
352 pipeline_info = pipeline_info.depth_stencil_state(&depth_stencil);
353 }
354
355 let pipelines = unsafe {
356 self.device
357 .create_graphics_pipelines(self.pipeline_cache, &[pipeline_info], None)
358 .map_err(|(_, e)| {
359 HalError::PipelineError(format!("create_graphics_pipelines failed: {e}"))
360 })?
361 };
362
363 Ok(pipelines[0])
364 }
365
366 pub fn destroy_shader(&self, shader: VulkanShader) {
367 unsafe {
368 self.device.destroy_shader_module(shader.module, None);
369 }
370 }
371
372 pub fn destroy_compute_pipeline(&self, pipeline: VulkanComputePipeline) {
373 unsafe {
374 self.device
375 .destroy_descriptor_pool(pipeline.descriptor_pool, None);
376 self.device
377 .destroy_descriptor_set_layout(pipeline.descriptor_set_layout, None);
378 self.device.destroy_pipeline(pipeline.pipeline, None);
379 self.device.destroy_pipeline_layout(pipeline.layout, None);
380 }
381 }
382}