1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
//! GPU-visible buffers in plain system memory (Vulkan only).
//!
//! wgpu places every `MAP_WRITE` buffer — including the queue's own
//! `write_buffer` staging — through its allocator's "CPU to GPU" location,
//! which prefers DEVICE_LOCAL | HOST_VISIBLE memory: without resizable BAR
//! that is the 256 MB PCIe window, so the CPU writes every byte across the
//! bus itself (write-combined stores into the card). Measured on an RTX 3090
//! (PCIe 4.0 x16, 256 MB BAR) that path tops out near the rate the prompt
//! frames' expert admissions ran at.
//!
//! A buffer allocated here lives in a HOST_VISIBLE, non-DEVICE_LOCAL memory
//! type (the host heap), is mapped once for its lifetime and can be the
//! source of `copy_buffer_to_buffer`: the card's copy engine then pulls the
//! bytes over PCIe (DMA), and the CPU only ever writes system RAM.
/// A system-memory buffer mapped for its whole life. `ptr` stays valid
/// while `buffer` lives (wgpu frees the memory when the buffer drops).
pub struct SysBuf {
pub buffer: wgpu::Buffer,
ptr: *mut u8,
size: u64,
/// HOST_CACHED memory (CPU reads are cheap, the card snoops).
pub cached: bool,
}
// SAFETY: the pointer is to a persistent mapping of device memory owned by
// `buffer`; callers write disjoint ranges.
unsafe impl Send for SysBuf {}
unsafe impl Sync for SysBuf {}
impl SysBuf {
pub fn size(&self) -> u64 {
self.size
}
/// Raw pointer to the mapping (valid for `size` bytes).
pub fn as_ptr(&self) -> *mut u8 {
self.ptr
}
/// Copy `src` to `off`. The caller keeps writers to one range apart and
/// keeps a range from being rewritten while a copy out of it is pending.
pub fn write(&self, off: u64, src: &[u8]) -> bool {
if off.checked_add(src.len() as u64).is_none_or(|e| e > self.size) {
return false;
}
// SAFETY: bounds checked above; the mapping is live while self is.
unsafe {
std::ptr::copy_nonoverlapping(src.as_ptr(), self.ptr.add(off as usize), src.len());
}
true
}
}
/// Allocate `size` bytes of host-heap memory as a `COPY_SRC | COPY_DST`
/// buffer: HOST_VISIBLE | HOST_COHERENT, with HOST_CACHED when `cached`,
/// never DEVICE_LOCAL. None off Vulkan (and off Linux) or when the driver
/// refuses.
#[cfg(target_os = "linux")]
pub fn sysmem_buffer(device: &wgpu::Device, size: u64, cached: bool) -> Option<SysBuf> {
use ash::vk;
if size == 0 || size > device.limits().max_buffer_size {
return None;
}
// SAFETY: plain Vulkan object creation on the device wgpu owns; every
// failure path releases what was created.
unsafe {
let hal = device.as_hal::<wgpu::hal::api::Vulkan>()?;
let raw = hal.raw_device();
let props = hal
.shared_instance()
.raw_instance()
.get_physical_device_memory_properties(hal.raw_physical_device());
let info = vk::BufferCreateInfo::default()
.size(size)
.usage(vk::BufferUsageFlags::TRANSFER_SRC | vk::BufferUsageFlags::TRANSFER_DST)
.sharing_mode(vk::SharingMode::EXCLUSIVE);
let buf = raw.create_buffer(&info, None).ok()?;
let req = raw.get_buffer_memory_requirements(buf);
let want = vk::MemoryPropertyFlags::HOST_VISIBLE
| vk::MemoryPropertyFlags::HOST_COHERENT
| if cached {
vk::MemoryPropertyFlags::HOST_CACHED
} else {
vk::MemoryPropertyFlags::empty()
};
let ty = (0..props.memory_type_count).find(|&i| {
let f = props.memory_types[i as usize].property_flags;
req.memory_type_bits & (1 << i) != 0
&& f.contains(want)
&& !f.contains(vk::MemoryPropertyFlags::DEVICE_LOCAL)
&& (cached || !f.contains(vk::MemoryPropertyFlags::HOST_CACHED))
});
let Some(ty) = ty else {
raw.destroy_buffer(buf, None);
return None;
};
let alloc = vk::MemoryAllocateInfo::default()
.allocation_size(req.size)
.memory_type_index(ty);
let Ok(mem) = raw.allocate_memory(&alloc, None) else {
raw.destroy_buffer(buf, None);
return None;
};
if raw.bind_buffer_memory(buf, mem, 0).is_err() {
raw.destroy_buffer(buf, None);
raw.free_memory(mem, None);
return None;
}
let Ok(ptr) = raw.map_memory(mem, 0, vk::WHOLE_SIZE, vk::MemoryMapFlags::empty()) else {
raw.destroy_buffer(buf, None);
raw.free_memory(mem, None);
return None;
};
let hb = wgpu::hal::vulkan::Buffer::from_raw_managed(buf, mem, 0, req.size);
drop(hal);
let buffer = device.create_buffer_from_hal::<wgpu::hal::api::Vulkan>(
hb,
&wgpu::BufferDescriptor {
label: Some("cortiq-sysmem"),
size,
usage: wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
},
);
Some(SysBuf {
buffer,
ptr: ptr as *mut u8,
size,
cached,
})
}
}
#[cfg(not(target_os = "linux"))]
pub fn sysmem_buffer(_device: &wgpu::Device, _size: u64, _cached: bool) -> Option<SysBuf> {
None
}