Fix offset calculation in kernel launch

This commit is contained in:
Andrzej Janik 2021-06-27 13:08:46 +02:00
commit d7d38256e0
2 changed files with 19 additions and 16 deletions

View file

@ -110,23 +110,24 @@ pub fn launch_kernel(
match (buffer_size, buffer_ptr) { match (buffer_size, buffer_ptr) {
(Some(buffer_size), Some(buffer_ptr)) => { (Some(buffer_size), Some(buffer_ptr)) => {
let sum_of_kernel_argument_sizes = let sum_of_kernel_argument_sizes =
func.arg_size.iter().fold(0, |offset, size_of_arg| { func.arg_size
size_of_arg + round_up_to_multiple(offset, *size_of_arg) .iter()
.fold(0, |sum_of_arg_sizes, size_of_arg| {
sum_of_arg_sizes + align_to_usize(*size_of_arg)
}); });
if buffer_size != sum_of_kernel_argument_sizes { if buffer_size != sum_of_kernel_argument_sizes {
return Err(CUresult::CUDA_ERROR_INVALID_VALUE); return Err(CUresult::CUDA_ERROR_INVALID_VALUE);
} }
let mut offset = 0; let mut offset = 0;
for (i, arg_size) in func.arg_size.iter().enumerate() { for (i, arg_size) in func.arg_size.iter().enumerate() {
let buffer_offset = round_up_to_multiple(offset, *arg_size);
unsafe { unsafe {
func.base.set_arg_raw( func.base.set_arg_raw(
i as u32, i as u32,
*arg_size, *arg_size,
buffer_ptr.add(buffer_offset) as *const _, buffer_ptr.add(offset) as *const _,
)? )?
}; };
offset = buffer_offset + *arg_size; offset += align_to_usize(*arg_size);
} }
} }
_ => return Err(CUresult::CUDA_ERROR_INVALID_VALUE), _ => return Err(CUresult::CUDA_ERROR_INVALID_VALUE),
@ -159,8 +160,9 @@ pub fn launch_kernel(
})? })?
} }
fn round_up_to_multiple(x: usize, multiple: usize) -> usize { fn align_to_usize(value: usize) -> usize {
((x + multiple - 1) / multiple) * multiple let multiple = std::mem::size_of::<usize>();
((value + multiple - 1) / multiple) * multiple
} }
pub(crate) fn get_attribute( pub(crate) fn get_attribute(

View file

@ -573,24 +573,24 @@ fn dump_arguments(
} }
match (buffer_size, buffer_ptr) { match (buffer_size, buffer_ptr) {
(Some(buffer_size), Some(buffer_ptr)) => { (Some(buffer_size), Some(buffer_ptr)) => {
let sum_of_kernel_argument_sizes = args.iter().fold(0, |offset, size_of_arg| { let sum_of_kernel_argument_sizes =
size_of_arg + round_up_to_multiple(offset, *size_of_arg) args.iter().fold(0, |sum_of_arg_sizes, size_of_arg| {
sum_of_arg_sizes + align_to_usize(*size_of_arg)
}); });
if buffer_size != sum_of_kernel_argument_sizes { if buffer_size != sum_of_kernel_argument_sizes {
return Err("Malformed `extra` parameter to kernel launch")?; return Err("Malformed `extra` parameter to kernel launch")?;
} }
let mut offset = 0; let mut offset = 0;
for (i, arg_size) in args.iter().enumerate() { for (i, arg_size) in args.iter().enumerate() {
let buffer_offset = round_up_to_multiple(offset, *arg_size);
unsafe { unsafe {
dump_argument_to_file( dump_argument_to_file(
&dump_dir, &dump_dir,
i, i,
*arg_size, *arg_size,
buffer_ptr.add(buffer_offset) as *const _, buffer_ptr.add(offset) as *const _,
)? )?
}; };
offset = buffer_offset + *arg_size; offset += align_to_usize(*arg_size);
} }
} }
_ => return Err("Malformed `extra` parameter to kernel launch")?, _ => return Err("Malformed `extra` parameter to kernel launch")?,
@ -599,8 +599,9 @@ fn dump_arguments(
Ok(()) Ok(())
} }
fn round_up_to_multiple(x: usize, multiple: usize) -> usize { fn align_to_usize(value: usize) -> usize {
((x + multiple - 1) / multiple) * multiple let multiple = std::mem::size_of::<usize>();
((value + multiple - 1) / multiple) * multiple
} }
unsafe fn dump_argument_to_file( unsafe fn dump_argument_to_file(