Improve DMatrix creation performance in python

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.dtype == np.float32 and data.flags['C_CONTIGUOUS'] and data.size <= 32768:
+        _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(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,56 @@
+import numpy as np
+import xgboost as xgb
+from collections import defaultdict
+import timeit
+import ctypes
+from xgboost.core import _LIB, DataSplitMode
+from xgboost.data import _check_call, _array_interface, c_bst_ulong, make_jcargs
+
+def measure_create_dmatrix(rows, cols, nthread, use_optimization):
+    data =  np.random.randn(rows, cols).astype(np.float32)
+    data = np.ascontiguousarray(data)
+
+    handle = ctypes.c_void_p()
+    missing = np.nan
+
+    start = timeit.default_timer()
+    if use_optimization:
+        _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(DataSplitMode.ROW),
+        )
+    else:
+        _LIB.XGDMatrixCreateFromDense(
+            _array_interface(data),
+            make_jcargs(
+                missing=float(missing),
+                nthread=int(nthread),
+                data_split_mode=int(DataSplitMode.ROW),
+            ),
+            ctypes.byref(handle),
+        )
+    end = timeit.default_timer()
+    return end - start
+
+COLS = 1000
+
+print(f"{'Threads':8} | {'Rows':8} | {'Cols':8} | {'Current (sec)':15} | {'Optimized (sec)':15} | {'Ratio':12}")
+
+for nthread in [1, 2, 4, 8]:
+    for rows in [1, 4, 16, 64, 256, 1024, 4096, 16384]:
+        repeats = 65536 // rows
+
+        current = 0
+        for i in range(repeats):
+            current += measure_create_dmatrix(rows=rows, cols=COLS, nthread=nthread, use_optimization=False)
+
+        optimized = 0
+        for i in range(repeats):
+            optimized += measure_create_dmatrix(rows=rows, cols=COLS, nthread=nthread, use_optimization=True)
+
+        print(f"{nthread:8} | {rows:8} | {COLS:8} | {current/repeats:15.4g} | {optimized/repeats:15.4g} | {optimized / current: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"}