fixed type for overflow warning

argonne-lcf · Oct 10, 2024 · b05eb73 · b05eb73
1 parent 6b50605
commit b05eb73
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/dlio_benchmark/data_generator/indexed_binary_generator.py b/dlio_benchmark/data_generator/indexed_binary_generator.py
@@ -73,16 +73,16 @@ def generate(self):
                     logging.info(f"{utcnow()} Starting metadata generation. ")
                 fh_off = MPI.File.Open(comm, out_path_spec_off_idx, amode)
                 fh_sz = MPI.File.Open(comm, out_path_spec_sz_idx, amode)
-                elements_per_loop = min(int(MB / np.dtype(np.uint32).itemsize), samples_per_rank)
+                off_type = np.uint64
+                elements_per_loop = min(int(MB / np.dtype(off_type).itemsize), samples_per_rank)
                 offsets_processed=0
                 for element_index in range(self.my_rank*samples_per_rank, samples_per_rank*(self.my_rank+1), elements_per_loop):
-                    myfmt = 'Q' * elements_per_loop
                     offsets = np.array(range(self.my_rank * elements_per_loop * sample_size, 
                                     (self.my_rank + 1) * elements_per_loop * sample_size, 
-                                    sample_size), dtype=np.uint32)
+                                    sample_size), dtype=off_type)
 
-                    sizes = np.array([sample_size] * elements_per_loop, dtype=np.uint32)
-                    offset = element_index * np.dtype(np.uint32).itemsize
+                    sizes = np.array([sample_size] * elements_per_loop, dtype=off_type)
+                    offset = element_index * np.dtype(off_type).itemsize
                     fh_off.Write_at_all(offset, offsets)
                     fh_sz.Write_at_all(offset, sizes)
                     offsets_processed += elements_per_loop

diff --git a/dlio_benchmark/reader/indexed_binary_mmap_reader.py b/dlio_benchmark/reader/indexed_binary_mmap_reader.py
@@ -57,10 +57,10 @@ def load_index_file(self, global_sample_idx, filename, sample_index):
             self.file_map_ibr[filename] = []
             bin_buffer_mmap = np.memmap(offset_file, mode='r', order='C')
             bin_buffer = memoryview(bin_buffer_mmap)
-            self.file_map_ibr[filename].append(np.frombuffer(bin_buffer, dtype=np.uint32))
+            self.file_map_ibr[filename].append(np.frombuffer(bin_buffer, dtype=np.uint64))
             bin_buffer_mmap = np.memmap(sz_file, mode='r', order='C')
             bin_buffer = memoryview(bin_buffer_mmap)
-            self.file_map_ibr[filename].append(np.frombuffer(bin_buffer, dtype=np.uint32))
+            self.file_map_ibr[filename].append(np.frombuffer(bin_buffer, dtype=np.uint64))
             bin_buffer_mmap = np.memmap(filename, mode='r', order='C')
             bin_buffer = memoryview(bin_buffer_mmap)
             self.buffer_map[filename] = np.frombuffer(bin_buffer, dtype=np.uint8)