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%
---------+----------+--------------+-----------------+-----------------+--------------

Author: arieleizenberg

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
 
 

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%}"
+        )

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"}