benchmark/common: correctly aligning argument of benchmark_kernel

src/pipedream/benchmark/common.py

@@ -1666,6 +1666,23 @@ class _Benchmark_Runner:
             benchmarks=benchmark_functions,
         )
 
+    # Assign and pre-fil with random values an array
+    @staticmethod
+    def aligned_alloc_byte_array(length, align: int, rand_values=False):
+        buf = (ctypes.c_byte * (length + align - 1))()
+        if rand_values:
+            random.seed(42)
+            for i in range(length + align - 1):
+                buf[i] = random.getrandbits(8)
+
+        addr = ctypes.addressof(buf)
+        offset = (align - (addr % align)) % align
+        assert (addr + offset) % align == 0
+        buf = (ctypes.c_byte * length).from_buffer(buf, offset)
+        addr = ctypes.addressof(buf)
+        assert addr % align == 0
+        return buf
+
     def _run_benchmark(
         self,
         *,
@@ -1737,45 +1754,20 @@ class _Benchmark_Runner:
             result_array.fill(42)
 
             ## values used to initialise vector registers
-            ## align on one page, because we are never too cautious
-            ALIGNEMENT = benchmark.arch.max_vector_size
-            init_values = numpy.ndarray(
-                shape=[
-                    benchmark.arch.nb_vector_reg * ALIGNEMENT - 1,
-                    benchmark.arch.max_vector_size,
-                ],
-                dtype=ctypes.c_byte,
-                order="C",
+            ## align on max_vector_size * 2, we are never too cautious
+            ALIGNMENT = benchmark.arch.max_vector_size
+            init_values = self.aligned_alloc_byte_array(
+                benchmark.arch.nb_vector_reg * benchmark.arch.max_vector_size,
+                ALIGNMENT,
+                rand_values=True,
             )
-            load_area = numpy.ndarray(
-                shape=[
-                    self.memory_size(benchmark),
-                    1,
-                ],
-                dtype=ctypes.c_byte,
-                order="C",
-            )
-            store_area = numpy.ndarray(
-                shape=[
-                    self.memory_size(benchmark),
-                    1,
-                ],
-                dtype=ctypes.c_byte,
-                order="C",
+            load_area = self.aligned_alloc_byte_array(
+                self.memory_size(benchmark),
+                ALIGNMENT,
             )
-
-            random.seed(42)
-            for i in range(benchmark.arch.nb_vector_reg * ALIGNEMENT - 1):
-                for j in range(benchmark.arch.max_vector_size):
-                    init_values[i, j] = random.getrandbits(8)
-
-            data = init_values.ctypes.data_as(ctypes.POINTER(ctypes.c_byte))
-            load_area_data = load_area.ctypes.data_as(ctypes.POINTER(ctypes.c_byte))
-            store_area_data = store_area.ctypes.data_as(ctypes.POINTER(ctypes.c_byte))
-            addr = ctypes.addressof(data)
-            offset = 0 if addr % ALIGNEMENT == 0 else ALIGNEMENT - (addr % ALIGNEMENT)
-            aligned_init_values = ctypes.POINTER(ctypes.c_byte).from_address(
-                addr + offset
+            store_area = self.aligned_alloc_byte_array(
+                self.memory_size(benchmark),
+                ALIGNMENT,
             )
 
             time_before = time.perf_counter()
@@ -1783,9 +1775,9 @@ class _Benchmark_Runner:
                 event_set_id,
                 num_events,
                 result_array.ctypes.data_as(ctypes.POINTER(ctypes.c_longlong)),
-                aligned_init_values,
-                load_area_data,
-                store_area_data,
+                init_values,
+                load_area,
+                store_area,
             )
             time_after = time.perf_counter()