Use matrix for gradient. (#9508)

- Use the `linalg::Matrix` for storing gradients.
- New API for the custom objective.
- Custom objective for multi-class/multi-target is now required to return the correct shape.
- Custom objective for Python can accept arrays with any strides. (row-major, column-major)

R-package/R/utils.R

@@ -154,7 +154,14 @@ xgb.iter.update <- function(booster_handle, dtrain, iter, obj) {
     pred <- predict(booster_handle, dtrain, outputmargin = TRUE, training = TRUE,
                     ntreelimit = 0)
     gpair <- obj(pred, dtrain)
-    .Call(XGBoosterBoostOneIter_R, booster_handle, dtrain, gpair$grad, gpair$hess)
+    n_samples <- dim(dtrain)[1]
+    # We still require row-major in R as I'm not quite sure sure how to get the stride of
+    # the matrix in C.
+    gpair$grad <- matrix(gpair$grad, nrow = n_samples, byrow = TRUE)
+    gpair$hess <- matrix(gpair$hess, nrow = n_samples, byrow = TRUE)
+    .Call(
+      XGBoosterBoostOneIter_R, booster_handle, dtrain, iter, gpair$grad, gpair$hess
+    )
   }
   return(TRUE)
 }

R-package/src/init.c

@@ -16,7 +16,7 @@ Check these declarations against the C/Fortran source code.
 */
 
 /* .Call calls */
-extern SEXP XGBoosterBoostOneIter_R(SEXP, SEXP, SEXP, SEXP);
+extern SEXP XGBoosterTrainOneIter_R(SEXP, SEXP, SEXP, SEXP, SEXP);
 extern SEXP XGBoosterCreate_R(SEXP);
 extern SEXP XGBoosterCreateInEmptyObj_R(SEXP, SEXP);
 extern SEXP XGBoosterDumpModel_R(SEXP, SEXP, SEXP, SEXP);
@@ -53,7 +53,7 @@ extern SEXP XGBGetGlobalConfig_R(void);
 extern SEXP XGBoosterFeatureScore_R(SEXP, SEXP);
 
 static const R_CallMethodDef CallEntries[] = {
-  {"XGBoosterBoostOneIter_R",     (DL_FUNC) &XGBoosterBoostOneIter_R,     4},
+  {"XGBoosterBoostOneIter_R",     (DL_FUNC) &XGBoosterTrainOneIter_R,     5},
   {"XGBoosterCreate_R",           (DL_FUNC) &XGBoosterCreate_R,           1},
   {"XGBoosterCreateInEmptyObj_R", (DL_FUNC) &XGBoosterCreateInEmptyObj_R, 2},
   {"XGBoosterDumpModel_R",        (DL_FUNC) &XGBoosterDumpModel_R,        4},

R-package/src/xgboost_R.cc

@@ -48,13 +48,6 @@
 
 using dmlc::BeginPtr;
 
-xgboost::Context const *BoosterCtx(BoosterHandle handle) {
-  CHECK_HANDLE();
-  auto *learner = static_cast<xgboost::Learner *>(handle);
-  CHECK(learner);
-  return learner->Ctx();
-}
-
 xgboost::Context const *DMatrixCtx(DMatrixHandle handle) {
   CHECK_HANDLE();
   auto p_m = static_cast<std::shared_ptr<xgboost::DMatrix> *>(handle);
@@ -394,21 +387,25 @@ XGB_DLL SEXP XGBoosterUpdateOneIter_R(SEXP handle, SEXP iter, SEXP dtrain) {
   return R_NilValue;
 }
 
-XGB_DLL SEXP XGBoosterBoostOneIter_R(SEXP handle, SEXP dtrain, SEXP grad, SEXP hess) {
+XGB_DLL SEXP XGBoosterTrainOneIter_R(SEXP handle, SEXP dtrain, SEXP iter, SEXP grad, SEXP hess) {
   R_API_BEGIN();
-  CHECK_EQ(length(grad), length(hess))
-      << "gradient and hess must have same length";
-  int len = length(grad);
-  std::vector<float> tgrad(len), thess(len);
-  auto ctx = BoosterCtx(R_ExternalPtrAddr(handle));
-  xgboost::common::ParallelFor(len, ctx->Threads(), [&](xgboost::omp_ulong j) {
-    tgrad[j] = REAL(grad)[j];
-    thess[j] = REAL(hess)[j];
-  });
-  CHECK_CALL(XGBoosterBoostOneIter(R_ExternalPtrAddr(handle),
-                                 R_ExternalPtrAddr(dtrain),
-                                 BeginPtr(tgrad), BeginPtr(thess),
-                                 len));
+  CHECK_EQ(length(grad), length(hess)) << "gradient and hess must have same length";
+  SEXP gdim = getAttrib(grad, R_DimSymbol);
+  auto n_samples = static_cast<std::size_t>(INTEGER(gdim)[0]);
+  auto n_targets = static_cast<std::size_t>(INTEGER(gdim)[1]);
+
+  SEXP hdim = getAttrib(hess, R_DimSymbol);
+  CHECK_EQ(INTEGER(hdim)[0], n_samples) << "mismatched size between gradient and hessian";
+  CHECK_EQ(INTEGER(hdim)[1], n_targets) << "mismatched size between gradient and hessian";
+  double const *d_grad = REAL(grad);
+  double const *d_hess = REAL(hess);
+
+  auto ctx = xgboost::detail::BoosterCtx(R_ExternalPtrAddr(handle));
+  auto [s_grad, s_hess] =
+      xgboost::detail::MakeGradientInterface(ctx, d_grad, d_hess, n_samples, n_targets);
+  CHECK_CALL(XGBoosterTrainOneIter(R_ExternalPtrAddr(handle), R_ExternalPtrAddr(dtrain),
+                                   asInteger(iter), s_grad.c_str(), s_hess.c_str()));
+
   R_API_END();
   return R_NilValue;
 }
@@ -460,7 +457,7 @@ XGB_DLL SEXP XGBoosterPredictFromDMatrix_R(SEXP handle, SEXP dmat, SEXP json_con
     len *= out_shape[i];
   }
   r_out_result = PROTECT(allocVector(REALSXP, len));
-  auto ctx = BoosterCtx(R_ExternalPtrAddr(handle));
+  auto ctx = xgboost::detail::BoosterCtx(R_ExternalPtrAddr(handle));
   xgboost::common::ParallelFor(len, ctx->Threads(), [&](xgboost::omp_ulong i) {
     REAL(r_out_result)[i] = out_result[i];
   });
@@ -669,7 +666,7 @@ XGB_DLL SEXP XGBoosterFeatureScore_R(SEXP handle, SEXP json_config) {
   }
 
   out_scores_sexp = PROTECT(allocVector(REALSXP, len));
-  auto ctx = BoosterCtx(R_ExternalPtrAddr(handle));
+  auto ctx = xgboost::detail::BoosterCtx(R_ExternalPtrAddr(handle));
   xgboost::common::ParallelFor(len, ctx->Threads(), [&](xgboost::omp_ulong i) {
     REAL(out_scores_sexp)[i] = out_scores[i];
   });

R-package/src/xgboost_R.h

@@ -161,12 +161,13 @@ XGB_DLL SEXP XGBoosterUpdateOneIter_R(SEXP ext, SEXP iter, SEXP dtrain);
  * \brief update the model, by directly specify gradient and second order gradient,
  *        this can be used to replace UpdateOneIter, to support customized loss function
  * \param handle handle
+ * \param iter The current training iteration.
  * \param dtrain training data
  * \param grad gradient statistics
  * \param hess second order gradient statistics
  * \return R_NilValue
  */
-XGB_DLL SEXP XGBoosterBoostOneIter_R(SEXP handle, SEXP dtrain, SEXP grad, SEXP hess);
+XGB_DLL SEXP XGBoosterTrainOneIter_R(SEXP handle, SEXP dtrain, SEXP iter, SEXP grad, SEXP hess);
 
 /*!
  * \brief get evaluation statistics for xgboost