Extend array interface to handle ndarray.

include/xgboost/data.h

@@ -11,12 +11,14 @@
 #include <dmlc/data.h>
 #include <dmlc/serializer.h>
 #include <xgboost/base.h>
-#include <xgboost/span.h>
 #include <xgboost/host_device_vector.h>
+#include <xgboost/linalg.h>
+#include <xgboost/span.h>
+#include <xgboost/string_view.h>
 
+#include <algorithm>
 #include <memory>
 #include <numeric>
-#include <algorithm>
 #include <string>
 #include <utility>
 #include <vector>
@@ -157,7 +159,7 @@ class MetaInfo {
    *
    *        Right now only 1 column is permitted.
    */
-  void SetInfo(const char* key, std::string const& interface_str);
+  void SetInfo(StringView key, std::string const& interface_str);
 
   void GetInfo(char const* key, bst_ulong* out_len, DataType dtype,
                const void** out_dptr) const;

jvm-packages/xgboost4j-gpu/src/native/xgboost4j-gpu.cu

@@ -35,13 +35,12 @@ template <typename T> T CheckJvmCall(T const &v, JNIEnv *jenv) {
 }
 
 template <typename VCont>
-void CopyColumnMask(xgboost::ArrayInterface const &interface,
+void CopyColumnMask(xgboost::ArrayInterface<1> const &interface,
                     std::vector<Json> const &columns, cudaMemcpyKind kind,
                     size_t c, VCont *p_mask, Json *p_out, cudaStream_t stream) {
   auto &mask = *p_mask;
   auto &out = *p_out;
-  auto size = sizeof(typename VCont::value_type) * interface.num_rows *
-              interface.num_cols;
+  auto size = sizeof(typename VCont::value_type) * interface.n;
   mask.resize(size);
   CHECK(RawPtr(mask));
   CHECK(size);
@@ -67,11 +66,11 @@ void CopyColumnMask(xgboost::ArrayInterface const &interface,
     LOG(FATAL) << "Invalid shape of mask";
   }
   out["mask"]["typestr"] = String("<t1");
-  out["mask"]["version"] = Integer(1);
+  out["mask"]["version"] = Integer(3);
 }
 
 template <typename DCont, typename VCont>
-void CopyInterface(std::vector<xgboost::ArrayInterface> &interface_arr,
+void CopyInterface(std::vector<xgboost::ArrayInterface<1>> &interface_arr,
                    std::vector<Json> const &columns, cudaMemcpyKind kind,
                    std::vector<DCont> *p_data, std::vector<VCont> *p_mask,
                    std::vector<xgboost::Json> *p_out, cudaStream_t stream) {
@@ -81,7 +80,7 @@ void CopyInterface(std::vector<xgboost::ArrayInterface> &interface_arr,
   for (size_t c = 0; c < interface_arr.size(); ++c) {
     auto &interface = interface_arr.at(c);
     size_t element_size = interface.ElementSize();
-    size_t size = element_size * interface.num_rows * interface.num_cols;
+    size_t size = element_size * interface.n;
 
     auto &data = (*p_data)[c];
     auto &mask = (*p_mask)[c];
@@ -95,25 +94,24 @@ void CopyInterface(std::vector<xgboost::ArrayInterface> &interface_arr,
         Json{Boolean{false}}};
 
     out["data"] = Array(std::move(j_data));
-    out["shape"] = Array(std::vector<Json>{Json(Integer(interface.num_rows)),
-                                           Json(Integer(interface.num_cols))});
+    out["shape"] = Array(std::vector<Json>{Json(Integer(interface.Shape(0)))});
 
     if (interface.valid.Data()) {
       CopyColumnMask(interface, columns, kind, c, &mask, &out, stream);
     }
     out["typestr"] = String("<f4");
-    out["version"] = Integer(1);
+    out["version"] = Integer(3);
   }
 }
 
 void CopyMetaInfo(Json *p_interface, dh::device_vector<float> *out, cudaStream_t stream) {
   auto &j_interface = *p_interface;
   CHECK_EQ(get<Array const>(j_interface).size(), 1);
   auto object = get<Object>(get<Array>(j_interface)[0]);
-  ArrayInterface interface(object);
-  out->resize(interface.num_rows);
+  ArrayInterface<1> interface(object);
+  out->resize(interface.Shape(0));
   size_t element_size = interface.ElementSize();
-  size_t size = element_size * interface.num_rows;
+  size_t size = element_size * interface.n;
   dh::safe_cuda(cudaMemcpyAsync(RawPtr(*out), interface.data, size,
                                 cudaMemcpyDeviceToDevice, stream));
   j_interface[0]["data"][0] = reinterpret_cast<Integer::Int>(RawPtr(*out));
@@ -285,11 +283,11 @@ class DataIteratorProxy {
 
     Json features = json_interface["features_str"];
     auto json_columns = get<Array const>(features);
-    std::vector<ArrayInterface> interfaces;
+    std::vector<ArrayInterface<1>> interfaces;
 
     // Stage the data
     for (auto &json_col : json_columns) {
-      auto column = ArrayInterface(get<Object const>(json_col));
+      auto column = ArrayInterface<1>(get<Object const>(json_col));
       interfaces.emplace_back(column);
     }
     Json::Dump(features, &interface_str);
@@ -342,9 +340,9 @@ class DataIteratorProxy {
     // Data
     auto const &json_interface = host_columns_.at(it_)->interfaces;
 
-    std::vector<ArrayInterface> in;
+    std::vector<ArrayInterface<1>> in;
     for (auto interface : json_interface) {
-      auto column = ArrayInterface(get<Object const>(interface));
+      auto column = ArrayInterface<1>(get<Object const>(interface));
       in.emplace_back(column);
     }
     std::vector<Json> out;

src/common/device_helpers.cuh

@@ -136,7 +136,7 @@ inline ncclResult_t ThrowOnNcclError(ncclResult_t code, const char *file,
 }
 #endif
 
-inline int32_t CudaGetPointerDevice(void* ptr) {
+inline int32_t CudaGetPointerDevice(void const *ptr) {
   int32_t device = -1;
   cudaPointerAttributes attr;
   dh::safe_cuda(cudaPointerGetAttributes(&attr, ptr));

src/data/adapter.h

@@ -254,20 +254,20 @@ class ArrayAdapterBatch : public detail::NoMetaInfo {
   static constexpr bool kIsRowMajor = true;
 
  private:
-  ArrayInterface array_interface_;
+  ArrayInterface<2> array_interface_;
 
   class Line {
-    ArrayInterface array_interface_;
+    ArrayInterface<2> array_interface_;
     size_t ridx_;
 
    public:
-    Line(ArrayInterface array_interface, size_t ridx)
+    Line(ArrayInterface<2> array_interface, size_t ridx)
         : array_interface_{std::move(array_interface)}, ridx_{ridx} {}
 
-    size_t Size() const { return array_interface_.num_cols; }
+    size_t Size() const { return array_interface_.Shape(1); }
 
     COOTuple GetElement(size_t idx) const {
-      return {ridx_, idx, array_interface_.GetElement(ridx_, idx)};
+      return {ridx_, idx, array_interface_(ridx_, idx)};
     }
   };
 
@@ -277,11 +277,11 @@ class ArrayAdapterBatch : public detail::NoMetaInfo {
     return Line{array_interface_, idx};
   }
 
-  size_t NumRows() const { return array_interface_.num_rows; }
-  size_t NumCols() const { return array_interface_.num_cols; }
+  size_t NumRows() const { return array_interface_.Shape(0); }
+  size_t NumCols() const { return array_interface_.Shape(1); }
   size_t Size() const { return this->NumRows(); }
 
-  explicit ArrayAdapterBatch(ArrayInterface array_interface)
+  explicit ArrayAdapterBatch(ArrayInterface<2> array_interface)
       : array_interface_{std::move(array_interface)} {}
 };
 
@@ -294,43 +294,42 @@ class ArrayAdapter : public detail::SingleBatchDataIter<ArrayAdapterBatch> {
  public:
   explicit ArrayAdapter(StringView array_interface) {
     auto j = Json::Load(array_interface);
-    array_interface_ = ArrayInterface(get<Object const>(j));
+    array_interface_ = ArrayInterface<2>(get<Object const>(j));
     batch_ = ArrayAdapterBatch{array_interface_};
   }
   ArrayAdapterBatch const& Value() const override { return batch_; }
-  size_t NumRows() const { return array_interface_.num_rows; }
-  size_t NumColumns() const { return array_interface_.num_cols; }
+  size_t NumRows() const { return array_interface_.Shape(0); }
+  size_t NumColumns() const { return array_interface_.Shape(1); }
 
  private:
   ArrayAdapterBatch batch_;
-  ArrayInterface array_interface_;
+  ArrayInterface<2> array_interface_;
 };
 
 class CSRArrayAdapterBatch : public detail::NoMetaInfo {
-  ArrayInterface indptr_;
-  ArrayInterface indices_;
-  ArrayInterface values_;
+  ArrayInterface<1> indptr_;
+  ArrayInterface<1> indices_;
+  ArrayInterface<1> values_;
   bst_feature_t n_features_;
 
   class Line {
-    ArrayInterface indices_;
-    ArrayInterface values_;
+    ArrayInterface<1> indices_;
+    ArrayInterface<1> values_;
     size_t ridx_;
     size_t offset_;
 
    public:
-    Line(ArrayInterface indices, ArrayInterface values, size_t ridx,
+    Line(ArrayInterface<1> indices, ArrayInterface<1> values, size_t ridx,
          size_t offset)
         : indices_{std::move(indices)}, values_{std::move(values)}, ridx_{ridx},
           offset_{offset} {}
 
     COOTuple GetElement(size_t idx) const {
-      return {ridx_, indices_.GetElement<size_t>(offset_ + idx, 0),
-              values_.GetElement(offset_ + idx, 0)};
+      return {ridx_, TypedIndex<size_t, 1>{indices_}(offset_ + idx), values_(offset_ + idx)};
     }
 
     size_t Size() const {
-      return values_.num_rows * values_.num_cols;
+      return values_.Shape(0);
     }
   };
 
@@ -339,37 +338,36 @@ class CSRArrayAdapterBatch : public detail::NoMetaInfo {
 
  public:
   CSRArrayAdapterBatch() = default;
-  CSRArrayAdapterBatch(ArrayInterface indptr, ArrayInterface indices,
-                       ArrayInterface values, bst_feature_t n_features)
-      : indptr_{std::move(indptr)}, indices_{std::move(indices)},
-        values_{std::move(values)}, n_features_{n_features} {
-    indptr_.AsColumnVector();
-    values_.AsColumnVector();
-    indices_.AsColumnVector();
+  CSRArrayAdapterBatch(ArrayInterface<1> indptr, ArrayInterface<1> indices,
+                       ArrayInterface<1> values, bst_feature_t n_features)
+      : indptr_{std::move(indptr)},
+        indices_{std::move(indices)},
+        values_{std::move(values)},
+        n_features_{n_features} {
   }
 
   size_t NumRows() const {
-    size_t size = indptr_.num_rows * indptr_.num_cols;
+    size_t size = indptr_.Shape(0);
     size = size == 0 ? 0 : size - 1;
     return size;
   }
   size_t NumCols() const { return n_features_; }
   size_t Size() const { return this->NumRows(); }
 
   Line const GetLine(size_t idx) const {
-    auto begin_offset = indptr_.GetElement<size_t>(idx, 0);
-    auto end_offset = indptr_.GetElement<size_t>(idx + 1, 0);
+    auto begin_no_stride = TypedIndex<size_t, 1>{indptr_}(idx);
+    auto end_no_stride = TypedIndex<size_t, 1>{indptr_}(idx + 1);
 
     auto indices = indices_;
     auto values = values_;
+    // Slice indices and values, stride remains unchanged since this is slicing by
+    // specific index.
+    auto offset = indices.strides[0] * begin_no_stride;
 
-    values.num_cols = end_offset - begin_offset;
-    values.num_rows = 1;
+    indices.shape[0] = end_no_stride - begin_no_stride;
+    values.shape[0] = end_no_stride - begin_no_stride;
 
-    indices.num_cols = values.num_cols;
-    indices.num_rows = values.num_rows;
-
-    return Line{indices, values, idx, begin_offset};
+    return Line{indices, values, idx, offset};
   }
 };
 
@@ -391,17 +389,17 @@ class CSRArrayAdapter : public detail::SingleBatchDataIter<CSRArrayAdapterBatch>
     return batch_;
   }
   size_t NumRows() const {
-    size_t size = indptr_.num_cols * indptr_.num_rows;
+    size_t size = indptr_.Shape(0);
     size = size == 0 ? 0 : size - 1;
     return  size;
   }
   size_t NumColumns() const { return num_cols_; }
 
  private:
   CSRArrayAdapterBatch batch_;
-  ArrayInterface indptr_;
-  ArrayInterface indices_;
-  ArrayInterface values_;
+  ArrayInterface<1> indptr_;
+  ArrayInterface<1> indices_;
+  ArrayInterface<1> values_;
   size_t num_cols_;
 };
 

src/data/array_interface.cu

@@ -1,21 +1,56 @@
 /*!
  * Copyright 2021 by Contributors
  */
-#include "array_interface.h"
 #include "../common/common.h"
+#include "array_interface.h"
 
 namespace xgboost {
 void ArrayInterfaceHandler::SyncCudaStream(int64_t stream) {
   switch (stream) {
-  case 0:
-    LOG(FATAL) << "Invalid stream ID in array interface: " << stream;
-  case 1:
-    // default legacy stream
-    break;
-  case 2:
-    // default per-thread stream
-  default:
-    dh::safe_cuda(cudaStreamSynchronize(reinterpret_cast<cudaStream_t>(stream)));
+    case 0:
+      /**
+       * disallowed by the `__cuda_array_interface__`.  Quote:
+       *
+       *   This is disallowed as it would be ambiguous between None and the default
+       *   stream, and also between the legacy and per-thread default streams. Any use
+       *   case where 0 might be given should either use None, 1, or 2 instead for
+       *   clarity.
+       */
+      LOG(FATAL) << "Invalid stream ID in array interface: " << stream;
+    case 1:
+      // default legacy stream
+      break;
+    case 2:
+      // default per-thread stream
+    default:
+      dh::safe_cuda(cudaStreamSynchronize(reinterpret_cast<cudaStream_t>(stream)));
+  }
+}
+
+bool ArrayInterfaceHandler::IsCudaPtr(void const* ptr) {
+  if (!ptr) {
+    return false;
+  }
+  cudaPointerAttributes attr;
+  auto err = cudaPointerGetAttributes(&attr, ptr);
+  // reset error
+  CHECK_EQ(err, cudaGetLastError());
+  if (err == cudaErrorInvalidValue) {
+    // CUDA < 11
+    return false;
+  } else if (err == cudaSuccess) {
+    // CUDA >= 11
+    switch (attr.type) {
+      case cudaMemoryTypeUnregistered:
+      case cudaMemoryTypeHost:
+        return false;
+      default:
+        return true;
+    }
+    return true;
+  } else {
+    // other errors, `cudaErrorNoDevice`, `cudaErrorInsufficientDriver` etc.
+    return false;
   }
 }
 }  // namespace xgboost