1 module compute 2 3 def begin = 4 let vulkan : Vulkan @auto 5 assert vulkan.gpu_state.is_none 6 7 vulkan.gpu_state = GpuState/Compute 8 command_buffer = vulkan.command_buffer.object 9 shader = null 10 pipeline = null 11 local_group_size = [0, 0, 0] 12 descriptor_set_states = memory@zero 13 images_back_to_read = Set<Image>.new 14 pipeline_push_constant_stages = ShaderStageFlags@zero 15 16 def bind_pipeline (pipeline : ComputePipeline) = 17 let vulkan : Vulkan @auto 18 let state = vulkan.gpu_state.compute_mut 19 if pipeline == state.pipeline then 20 return 21 22 state.pipeline = pipeline 23 24 vk/cmd_bind_pipeline state.command_buffer PipelineBindPoint/Compute 25 pipeline.object 26 27 if state.shader <> pipeline.shader then 28 state.shader = pipeline.shader 29 local_group_size = pipeline.local_group_size 30 pipeline_push_constant_stages = pipeline.push_constant_stages 31 32 def bind_descriptor_set (descriptor_set : DescriptorSet) (index : u32) = 33 let vulkan : Vulkan @auto 34 let state = vulkan.gpu_state.compute_mut 35 let item = state.descriptor_set_states[ref index] 36 item.descriptor_set = descriptor_set 37 bound = false 38 39 let barriers = List<ImageMemoryBarrier>.new 40 let mut src_stage_flags = PipelineStageFlags@zero 41 42 for image in descriptor_set.sampled_images do 43 if image.layout <> ImageLayout/ShaderReadOnlyOptimal then 44 src_stage_flags |= PipelineStageFlags/ComputeShader 45 46 let image_memory_barrier = ImageMemoryBarrier 47 src_access_mask = AccessFlags/ShaderRead | AccessFlags/ShaderWrite 48 dst_access_mask = AccessFlags/ShaderRead | AccessFlags/ShaderWrite 49 f@ image.barrier_to ImageLayout/ShaderReadOnlyOptimal 50 51 barriers.add image_memory_barrier 52 image.layout = ImageLayout/ShaderReadOnlyOptimal 53 54 state.images_back_to_read.remove image 55 vulkan.compute_images.add image 56 57 for image in descriptor_set.storage_images do 58 if image.layout <> ImageLayout/General then 59 let mut src_access_flags = AccessFlags@zero 60 61 if vulkan.compute_images.contains image then 62 src_stage_flags |= PipelineStageFlags/ComputeShader 63 src_access_flags |= AccessFlags/ShaderRead | AccessFlags/ShaderWrite 64 vulkan.compute_images.remove image 65 66 if vulkan.raster_images.contains image then 67 src_stage_flags |= PipelineStageFlags/FragmentShader 68 | PipelineStageFlags/VertexShader 69 70 src_access_flags |= AccessFlags/ShaderRead | AccessFlags/ShaderWrite 71 vulkan.raster_images.remove image 72 73 if vulkan.transfer_images.contains image then 74 src_stage_flags |= PipelineStageFlags/Transfer 75 src_access_flags |= AccessFlags/TransferWrite 76 | AccessFlags/TransferRead 77 78 vulkan.transfer_images.remove image 79 80 let image_memory_barrier = ImageMemoryBarrier 81 src_access_mask = src_access_flags 82 dst_access_mask = AccessFlags/ShaderRead | AccessFlags/ShaderWrite 83 f@ image.barrier_to ImageLayout/General 84 85 barriers.add image_memory_barrier 86 image.layout = ImageLayout/General 87 88 state.images_back_to_read.insert image 89 90 if barriers.size > 0 then 91 if src_stage_flags == PipelineStageFlags@zero then 92 src_stage_flags = PipelineStageFlags/TopOfPipe 93 94 vk/cmd_pipeline_barrier state.command_buffer src_stage_flags 95 PipelineStageFlags/ComputeShader 96 DependencyFlags@zero 0 null 0 null 97 barriers.size barriers.as_ptr 98 99 def set_push_constant<T> (data : Ptr<T>) = 100 let vulkan : Vulkan @auto 101 let state = vulkan.gpu_state.compute 102 let size = sizeof T 103 104 vk/cmd_push_constants state.command_buffer state.pipeline.pipeline_layout 105 state.pipeline_push_constant_stages 0 size data 106 107 def end (post_barrier : BarrierFlags) = 108 let vulkan : Vulkan @auto 109 110 let mut stage_flags = PipelineStageFlags@zero 111 let mut access_flags = AccessFlags@zero 112 113 if post_barrier.has_compute then 114 stage_flags |= PipelineStageFlags/ComputeShader 115 access_flags |= AccessFlags/ShaderRead | AccessFlags/ShaderWrite 116 117 if post_barrier.has_raster then 118 stage_flags |= PipelineStageFlags/VertexInput 119 | PipelineStageFlags/VertexShader 120 | PipelineStageFlags/FragmentShader 121 | PipelineStageFlags/DrawIndirect 122 access_flags |= AccessFlags/ShaderRead 123 | AccessFlags/ShaderWrite 124 | AccessFlags/IndexRead 125 | AccessFlags/VertexAttributeRead 126 | AccessFlags/IndirectCommandRead 127 128 if post_barrier.has_transfer then 129 stage_flags |= PipelineStageFlags/Transfer 130 access_flags |= AccessFlags/TransferWrite | AccessFlags/TransferRead 131 132 if stage_flags == PipelineStageFlags@zero then 133 stage_flags = PipelineStageFlags/BottomOfPipe 134 135 let state = vulkan.gpu_state.compute 136 let image_barriers = 137 state.images_back_to_read 138 |> map { image -> ImageMemoryBarrier 139 src_access_mask = AccessFlags/ShaderRead | AccessFlags/ShaderWrite 140 dst_access_mask = access_flags 141 f@ image.barrier_to ImageLayout/ShaderReadOnlyOptimal } 142 143 for image in state.images_back_to_read do 144 image.layout = ImageLayout/ShaderReadOnlyOptimal 145 146 let memory_barrier = MemoryBarrier 147 type = StructureType/MemoryBarrier 148 next = null 149 src_access_mask = AccessFlags/ShaderWrite 150 dst_access_mask = access_flags 151 152 vk/cmd_pipeline_barrier state.command_buffer PipelineStageFlags/ComputeShader 153 stage_flags DependencyFlags@zero 1 154 memory_barrier@ptr 0 null image_barriers.size 155 image_barriers.as_ptr 156 157 vulkan.gpu_state = GpuState/None 158 159 def dispatch (x : u32 160 y : u32 161 z : u32) = 162 let vulkan : Vulkan @auto 163 let state = vulkan.gpu_state.compute_mut 164 let group_size = state.local_group_size 165 let group_count_x = (x - 1) / group_size[0] + 1 166 let group_count_y = (y - 1) / group_size[1] + 1 167 let group_count_z = (z - 1) / group_size[2] + 1 168 169 bind_descriptor_sets state.command_buffer PipelineBindPoint/Compute 170 state.pipeline.pipeline_layout 171 state.descriptor_set_states@mut_ref 172 173 vk/cmd_dispatch state.command_buffer group_count_x group_count_y 174 group_count_z 175 176 def resolve_depth (src_descriptor_set : DescriptorSet 177 dst_descriptor_set : DescriptorSet 178 size : Size 179 samples : i32 180 post_barrier : BarrierFlags) = 181 let vulkan : Vulkan @auto 182 183 let push_constant = ResolvePushConstant 184 size = [size.width as i32, size.height as i32] 185 =samples 186 pad = 0 187 188 begin 189 190 bind_pipeline vulkan.resolve.pipeline 191 bind_descriptor_set src_descriptor_set 0 192 bind_descriptor_set dst_descriptor_set 1 193 set_push_constant push_constant@ptr 194 dispatch size.width size.height 1 195 196 end post_barrier 197