MCPcopy Create free account
hub / github.com/evilsocket/cake / dispatch_compute

Method dispatch_compute

cake-core/src/backends/vulkan/mod.rs:692–799  ·  view source on GitHub ↗

Core dispatch: bind pipeline, bind buffers, dispatch workgroups, fence wait, read output.

(
        &self,
        entry: &str,
        storage_buffers: &[vk::Buffer], // bindings 0, 1, 2, [4]
        output_buf: &MappedBuffer,
        output_count: usize,
        params: &[u32],
        w

Source from the content-addressed store, hash-verified

690
691 /// Core dispatch: bind pipeline, bind buffers, dispatch workgroups, fence wait, read output.
692 fn dispatch_compute(
693 &self,
694 entry: &str,
695 storage_buffers: &[vk::Buffer], // bindings 0, 1, 2, [4]
696 output_buf: &MappedBuffer,
697 output_count: usize,
698 params: &[u32],
699 workgroups: (u32, u32, u32),
700 ) -> Vec<f32> {
701 let _lock = self.dispatch_lock.lock().unwrap();
702 let (pipeline, pipe_layout, _ds_layout, descriptor_set, _num_bindings) =
703 self.pipelines.get(entry).unwrap_or_else(|| panic!("unknown pipeline: {entry}"));
704 let descriptor_set = *descriptor_set;
705
706 // Write params
707 self.params_buf.write_u32(params);
708
709 unsafe {
710 // Build descriptor writes with stack-allocated arrays (max 5 storage + 1 uniform = 6)
711 let num_storage = storage_buffers.len();
712 let mut buf_infos: [vk::DescriptorBufferInfo; 6] = Default::default();
713 for (i, &buf) in storage_buffers.iter().enumerate() {
714 buf_infos[i] = vk::DescriptorBufferInfo::default()
715 .buffer(buf)
716 .offset(0)
717 .range(vk::WHOLE_SIZE);
718 }
719 buf_infos[num_storage] = vk::DescriptorBufferInfo::default()
720 .buffer(self.params_buf.buffer)
721 .offset(0)
722 .range(16);
723
724 let mut writes: [vk::WriteDescriptorSet; 6] = Default::default();
725 for i in 0..num_storage {
726 let binding = if i == 3 { 4 } else { i as u32 };
727 writes[i] = vk::WriteDescriptorSet::default()
728 .dst_set(descriptor_set)
729 .dst_binding(binding)
730 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
731 .buffer_info(std::slice::from_ref(&buf_infos[i]));
732 }
733 writes[num_storage] = vk::WriteDescriptorSet::default()
734 .dst_set(descriptor_set)
735 .dst_binding(3)
736 .descriptor_type(vk::DescriptorType::UNIFORM_BUFFER)
737 .buffer_info(std::slice::from_ref(&buf_infos[num_storage]));
738
739 self.vk_device.update_descriptor_sets(&writes[..num_storage + 1], &[]);
740
741 // Reset fence + command buffer.
742 // Fence is always signaled here: either from SIGNALED init flag (first call)
743 // or from the wait_for_fences at the bottom of the previous dispatch.
744 self.vk_device.reset_fences(&[self.fence]).expect("reset_fences");
745 self.vk_device
746 .reset_command_buffer(self.command_buffer, vk::CommandBufferResetFlags::empty())
747 .expect("reset_command_buffer");
748
749 // Record

Callers 7

init_vulkanMethod · 0.80
dispatch_binary_vec4Method · 0.80
dispatch_ternary_vec4Method · 0.80
dispatch_unary_vec4Method · 0.80
dispatch_softmaxMethod · 0.80
gpu_gemvMethod · 0.80
gpu_gemmMethod · 0.80

Calls 3

write_u32Method · 0.80
read_f32Method · 0.80
getMethod · 0.45

Tested by

no test coverage detected