From d7d38256e0580c3cfb649a641d9ed62c6ff0fc20 Mon Sep 17 00:00:00 2001
From: Andrzej Janik <vosen@vosen.pl>
Date: Sun, 27 Jun 2021 13:08:46 +0200
Subject: [PATCH] Fix offset calculation in kernel launch

---
 zluda/src/impl/function.rs | 18 ++++++++++--------
 zluda_dump/src/lib.rs      | 17 +++++++++--------
 2 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/zluda/src/impl/function.rs b/zluda/src/impl/function.rs
index 4f2006c..638d08a 100644
--- a/zluda/src/impl/function.rs
+++ b/zluda/src/impl/function.rs
@@ -110,23 +110,24 @@ pub fn launch_kernel(
             match (buffer_size, buffer_ptr) {
                 (Some(buffer_size), Some(buffer_ptr)) => {
                     let sum_of_kernel_argument_sizes =
-                        func.arg_size.iter().fold(0, |offset, size_of_arg| {
-                            size_of_arg + round_up_to_multiple(offset, *size_of_arg)
-                        });
+                        func.arg_size
+                            .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 {
                         return Err(CUresult::CUDA_ERROR_INVALID_VALUE);
                     }
                     let mut offset = 0;
                     for (i, arg_size) in func.arg_size.iter().enumerate() {
-                        let buffer_offset = round_up_to_multiple(offset, *arg_size);
                         unsafe {
                             func.base.set_arg_raw(
                                 i as u32,
                                 *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),
@@ -159,8 +160,9 @@ pub fn launch_kernel(
     })?
 }
 
-fn round_up_to_multiple(x: usize, multiple: usize) -> usize {
-    ((x + multiple - 1) / multiple) * multiple
+fn align_to_usize(value: usize) -> usize {
+    let multiple = std::mem::size_of::<usize>();
+    ((value + multiple - 1) / multiple) * multiple
 }
 
 pub(crate) fn get_attribute(
diff --git a/zluda_dump/src/lib.rs b/zluda_dump/src/lib.rs
index b5a1e3a..450a00a 100644
--- a/zluda_dump/src/lib.rs
+++ b/zluda_dump/src/lib.rs
@@ -573,24 +573,24 @@ fn dump_arguments(
         }
         match (buffer_size, buffer_ptr) {
             (Some(buffer_size), Some(buffer_ptr)) => {
-                let sum_of_kernel_argument_sizes = args.iter().fold(0, |offset, size_of_arg| {
-                    size_of_arg + round_up_to_multiple(offset, *size_of_arg)
-                });
+                let sum_of_kernel_argument_sizes =
+                    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 {
                     return Err("Malformed `extra` parameter to kernel launch")?;
                 }
                 let mut offset = 0;
                 for (i, arg_size) in args.iter().enumerate() {
-                    let buffer_offset = round_up_to_multiple(offset, *arg_size);
                     unsafe {
                         dump_argument_to_file(
                             &dump_dir,
                             i,
                             *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")?,
@@ -599,8 +599,9 @@ fn dump_arguments(
     Ok(())
 }
 
-fn round_up_to_multiple(x: usize, multiple: usize) -> usize {
-    ((x + multiple - 1) / multiple) * multiple
+fn align_to_usize(value: usize) -> usize {
+    let multiple = std::mem::size_of::<usize>();
+    ((value + multiple - 1) / multiple) * multiple
 }
 
 unsafe fn dump_argument_to_file(