[GraphBolt] Add gb.numpy_save_aligned.

python/dgl/graphbolt/__init__.py

@@ -51,7 +51,11 @@ def load_graphbolt():
 from .negative_sampler import *
 from .sampled_subgraph import *
 from .subgraph_sampler import *
-from .external_utils import add_reverse_edges, exclude_seed_edges
+from .external_utils import (
+    add_reverse_edges,
+    exclude_seed_edges,
+    numpy_save_aligned,
+)
 from .internal import (
     compact_csc_format,
     unique_and_compact,

python/dgl/graphbolt/external_utils.py

@@ -2,6 +2,7 @@
 
 from typing import Dict, Union
 
+import numpy
 import torch
 
 from .minibatch import MiniBatch
@@ -101,3 +102,14 @@ def exclude_seed_edges(
         for subgraph in minibatch.sampled_subgraphs
     ]
     return minibatch
+
+
+def numpy_save_aligned(*args, **kwargs):
+    """A wrapper for numpy.save(), ensures the array is stored 4KiB aligned."""
+    has_array_align = hasattr(numpy.lib.format, "ARRAY_ALIGN")
+    if has_array_align:
+        default_alignment = numpy.lib.format.ARRAY_ALIGN
+        numpy.lib.format.ARRAY_ALIGN = 4096
+    numpy.save(*args, **kwargs)
+    if has_array_align:
+        numpy.lib.format.ARRAY_ALIGN = default_alignment

python/dgl/graphbolt/internal/utils.py

@@ -11,6 +11,8 @@
 import torch
 from numpy.lib.format import read_array_header_1_0, read_array_header_2_0
 
+from ..external_utils import numpy_save_aligned
+
 
 def _read_torch_data(path):
     return torch.load(path)
@@ -54,7 +56,7 @@ def save_data(data, path, fmt):
                 "so it will be copied to contiguous memory."
             )
             data = np.ascontiguousarray(data)
-        np.save(path, data)
+        numpy_save_aligned(path, data)
     elif fmt == "torch":
         if not data.is_contiguous():
             Warning(