Optimize device regex, test reusing a ProxyDMatrix. (dmlc#11273)

Author: trivialfis

demo/c-api/inference/inference.c

@@ -1,7 +1,10 @@
-/*!
- * Copyright 2021 XGBoost contributors
+/**
+ * Copyright 2021-2025, XGBoost contributors
  *
- * \brief A simple example of using prediction functions.
+ * @brief A simple example of using prediction functions.
+ *
+ * See more examples in test_c_api.cc on how to reuse a ProxyDMatrix object for reducing
+ * the latency of DMatrix creation.
  */
 #include <stddef.h>
 #include <stdlib.h>

src/context.cc

@@ -1,5 +1,5 @@
 /**
- * Copyright 2014-2024, XGBoost Contributors
+ * Copyright 2014-2025, XGBoost Contributors
  *
  * \brief Context object used for controlling runtime parameters.
  */
@@ -11,12 +11,17 @@
 #include <optional>   // for optional
 #include <regex>      // for regex_replace, regex_match
 
-#include "common/common.h"         // AssertGPUSupport
 #include "common/cuda_rt_utils.h"  // for AllVisibleGPUs
 #include "common/error_msg.h"      // WarnDeprecatedGPUId
 #include "common/threading_utils.h"
 #include "xgboost/string_view.h"
 
+#if !defined(XGBOOST_USE_CUDA)
+
+#include "common/common.h"  // for AssertGPUSupport
+
+#endif  // !defined(XGBOOST_USE_CUDA)
+
 namespace xgboost {
 
 DMLC_REGISTER_PARAMETER(Context);
@@ -108,7 +113,8 @@ DeviceOrd CUDAOrdinal(DeviceOrd device, bool) {
   bool valid = substr == "cpu" || substr == "cud" || substr == "gpu" || substr == "syc";
   CHECK(valid) << msg;
 #else
-  std::regex pattern{"gpu(:[0-9]+)?|cuda(:[0-9]+)?|cpu|sycl(:cpu|:gpu)?(:-1|:[0-9]+)?"};
+  thread_local static std::regex pattern{
+      "gpu(:[0-9]+)?|cuda(:[0-9]+)?|cpu|sycl(:cpu|:gpu)?(:-1|:[0-9]+)?"};
   if (!std::regex_match(input, pattern)) {
     fatal();
   }

tests/cpp/c_api/test_c_api.cc

@@ -1,5 +1,5 @@
 /**
- * Copyright 2019-2024 XGBoost contributors
+ * Copyright 2019-2025, XGBoost contributors
  */
 #include <gtest/gtest.h>
 #include <xgboost/c_api.h>
@@ -8,6 +8,7 @@
 #include <xgboost/learner.h>
 #include <xgboost/version_config.h>
 
+#include <algorithm>   // for copy_n
 #include <array>       // for array
 #include <cstddef>     // std::size_t
 #include <filesystem>  // std::filesystem
@@ -22,6 +23,7 @@
 #include "../../../src/data/batch_utils.h"          // for MatchingPageBytes
 #include "../../../src/data/gradient_index.h"       // for GHistIndexMatrix
 #include "../../../src/data/iterative_dmatrix.h"    // for IterativeDMatrix
+#include "../../../src/data/proxy_dmatrix.h"        // for DMatrixProxy
 #include "../../../src/data/sparse_page_dmatrix.h"  // for SparsePageDMatrix
 #include "../helpers.h"
 
@@ -599,4 +601,112 @@ TEST(CAPI, GPUXGDMatrixGetQuantileCut) {
   TestXGDMatrixGetQuantileCut(&ctx);
 }
 #endif  // defined(XGBOOST_USE_CUDA)
+
+TEST(CAPI, PredictReuseProxy) {
+  // Configuration for creating DMatrix
+  Json fmat_cfg{Object{}};
+  fmat_cfg["missing"] = std::numeric_limits<float>::quiet_NaN();
+  auto sfmat_cfg = Json::Dump(fmat_cfg);
+
+  // Configuration for prediction
+  Json config{Object{}};
+  config["type"] = Integer{0};
+  config["iteration_begin"] = config["iteration_end"] = Integer{0};
+  config["missing"] = Number{std::numeric_limits<float>::quiet_NaN()};
+  config["strict_shape"] = Boolean{true};
+  config["training"] = Boolean{false};
+  auto scfg = Json::Dump(config);
+
+  HostDeviceVector<float> storage;
+  bst_idx_t n_samples = 1024;
+  auto inf = RandomDataGenerator{n_samples, 256, 0.0}.GenerateArrayInterface(&storage);
+  HostDeviceVector<float> storage_y;
+  auto y_inf = RandomDataGenerator{n_samples, 1, 0.0}.GenerateArrayInterface(&storage_y);
+
+  // Create a DMatrix for training
+  DMatrixHandle fmat_hdl{nullptr};
+  ASSERT_EQ(XGDMatrixCreateFromDense(inf.c_str(), sfmat_cfg.c_str(), &fmat_hdl), 0);
+  ASSERT_EQ(XGDMatrixSetInfoFromInterface(fmat_hdl, "label", y_inf.c_str()), 0);
+
+  // Create booster and train.
+  std::array<DMatrixHandle, 1> mats{fmat_hdl};
+  BoosterHandle booster_hdl;
+  ASSERT_EQ(XGBoosterCreate(mats.data(), 1, &booster_hdl), 0);
+
+  for (std::int32_t i = 0; i < 3; ++i) {
+    ASSERT_EQ(XGBoosterUpdateOneIter(booster_hdl, i, fmat_hdl), 0);
+  }
+
+  // Create a proxy that can be reused.
+  DMatrixHandle proxy_hdl{nullptr};
+  ASSERT_EQ(XGProxyDMatrixCreate(&proxy_hdl), 0);
+
+  bst_ulong const *outshape{nullptr};
+  bst_ulong outdim{0};
+  float const *result{nullptr};
+
+  {
+    // Prediction with DMatrix
+    ASSERT_EQ(XGBoosterPredictFromDMatrix(booster_hdl, fmat_hdl, scfg.c_str(), &outshape, &outdim,
+                                          &result),
+              0);
+    bst_ulong n_samples_ret = 0;
+    ASSERT_EQ(XGDMatrixNumRow(fmat_hdl, &n_samples_ret), 0);
+    std::vector<float> vec_0(n_samples_ret);
+    ASSERT_EQ(vec_0.size(), n_samples);
+    ASSERT_EQ(outdim, 2);
+    std::copy_n(result, vec_0.size(), vec_0.begin());
+
+    // In-place predict
+    ASSERT_EQ(XGBoosterPredictFromDense(booster_hdl, inf.c_str(), scfg.c_str(), proxy_hdl,
+                                        &outshape, &outdim, &result),
+              0);
+    ASSERT_EQ(XGDMatrixNumRow(proxy_hdl, &n_samples_ret), 0);
+    std::vector<float> vec_1(n_samples_ret);
+    ASSERT_EQ(vec_1.size(), n_samples);
+    ASSERT_EQ(outdim, 2);
+    std::copy_n(result, vec_1.size(), vec_1.begin());
+
+    // Same result
+    ASSERT_EQ(vec_0, vec_1);
+  }
+
+  {
+    bst_idx_t n_samples = 512;
+
+    // Prediction with DMatrix
+    auto inf = RandomDataGenerator{n_samples, 256, 0.0}.GenerateArrayInterface(&storage);
+    DMatrixHandle fmat_hdl{nullptr};
+    ASSERT_EQ(XGDMatrixCreateFromDense(inf.c_str(), sfmat_cfg.c_str(), &fmat_hdl), 0);
+
+    ASSERT_EQ(XGBoosterPredictFromDMatrix(booster_hdl, fmat_hdl, scfg.c_str(), &outshape, &outdim,
+                                          &result),
+              0);
+    bst_ulong n_samples_ret = 0;
+    ASSERT_EQ(XGDMatrixNumRow(fmat_hdl, &n_samples_ret), 0);
+    std::vector<float> vec_0(n_samples_ret);
+    ASSERT_EQ(vec_0.size(), n_samples);
+    ASSERT_EQ(outdim, 2);
+    std::copy_n(result, vec_0.size(), vec_0.begin());
+
+    // In-place predict, same proxy as before
+    ASSERT_EQ(XGBoosterPredictFromDense(booster_hdl, inf.c_str(), scfg.c_str(), proxy_hdl,
+                                        &outshape, &outdim, &result),
+              0);
+    ASSERT_EQ(XGDMatrixNumRow(proxy_hdl, &n_samples_ret), 0);
+    std::vector<float> vec_1(n_samples_ret);
+    ASSERT_EQ(vec_1.size(), n_samples);
+    ASSERT_EQ(outdim, 2);
+    std::copy_n(result, vec_1.size(), vec_1.begin());
+
+    // Same result
+    ASSERT_EQ(vec_0, vec_1);
+
+    ASSERT_EQ(XGDMatrixFree(fmat_hdl), 0);
+  }
+
+  ASSERT_EQ(XGDMatrixFree(fmat_hdl), 0);
+  ASSERT_EQ(XGBoosterFree(booster_hdl), 0);
+  ASSERT_EQ(XGDMatrixFree(proxy_hdl), 0);
+}
 }  // namespace xgboost