Implement unified update prediction cache for (gpu_)hist. (#6860)

* Implement utilites for linalg.
* Unify the update prediction cache functions.
* Implement update prediction cache for multi-class gpu hist.

include/xgboost/linalg.h

@@ -0,0 +1,105 @@
+/*!
+ * Copyright 2021 by Contributors
+ * \file linalg.h
+ * \brief  Linear algebra related utilities.
+ */
+#ifndef XGBOOST_LINALG_H_
+#define XGBOOST_LINALG_H_
+
+#include <xgboost/span.h>
+#include <xgboost/host_device_vector.h>
+#include <xgboost/generic_parameters.h>
+
+#include <array>
+#include <algorithm>
+#include <utility>
+
+namespace xgboost {
+/*!
+ * \brief A veiw over a matrix on contigious storage.
+ *
+ * \tparam T data type of matrix
+ */
+template <typename T> class MatrixView {
+  int32_t device_;
+  common::Span<T> values_;
+  size_t strides_[2];
+  size_t shape_[2];
+
+  template <typename Vec> static auto InferValues(Vec *vec, int32_t device) {
+    return device == GenericParameter::kCpuId ? vec->HostSpan()
+                                              : vec->DeviceSpan();
+  }
+
+ public:
+  /*!
+   * \param vec     storage.
+   * \param strides Strides for matrix.
+   * \param shape   Rows anc columns.
+   * \param device  Where the data is stored in.
+   */
+  MatrixView(HostDeviceVector<T> *vec, std::array<size_t, 2> strides,
+             std::array<size_t, 2> shape, int32_t device)
+      : device_{device}, values_{InferValues(vec, device)} {
+    std::copy(strides.cbegin(), strides.cend(), strides_);
+    std::copy(shape.cbegin(), shape.cend(), shape_);
+  }
+  MatrixView(HostDeviceVector<std::remove_const_t<T>> const *vec,
+             std::array<size_t, 2> strides, std::array<size_t, 2> shape,
+             int32_t device)
+      : device_{device}, values_{InferValues(vec, device)} {
+    std::copy(strides.cbegin(), strides.cend(), strides_);
+    std::copy(shape.cbegin(), shape.cend(), shape_);
+  }
+  /*! \brief Row major constructor. */
+  MatrixView(HostDeviceVector<T> *vec, std::array<size_t, 2> shape,
+             int32_t device)
+      : device_{device}, values_{InferValues(vec, device)} {
+    std::copy(shape.cbegin(), shape.cend(), shape_);
+    strides_[0] = shape[1];
+    strides_[1] = 1;
+  }
+  MatrixView(HostDeviceVector<std::remove_const_t<T>> const *vec,
+             std::array<size_t, 2> shape, int32_t device)
+      : device_{device}, values_{InferValues(vec, device)} {
+    std::copy(shape.cbegin(), shape.cend(), shape_);
+    strides_[0] = shape[1];
+    strides_[1] = 1;
+  }
+
+  XGBOOST_DEVICE T const &operator()(size_t r, size_t c) const {
+    return values_[strides_[0] * r + strides_[1] * c];
+  }
+  XGBOOST_DEVICE T &operator()(size_t r, size_t c) {
+    return values_[strides_[0] * r + strides_[1] * c];
+  }
+
+  auto Strides() const { return strides_; }
+  auto Shape() const { return shape_; }
+  auto Values() const { return values_; }
+  auto Size() const { return shape_[0] * shape_[1]; }
+  auto DeviceIdx() const { return device_; }
+};
+
+/*! \brief A slice for 1 column of MatrixView.  Can be extended to row if needed. */
+template <typename T> class VectorView {
+  MatrixView<T> matrix_;
+  size_t column_;
+
+ public:
+  explicit VectorView(MatrixView<T> matrix, size_t column)
+      : matrix_{std::move(matrix)}, column_{column} {}
+
+  XGBOOST_DEVICE T &operator[](size_t i) {
+    return matrix_(i, column_);
+  }
+
+  XGBOOST_DEVICE T const &operator[](size_t i) const {
+    return matrix_(i, column_);
+  }
+
+  size_t Size() { return matrix_.Shape()[0]; }
+  int32_t DeviceIdx() const { return matrix_.DeviceIdx(); }
+};
+}       // namespace xgboost
+#endif  // XGBOOST_LINALG_H_

include/xgboost/tree_updater.h

@@ -15,6 +15,7 @@
 #include <xgboost/generic_parameters.h>
 #include <xgboost/host_device_vector.h>
 #include <xgboost/model.h>
+#include <xgboost/linalg.h>
 
 #include <functional>
 #include <vector>
@@ -70,14 +71,8 @@ class TreeUpdater : public Configurable {
    *         the prediction cache. If true, the prediction cache will have been
    *         updated by the time this function returns.
    */
-  virtual bool UpdatePredictionCache(const DMatrix* /*data*/,
-                                     HostDeviceVector<bst_float>* /*out_preds*/) {
-    return false;
-  }
-
-  virtual bool UpdatePredictionCacheMulticlass(const DMatrix* /*data*/,
-                                               HostDeviceVector<bst_float>* /*out_preds*/,
-                                               const int /*gid*/, const int /*ngroup*/) {
+  virtual bool UpdatePredictionCache(const DMatrix * /*data*/,
+                                     VectorView<float> /*out_preds*/) {
     return false;
   }