Improve DMatrix creation performance in python

The xgboost python python package serializes numpy arrays as json. This can take up a considerable amount of time in production workloads. This patch optimizes the specific case where the numpy is already in "C" contigous 32-bit floating point format, and can be loaded directly without the json layer. This can improve performance up to 35% in some cases, as can be seen by the microbenchmark added in xgboost/tests/python/microbench_numpy.py: Rows | Cols | Threads | Contiguous | Non-contiguous | Ratio ---------+----------+--------------+-----------------+-----------------+-------------- 15000 | 100 | 0 | 0.01686 | 0.01988 | 84.8% 15000 | 100 | 1 | 0.02897 | 0.04424 | 65.5% 15000 | 100 | 2 | 0.02579 | 0.0392 | 65.8% 15000 | 100 | 10 | 0.01581 | 0.02058 | 76.8% ---------+----------+--------------+-----------------+-----------------+-------------- 2 | 2000 | 0 | 0.001055 | 0.001205 | 87.6% 2 | 2000 | 1 | 0.0004465 | 0.0005689 | 78.5% 2 | 2000 | 2 | 0.0004609 | 0.000615 | 74.9% 2 | 2000 | 10 | 0.0005087 | 0.0005623 | 90.5% ---------+----------+--------------+-----------------+-----------------+--------------
dmlc · Jun 10, 2024 · 2e3adc3 · 2e3adc3
1 parent 0c44067
commit 2e3adc3
Show file tree

Hide file tree

Showing 3 changed files with 130 additions and 10 deletions.
diff --git a/python-package/xgboost/data.py b/python-package/xgboost/data.py
@@ -252,17 +252,30 @@ def _from_numpy_array(
     _check_data_shape(data)
     data, _ = _ensure_np_dtype(data, data.dtype)
     handle = ctypes.c_void_p()
-    _check_call(
-        _LIB.XGDMatrixCreateFromDense(
-            _array_interface(data),
-            make_jcargs(
-                missing=float(missing),
-                nthread=int(nthread),
-                data_split_mode=int(data_split_mode),
-            ),
-            ctypes.byref(handle),
+    if isinstance(data, np.ndarray) and data.flags['C_CONTIGUOUS'] and data.dtype == np.float32:
+        _check_call(
+            _LIB.XGDMatrixCreateFromMat_omp(
+                data.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
+                c_bst_ulong(data.shape[0]),
+                c_bst_ulong(data.shape[1]),
+                ctypes.c_float(missing),
+                ctypes.byref(handle),
+                ctypes.c_int(nthread),
+                ctypes.c_int(data_split_mode),
+            )
+        )
+    else:
+        _check_call(
+            _LIB.XGDMatrixCreateFromDense(
+                _array_interface(data),
+                make_jcargs(
+                    missing=float(missing),
+                    nthread=int(nthread),
+                    data_split_mode=int(data_split_mode),
+                ),
+                ctypes.byref(handle),
+            )
         )
-    )
     return handle, feature_names, feature_types
 
 

diff --git a/tests/python/microbench_numpy.py b/tests/python/microbench_numpy.py
@@ -0,0 +1,66 @@
+import numpy as np
+import xgboost as xgb
+from collections import defaultdict
+import timeit
+
+
+def create_dmatrix(X, y, test_contiguous, nthread):
+    if test_contiguous:
+        X = np.ascontiguousarray(X).astype(np.float32)
+        y = np.ascontiguousarray(y).astype(np.float32)
+        assert X.flags["C_CONTIGUOUS"]
+    else:
+        X = np.asfortranarray(X)
+        y = np.asfortranarray(y)
+        assert not X.flags["C_CONTIGUOUS"]
+
+    dm = xgb.DMatrix(X, y, nthread=nthread)
+    return dm
+
+
+def benchmark_dmatrix_creation(test_contiguous, nthread, rows, cols):
+    X = np.random.randn(rows, cols)
+    y = np.random.randn(
+        rows,
+    ).astype(np.float32)
+
+    start_time = timeit.default_timer()
+    dm = create_dmatrix(X, y, test_contiguous, nthread)
+    end_time = timeit.default_timer()
+
+    np.testing.assert_allclose(dm.get_data().toarray(), X, rtol=1e-7)
+    np.testing.assert_array_equal(dm.get_label(), y)
+    assert dm.num_row() == rows
+    assert dm.num_col() == cols
+
+    total_time = end_time - start_time
+
+    return total_time
+
+
+REPEATS = 10
+
+contiguous = defaultdict(float)
+noncontiguous = defaultdict(float)
+
+print(
+    f"{'Rows':8} | {'Cols':8} | {'Threads':12} | {'Contiguous':15} | {'Non-contiguous':15} | {'Ratio':12}"
+)
+
+for rows, cols, repeats in ((15000, 100, 10), (2, 2000, 200)):
+    for nthread in (0, 1, 2, 10):
+        for _ in range(repeats):
+            contiguous[nthread] += benchmark_dmatrix_creation(
+                test_contiguous=True, nthread=nthread, rows=rows, cols=cols
+            )
+            noncontiguous[nthread] += benchmark_dmatrix_creation(
+                test_contiguous=False, nthread=nthread, rows=rows, cols=cols
+            )
+
+    contiguous = {k: v / repeats for k, v in contiguous.items()}
+    noncontiguous = {k: v / repeats for k, v in noncontiguous.items()}
+
+    for k in contiguous:
+        print(
+            f"{rows:8} | {cols:8} | {k:12} | {contiguous[k]:15.4g} | {noncontiguous[k]:15.4g} | {contiguous[k] / noncontiguous[k]:12.1%}"
+        )
diff --git a/tests/python/test_basic.py b/tests/python/test_basic.py
@@ -210,6 +210,47 @@ def test_dmatrix_numpy_init_omp(self):
             assert dm.num_row() == row
             assert dm.num_col() == cols
 
+    def _test_dmatrix_numpy_init_omp_contiguous(self, test_contiguous: bool):
+        rows = [1000, 11326, 15000]
+        cols = 50
+        for row in rows:
+            X = np.random.randn(row, cols)
+            y = np.random.randn(row).astype("f")
+
+            # Ensure data is contiguous
+            if test_contiguous:
+                X = np.ascontiguousarray(X).astype(np.float32)
+                y = np.ascontiguousarray(y).astype(np.float32)
+                assert X.flags['C_CONTIGUOUS']
+            else:
+                X = np.asfortranarray(X)
+                y = np.asfortranarray(y)
+                assert not X.flags['C_CONTIGUOUS']
+
+            dm = xgb.DMatrix(X, y, nthread=0)
+            np.testing.assert_allclose(dm.get_data().toarray(), X, rtol=1e-7)
+            np.testing.assert_array_equal(dm.get_label(), y)
+            assert dm.num_row() == row
+            assert dm.num_col() == cols
+
+            dm = xgb.DMatrix(X, y, nthread=1)
+            np.testing.assert_allclose(dm.get_data().toarray(), X, rtol=1e-7)
+            np.testing.assert_array_equal(dm.get_label(), y)
+            assert dm.num_row() == row
+            assert dm.num_col() == cols
+
+            dm = xgb.DMatrix(X, y, nthread=10)
+            np.testing.assert_allclose(dm.get_data().toarray(), X, rtol=1e-7)
+            np.testing.assert_array_equal(dm.get_label(), y)
+            assert dm.num_row() == row
+            assert dm.num_col() == cols
+
+    def test_dmatrix_numpy_init_omp_contiguous(self):
+        return self._test_dmatrix_numpy_init_omp_contiguous(True)
+
+    def test_dmatrix_numpy_init_omp_not_contiguous(self):
+        return self._test_dmatrix_numpy_init_omp_contiguous(False)
+
     def test_cv(self):
         dm, _ = tm.load_agaricus(__file__)
         params = {"max_depth": 2, "eta": 1, "objective": "binary:logistic"}