use super::*;
use std::{
collections::VecDeque,
sync::{
Mutex,
atomic::{AtomicI64, AtomicU32, AtomicU64, Ordering},
},
time::{Duration, Instant},
};
#[derive(Clone)]
pub struct AqlQueue(Arc<QueueState>);
struct QueueState {
device: Device,
raw: *mut amdf_user_queue_t,
mapping: *mut amdf_user_queue_mapping_t,
info: amdf_user_queue_mapping_info_t,
scratch: Option<Buffer>,
maximum_private_bytes: u32,
budget: Option<crate::residency::MemoryBudget>,
producer: Mutex<Producer>,
}
struct Producer {
write: u64,
pending: VecDeque<(Arc<Work>, u64)>,
}
unsafe impl Send for QueueState {}
unsafe impl Sync for QueueState {}
struct Work {
signal: Buffer,
buffers: Vec<Buffer>,
_kernel: Kernel,
leases: Mutex<Vec<memory::DeviceUse>>,
active: AtomicU64,
retired: AtomicU64,
}
impl Work {
fn refresh(&self, point: u64) -> Result<bool> {
if self.retired.load(Ordering::Acquire) >= point {
return Ok(true);
}
let value = unsafe {
(&*self.signal.host_pointer().add(8).cast::<AtomicI64>()).load(Ordering::Acquire)
};
let completed = i64::MAX
.checked_sub(value)
.filter(|v| *v >= 0)
.ok_or_else(|| Error::DeviceLost("invalid AQL completion signal".into()))?
as u64;
if completed < point {
return Ok(false);
}
let mut leases = self
.leases
.lock()
.map_err(|_| Error::DeviceLost("AQL leases poisoned".into()))?;
if completed >= self.active.load(Ordering::Acquire) {
leases.clear();
}
self.retired.fetch_max(completed, Ordering::Release);
Ok(true)
}
}
impl Drop for QueueState {
fn drop(&mut self) {
let pending = std::mem::take(
&mut self
.producer
.get_mut()
.unwrap_or_else(|e| e.into_inner())
.pending,
);
let deadline = Instant::now() + Duration::from_secs(10);
for (work, point) in &pending {
loop {
match work.refresh(*point) {
Ok(true) => break,
Ok(false) if Instant::now() < deadline => std::thread::yield_now(),
_ => {
std::mem::forget(pending);
std::mem::forget(self.device.clone());
std::mem::forget(self.scratch.take());
return;
}
}
}
}
let api = self.device.0.endpoint.0.instance.api.core;
unsafe {
if (!self.mapping.is_null()
&& api.user_queue_mapping_destroy.unwrap()(self.mapping) != 0)
|| api.user_queue_destroy.unwrap()(self.raw) != 0
{
std::mem::forget(self.device.clone());
std::mem::forget(self.scratch.take());
}
}
}
}
#[derive(Clone)]
pub struct PreparedAql {
queue: AqlQueue,
work: Arc<Work>,
packets: Vec<[u32; 16]>,
next: Arc<Mutex<u64>>,
storage_bytes: usize,
}
#[derive(Clone)]
pub struct AqlCompletion {
queue: AqlQueue,
work: Arc<Work>,
point: u64,
}
impl AqlCompletion {
pub fn is_complete(&self) -> bool {
self.work.retired.load(Ordering::Acquire) >= self.point
}
pub fn refresh(&self) -> Result<bool> {
if !self.work.refresh(self.point)? {
let api = self.queue.0.device.0.endpoint.0.instance.api.core;
let mut info = amdf_user_queue_status_t {
type_: AMDF_STRUCTURE_TYPE_USER_QUEUE_STATUS,
structure_size: size_of::<amdf_user_queue_status_t>() as u32,
..Default::default()
};
check("AQL queue status", unsafe {
entry!(api, user_queue_query_status)(self.queue.0.raw, &mut info)
})?;
if info.state != AMDF_QUEUE_STATE_ACTIVE {
return Err(Error::DeviceLost(format!(
"AQL queue status {}",
info.terminal_status
)));
}
return Ok(false);
}
let mut producer = self
.queue
.0
.producer
.lock()
.map_err(|_| Error::DeviceLost("AQL producer poisoned".into()))?;
while let Some((work, point)) = producer.pending.front() {
if !work.refresh(*point)? {
break;
}
producer.pending.pop_front();
}
Ok(true)
}
pub fn wait_timeout(&self, timeout: Duration) -> Result<bool> {
let start = Instant::now();
loop {
if self.refresh()? {
return Ok(true);
}
if start.elapsed() >= timeout {
return Ok(false);
}
std::thread::yield_now();
}
}
pub fn wait(&self) -> Result<()> {
while !self.refresh()? {
std::thread::yield_now();
}
Ok(())
}
}
impl Device {
pub fn aql_queue(&self, maximum_private_bytes: u32) -> Result<AqlQueue> {
self.aql_queue_budgeted(maximum_private_bytes, None)
}
pub fn aql_queue_budgeted(
&self,
maximum_private_bytes: u32,
budget: Option<&crate::residency::MemoryBudget>,
) -> Result<AqlQueue> {
if self.endpoint().engine() != Engine::Gpu {
return Err(Error::Unsupported("AQL requires a GPU".into()));
}
let family = self
.endpoint()
.queue_capabilities()?
.into_iter()
.find(|f| {
f.command == QueueCommand::Aql
&& f.format_version == 1
&& f.user_publication
&& f.system_release
&& f.system_acquire
})
.ok_or_else(|| Error::Unsupported("no native AQL v1 family".into()))?;
let api = &self.0.endpoint.0.instance.api;
let _ = entry!(api.core, user_queue_destroy);
let _ = entry!(api.core, user_queue_mapping_destroy);
let mut scratch_info = amdf_gpu_queue_scratch_t::default();
let scratch = if maximum_private_bytes != 0 {
let mut endpoint = amdf_gpu_endpoint_info_t {
type_: HRX_AMDF_STRUCTURE_TYPE_GPU_ENDPOINT_INFO,
structure_size: size_of::<amdf_gpu_endpoint_info_t>() as u32,
..Default::default()
};
check("AQL scratch topology", unsafe {
entry!(api.gpu, endpoint_query_info)(self.0.endpoint.0.raw, &mut endpoint)
})?;
let engines = u64::from(endpoint.topology.xcc_count)
* u64::from(endpoint.topology.shader_engine_count_per_xcc);
if engines == 0 {
return Err(Error::Unsupported("empty AQL scratch topology".into()));
}
let slots = u64::from(endpoint.compute.compute_unit_count).div_ceil(engines)
* engines
* u64::from(endpoint.compute.maximum_scratch_wave_count_per_compute_unit);
let bytes = (u64::from(maximum_private_bytes) * 64)
.next_multiple_of(1024)
.checked_mul(slots)
.filter(|bytes| {
*bytes != 0 && *bytes <= 512 * 1024 * 1024 && slots <= u64::from(u32::MAX)
})
.ok_or_else(|| {
Error::Unsupported("AQL scratch exceeds the 512 MiB admission limit".into())
})?;
let buffer = self.fabric().allocate_owned(
bytes as usize,
self,
AMDF_MEMORY_ACCESS_READ | AMDF_MEMORY_ACCESS_WRITE,
4096,
false,
budget,
)?;
scratch_info = amdf_gpu_queue_scratch_t {
memory: buffer.0.raw,
byte_length: bytes,
maximum_private_segment_byte_length: maximum_private_bytes,
maximum_wave_count: slots as u32,
..Default::default()
};
Some(buffer)
} else {
None
};
let options = amdf_gpu_user_queue_create_info_t {
type_: HRX_AMDF_STRUCTURE_TYPE_GPU_USER_QUEUE_CREATE_INFO,
structure_size: size_of::<amdf_gpu_user_queue_create_info_t>() as u32,
queue_family_ordinal: family.ordinal,
priority: AMDF_QUEUE_PRIORITY_NORMAL,
producer_mode: AMDF_QUEUE_PRODUCER_MODE_SINGLE,
required_capabilities: AMDF_USER_QUEUE_CAPABILITY_HOST_PRODUCER as u64,
scratch: scratch_info,
..Default::default()
};
let mut raw = ptr::null_mut();
check("AQL queue create", unsafe {
entry!(api.gpu, user_queue_create)(self.0.raw, &options, &mut raw)
})?;
if raw.is_null() {
return Err(missing("AQL queue"));
}
let mut queue = QueueState {
device: self.clone(),
raw,
mapping: ptr::null_mut(),
info: Default::default(),
scratch,
maximum_private_bytes,
budget: budget.cloned(),
producer: Mutex::new(Producer {
write: 0,
pending: VecDeque::with_capacity(4096),
}),
};
check("AQL queue map", unsafe {
entry!(api.core, user_queue_map)(raw, ptr::null_mut(), &mut queue.mapping)
})?;
queue.info.type_ = AMDF_STRUCTURE_TYPE_USER_QUEUE_MAPPING_INFO;
queue.info.structure_size = size_of::<amdf_user_queue_mapping_info_t>() as u32;
check("AQL mapping info", unsafe {
entry!(api.core, user_queue_mapping_query_info)(queue.mapping, &mut queue.info)
})?;
let info = &queue.info;
if info.command_type != AMDF_QUEUE_COMMAND_TYPE_GPU_AQL
|| info.format_version != 1
|| info.index_bits != 64
|| info.doorbell_bits != 64
|| !info.ring_byte_length.is_power_of_two()
|| info.ring_byte_length < 128
|| [
info.ring_address,
info.read_index_address,
info.write_index_address,
info.doorbell_address,
]
.iter()
.any(|address| *address == 0 || address % 8 != 0)
{
return Err(Error::Unsupported("unexpected AQL mapping contract".into()));
}
Ok(AqlQueue(Arc::new(queue)))
}
}
impl AqlQueue {
pub fn device(&self) -> &Device {
&self.0.device
}
pub unsafe fn prepare(
&self,
kernel: &Kernel,
grid: [u32; 3],
block: [u16; 3],
arguments: &[Argument<'_>],
) -> Result<PreparedAql> {
if !Arc::ptr_eq(&kernel.device().0, &self.device().0) {
return Err(Error::Message(
"AQL kernel belongs to another device".into(),
));
}
if kernel.0.info.private_bytes > self.0.maximum_private_bytes {
return Err(Error::Unsupported(
"kernel exceeds fixed AQL scratch capacity".into(),
));
}
let mut extents = [0; 3];
for axis in 0..3 {
if block[axis] == 0
|| grid[axis] == 0
|| (kernel.workgroup_size()[axis] != 0
&& kernel.workgroup_size()[axis] != u32::from(block[axis]))
{
return Err(Error::Message(
"AQL launch geometry does not match kernel".into(),
));
}
extents[axis] = grid[axis]
.checked_mul(u32::from(block[axis]))
.ok_or_else(|| Error::Message("AQL grid extent overflow".into()))?;
}
if block.iter().map(|v| u64::from(*v)).product::<u64>() > 1024 {
return Err(Error::Message(
"AQL workgroup exceeds gfx1151 capacity".into(),
));
}
let specialized = kernel.for_aql(
grid,
self.0.info.ring_address,
self.0.info.ring_byte_length / 64 - 1,
arguments,
)?;
let clear = if let Some(state) = &specialized.0.race_state {
let (utility, clear_grid) =
state.clear_kernel(self.device(), kernel.shadow_budget())?;
Some(unsafe {
self.prepare(
&utility,
clear_grid,
[64, 1, 1],
&[Argument::Buffer(&state.shadow, 0)],
)
}?)
} else {
None
};
let kernel = &specialized;
let (packed, mut buffers) = kernel.arguments(arguments)?;
let fabric = self.device().fabric();
let kernarg = fabric.allocate_owned(
packed.len().max(64),
self.device(),
AMDF_MEMORY_ACCESS_READ | AMDF_MEMORY_ACCESS_WRITE,
64,
true,
self.0.budget.as_ref(),
)?;
kernarg.write(0, &packed)?;
let kernarg_address = kernarg.device_address(self.device())?;
let descriptor = kernel
.0
.code
.device_address(self.device())?
.checked_add(kernel.0.info.descriptor_offset)
.ok_or_else(|| Error::Message("AQL descriptor address overflow".into()))?;
let invalidation = [
0xc0065800u32,
0,
u32::MAX,
0xff,
0,
0,
0x0a,
1 | (1 << 4) | (1 << 5) | (1 << 7) | (1 << 8) | (1 << 9) | (1 << 14) | (1 << 15),
];
let publication = fabric.allocate_owned(
4096,
self.device(),
AMDF_MEMORY_ACCESS_READ | AMDF_MEMORY_ACCESS_EXECUTE,
4096,
false,
self.0.budget.as_ref(),
)?;
publication.write(
0,
&invalidation
.into_iter()
.flat_map(u32::to_le_bytes)
.collect::<Vec<_>>(),
)?;
let ib = publication.device_address(self.device())?;
if ib >= (1 << 48) || ib + 32 > (1 << 48) {
return Err(Error::Unsupported(
"AQL publication IB is outside the 48-bit address domain".into(),
));
}
let signal = fabric.allocate_owned(
64,
self.device(),
AMDF_MEMORY_ACCESS_READ | AMDF_MEMORY_ACCESS_WRITE,
64,
true,
self.0.budget.as_ref(),
)?;
signal.write(0, &1i64.to_le_bytes())?;
signal.write(8, &i64::MAX.to_le_bytes())?;
let signal_address = signal.device_address(self.device())?;
let mut publish = [0; 16];
publish[0] = (1 << 16) | (1 << 8);
publish[1] = 0xc0023f00;
publish[2] = ib as u32;
publish[3] = (ib >> 32) as u32;
publish[4] = (1 << 23) | 8;
publish[5] = 10;
let mut dispatch = [0; 16];
dispatch[0] = 2 | (1 << 8) | (2 << 9) | (2 << 11) | (3 << 16);
dispatch[1] = u32::from(block[0]) | (u32::from(block[1]) << 16);
dispatch[2] = u32::from(block[2]);
dispatch[3..6].copy_from_slice(&extents);
dispatch[6] = kernel.0.info.private_bytes;
dispatch[7] = kernel.0.info.local_bytes;
dispatch[8] = descriptor as u32;
dispatch[9] = (descriptor >> 32) as u32;
dispatch[10] = kernarg_address as u32;
dispatch[11] = (kernarg_address >> 32) as u32;
dispatch[14] = signal_address as u32;
dispatch[15] = (signal_address >> 32) as u32;
let mut packets = vec![publish];
if let Some(clear) = clear {
let mut packet = *clear.packets.last().unwrap();
packet[14..16].fill(0);
packets.push(packet);
buffers.extend(clear.work.buffers.iter().cloned());
}
packets.push(dispatch);
buffers.extend([kernarg, publication]);
buffers.sort_by_key(Buffer::identity);
buffers.dedup_by(|a, b| a.same_backing(b));
let storage_bytes = signal.len()
+ buffers
.iter()
.filter(|buffer| {
!arguments.iter().any(|argument| {
matches!(argument,
Argument::Buffer(bound, _) if buffer.same_backing(bound))
})
})
.map(Buffer::len)
.sum::<usize>();
let leases = Mutex::new(Vec::with_capacity(buffers.len()));
Ok(PreparedAql {
queue: self.clone(),
work: Arc::new(Work {
signal,
buffers,
_kernel: kernel.clone(),
leases,
active: AtomicU64::new(0),
retired: AtomicU64::new(0),
}),
packets,
next: Arc::new(Mutex::new(0)),
storage_bytes,
})
}
}
impl PreparedAql {
pub(crate) fn storage_bytes(&self) -> usize {
self.storage_bytes
}
pub unsafe fn dispatch(&self) -> Result<AqlCompletion> {
let mut next = self
.next
.lock()
.map_err(|_| Error::DeviceLost("AQL dispatch poisoned".into()))?;
let point = next
.checked_add(1)
.filter(|point| *point <= i64::MAX as u64)
.ok_or_else(|| Error::DeviceLost("AQL timeline exhausted".into()))?;
let mut producer = self
.queue
.0
.producer
.lock()
.map_err(|_| Error::DeviceLost("AQL producer poisoned".into()))?;
while let Some((work, point)) = producer.pending.front() {
if !work.refresh(*point)? {
break;
}
producer.pending.pop_front();
}
let info = &self.queue.0.info;
let capacity = info.ring_byte_length / 64;
let read =
unsafe { (&*(info.read_index_address as *const AtomicU64)).load(Ordering::Acquire) };
let end = producer
.write
.checked_add(self.packets.len() as u64)
.ok_or_else(|| Error::DeviceLost("AQL ring index exhausted".into()))?;
if producer.pending.len() == 4096 || end.saturating_sub(read) > capacity {
return Err(Error::Busy("AQL ring or pending capacity exhausted".into()));
}
let mut leases = self
.work
.leases
.lock()
.map_err(|_| Error::DeviceLost("AQL leases poisoned".into()))?;
if leases.is_empty() {
for buffer in &self.work.buffers {
match buffer.retain_use() {
Ok(lease) => leases.push(lease),
Err(error) => {
leases.clear();
return Err(error);
}
}
}
}
self.work.active.store(point, Ordering::Release);
producer.pending.push_back((self.work.clone(), point));
unsafe {
(&*(info.write_index_address as *const AtomicU64)).store(end, Ordering::Release);
for (offset, packet) in self.packets.iter().enumerate() {
let index = producer.write + offset as u64;
let slot =
(info.ring_address as *mut u32).add(((index & (capacity - 1)) * 16) as usize);
ptr::copy_nonoverlapping(packet.as_ptr().add(1), slot.add(1), 15);
(&*slot.cast::<AtomicU32>()).store(packet[0], Ordering::Release);
std::sync::atomic::fence(Ordering::SeqCst);
ptr::write_volatile(info.doorbell_address as *mut u64, index);
}
}
producer.write = end;
*next = point;
Ok(AqlCompletion {
queue: self.queue.clone(),
work: self.work.clone(),
point,
})
}
}
#[cfg(test)]
mod sanitizer_tests {
use super::*;
use crate::loom::{
Compiler, CompilerOptions, SanitizerChecks, SanitizerOptions, SanitizerReporting,
Specialization,
};
#[test]
#[ignore = "requires current native compiler/bridge and gfx1151"]
fn address_shadow_guards_null_wrap_and_every_native_load_width() -> Result<()> {
let device = Device::open(Engine::Gpu, 0)?;
let queue = device.aql_queue(0)?;
let compiler = Compiler::shared(
None,
CompilerOptions {
sanitizer: SanitizerOptions {
checks: SanitizerChecks {
access: true,
..Default::default()
},
reporting: SanitizerReporting::ReportOnly,
},
..Default::default()
},
)?;
for wave in [32, 64] {
for width in [1, 2, 4, 8, 16, 32, 64, 128] {
let source = format!(
r#"
amdgpu.target<gfx1151> @target {{subgroup_size = {wave}}}
kernel.def target(@target) @check_range() {{
%one = index.constant 1 : index
kernel.launch.config workgroups(%one, %one, %one) workgroup_size(%one, %one, %one) : index
}} launch(%input: buffer, %output: buffer) {{
%zero = index.constant 0 : offset
%index = index.constant 0 : index
%value = scalar.constant 9 : i32
%global = buffer.assume.memory_space<global> %input : buffer
%view = buffer.view %global[%zero] : buffer -> view<{width}xi8>
sanitizer.assert.access<read> %view[0] {{static_extents = [{width}]}} : view<{width}xi8>
%out = buffer.assume.memory_space<global> %output : buffer
%out_view = buffer.view %out[%zero] : buffer -> view<1xi32>
view.store %value, %out_view[%index] : i32, view<1xi32>
kernel.return
}}
"#
);
let artifact = compiler
.module(&source)
.compile(&Specialization::new("check_range"))?;
let kernel = unsafe {
device.load_sanitized(&artifact, &SanitizerRuntimeOptions::default())
}?;
let input = device
.fabric()
.allocate(width + 1, std::slice::from_ref(&device))?;
let output = device.fabric().allocate(4, std::slice::from_ref(&device))?;
let prepared = unsafe {
queue.prepare(
&kernel,
[1; 3],
[1; 3],
&[Argument::Buffer(&input, 0), Argument::Buffer(&output, 0)],
)
}?;
let packet = prepared.packets.last().unwrap();
let kernarg_address = u64::from(packet[10]) | (u64::from(packet[11]) << 32);
let kernarg = prepared
.work
.buffers
.iter()
.find(|b| b.device_address(&device).ok() == Some(kernarg_address))
.unwrap();
let base = input.device_address(&device)?;
for (pointer, valid) in [
(base, true),
(base + 1, true),
(base + 2, false),
(base - 1, false),
(0, false),
(1, false),
(u64::MAX - 1, false),
(u64::MAX - 7, false),
(1 << 63, false),
] {
kernarg.write(0, &pointer.to_le_bytes())?;
output.write(0, &17u32.to_le_bytes())?;
let done = unsafe { prepared.dispatch() }?;
assert!(
done.wait_timeout(Duration::from_secs(10))?,
"wave{wave} width{width} pointer{pointer:x}"
);
let reports = kernel.sanitizer_reports()?;
assert_eq!(reports.dropped, 0);
assert_eq!(
reports.reports.len(),
usize::from(!valid),
"wave{wave} width{width} pointer{pointer:x}: {reports:?}"
);
if !valid {
let address = reports.reports[0].address.as_ref().unwrap();
assert_eq!(address.fault_address, pointer);
assert_eq!(address.access_bytes, width as u64);
}
let mut actual = [0; 4];
output.read(0, &mut actual)?;
assert_eq!(u32::from_le_bytes(actual), if valid { 9 } else { 17 });
}
}
}
Ok(())
}
#[test]
#[ignore = "requires current native compiler/bridge and gfx1151"]
fn address_checks_keep_valid_and_invalid_lanes_separate() -> Result<()> {
let device = Device::open(Engine::Gpu, 0)?;
let queue = device.aql_queue(0)?;
let compiler = Compiler::shared(
None,
CompilerOptions {
sanitizer: SanitizerOptions {
checks: SanitizerChecks {
access: true,
..Default::default()
},
reporting: SanitizerReporting::ReportOnly,
},
..Default::default()
},
)?;
for wave in [32, 64] {
let source = format!(
r#"
amdgpu.target<gfx1151> @target {{subgroup_size = {wave}}}
kernel.def target(@target) @check_lanes() {{
%one = index.constant 1 : index
%threads = index.constant 64 : index
kernel.launch.config workgroups(%one, %one, %one) workgroup_size(%threads, %one, %one) : index
}} launch(%input: buffer, %output: buffer) {{
%zero = index.constant 0 : offset
%lane = kernel.workitem.id<x> : index
%value = scalar.constant 9 : i32
%global = buffer.assume.memory_space<global> %input : buffer
%view = buffer.view %global[%zero] : buffer -> view<64xi8>
sanitizer.assert.access<read> %view[%lane] : view<64xi8>
%out = buffer.assume.memory_space<global> %output : buffer
%out_view = buffer.view %out[%zero] : buffer -> view<64xi32>
view.store %value, %out_view[%lane] : i32, view<64xi32>
kernel.return
}}
"#
);
let artifact = compiler
.module(&source)
.compile(&Specialization::new("check_lanes"))?;
let kernel =
unsafe { device.load_sanitized(&artifact, &SanitizerRuntimeOptions::default()) }?;
let input = device
.fabric()
.allocate(33, std::slice::from_ref(&device))?;
let output = device
.fabric()
.allocate(256, std::slice::from_ref(&device))?;
output.write(0, &[0; 256])?;
let prepared = unsafe {
queue.prepare(
&kernel,
[1; 3],
[64, 1, 1],
&[Argument::Buffer(&input, 0), Argument::Buffer(&output, 0)],
)
}?;
let done = unsafe { prepared.dispatch() }?;
assert!(done.wait_timeout(Duration::from_secs(10))?);
let reports = kernel.sanitizer_reports()?;
assert_eq!(reports.dropped, 0);
assert!(!reports.reports.is_empty());
let base = input.device_address(&device)?;
for report in reports.reports {
let address = report.address.unwrap();
assert!((base + 33..base + 64).contains(&address.fault_address));
assert!((33..64).contains(&report.workitem_x));
}
if wave == 32 {
let mut first_wave = [0; 128];
output.read(0, &mut first_wave)?;
for word in first_wave.as_chunks::<4>().0 {
assert_eq!(*word, 9u32.to_le_bytes());
}
}
}
Ok(())
}
}