Small cleanup to gradient index and hist. (#7668)

* Code comments.
* Const accessor to index.
* Remove some weird variables in the `Index` class.
* Simplify the `MemStackAllocator`.

src/common/hist_util.h

@@ -197,19 +197,27 @@ enum BinTypeSize : uint32_t {
   kUint32BinsTypeSize = 4
 };
 
+/**
+ * \brief Optionally compressed gradient index. The compression works only with dense
+ *        data.
+ *
+ *   The main body of construction code is in gradient_index.cc, this struct is only a
+ *   storage class.
+ */
 struct Index {
-  Index() {
-    SetBinTypeSize(binTypeSize_);
-  }
+  Index() { SetBinTypeSize(binTypeSize_); }
   Index(const Index& i) = delete;
   Index& operator=(Index i) = delete;
   Index(Index&& i) = delete;
   Index& operator=(Index&& i) = delete;
   uint32_t operator[](size_t i) const {
-    if (offset_ptr_ != nullptr) {
-      return func_(data_ptr_, i) + offset_ptr_[i%p_];
+    if (!bin_offset_.empty()) {
+      // dense, compressed
+      auto fidx = i % bin_offset_.size();
+      // restore the index by adding back its feature offset.
+      return func_(data_.data(), i) + bin_offset_[fidx];
     } else {
-      return func_(data_ptr_, i);
+      return func_(data_.data(), i);
     }
   }
   void SetBinTypeSize(BinTypeSize binTypeSize) {
@@ -225,35 +233,32 @@ struct Index {
         func_ = &GetValueFromUint32;
         break;
       default:
-        CHECK(binTypeSize == kUint8BinsTypeSize  ||
-              binTypeSize == kUint16BinsTypeSize ||
+        CHECK(binTypeSize == kUint8BinsTypeSize || binTypeSize == kUint16BinsTypeSize ||
               binTypeSize == kUint32BinsTypeSize);
     }
   }
   BinTypeSize GetBinTypeSize() const {
     return binTypeSize_;
   }
-  template<typename T>
-  T* data() const {  // NOLINT
-    return static_cast<T*>(data_ptr_);
+  template <typename T>
+  T const* data() const {  // NOLINT
+    return reinterpret_cast<T const*>(data_.data());
   }
-  uint32_t* Offset() const {
-    return offset_ptr_;
+  template <typename T>
+  T* data() {  // NOLINT
+    return reinterpret_cast<T*>(data_.data());
   }
-  size_t OffsetSize() const {
-    return offset_.size();
+  uint32_t const* Offset() const { return bin_offset_.data(); }
+  size_t OffsetSize() const { return bin_offset_.size(); }
+  size_t Size() const { return data_.size() / (binTypeSize_); }
+
+  void Resize(const size_t n_bytes) {
+    data_.resize(n_bytes);
   }
-  size_t Size() const {
-    return data_.size() / (binTypeSize_);
-  }
-  void Resize(const size_t nBytesData) {
-    data_.resize(nBytesData);
-    data_ptr_ = reinterpret_cast<void*>(data_.data());
-  }
-  void ResizeOffset(const size_t nDisps) {
-    offset_.resize(nDisps);
-    offset_ptr_ = offset_.data();
-    p_ = nDisps;
+  // set the offset used in compression, cut_ptrs is the CSC indptr in HistogramCuts
+  void SetBinOffset(std::vector<uint32_t> const& cut_ptrs) {
+    bin_offset_.resize(cut_ptrs.size() - 1);  // resize to number of features.
+    std::copy_n(cut_ptrs.begin(), bin_offset_.size(), bin_offset_.begin());
   }
   std::vector<uint8_t>::const_iterator begin() const {  // NOLINT
     return data_.begin();
@@ -270,24 +275,23 @@ struct Index {
   }
 
  private:
-  static uint32_t GetValueFromUint8(void *t, size_t i) {
-    return reinterpret_cast<uint8_t*>(t)[i];
+  // Functions to decompress the index.
+  static uint32_t GetValueFromUint8(uint8_t const* t, size_t i) { return t[i]; }
+  static uint32_t GetValueFromUint16(uint8_t const* t, size_t i) {
+    return reinterpret_cast<uint16_t const*>(t)[i];
   }
-  static uint32_t GetValueFromUint16(void* t, size_t i) {
-    return reinterpret_cast<uint16_t*>(t)[i];
-  }
-  static uint32_t GetValueFromUint32(void* t, size_t i) {
-    return reinterpret_cast<uint32_t*>(t)[i];
+  static uint32_t GetValueFromUint32(uint8_t const* t, size_t i) {
+    return reinterpret_cast<uint32_t const*>(t)[i];
   }
 
-  using Func = uint32_t (*)(void*, size_t);
+  using Func = uint32_t (*)(uint8_t const*, size_t);
 
   std::vector<uint8_t> data_;
-  std::vector<uint32_t> offset_;  // size of this field is equal to number of features
-  void* data_ptr_;
+  // starting position of each feature inside the cut values (the indptr of the CSC cut matrix
+  // HistogramCuts without the last entry.) Used for bin compression.
+  std::vector<uint32_t> bin_offset_;
+
   BinTypeSize binTypeSize_ {kUint8BinsTypeSize};
-  size_t p_ {1};
-  uint32_t* offset_ptr_ {nullptr};
   Func func_;
 };
 
@@ -304,9 +308,11 @@ int32_t XGBOOST_HOST_DEV_INLINE BinarySearchBin(size_t begin, size_t end,
     }
     previous_middle = middle;
 
+    // index into all the bins
     auto gidx = data[middle];
 
     if (gidx >= fidx_begin && gidx < fidx_end) {
+      // Found the intersection.
       return static_cast<int32_t>(gidx);
     } else if (gidx < fidx_begin) {
       begin = middle;
@@ -636,42 +642,6 @@ class GHistBuilder {
   /*! \brief number of all bins over all features */
   uint32_t nbins_ { 0 };
 };
-
-/*!
- * \brief A C-style array with in-stack allocation. As long as the array is smaller than
- * MaxStackSize, it will be allocated inside the stack. Otherwise, it will be
- * heap-allocated.
- */
-template<typename T, size_t MaxStackSize>
-class MemStackAllocator {
- public:
-  explicit MemStackAllocator(size_t required_size): required_size_(required_size) {
-  }
-
-  T* Get() {
-    if (!ptr_) {
-      if (MaxStackSize >= required_size_) {
-        ptr_ = stack_mem_;
-      } else {
-        ptr_ =  reinterpret_cast<T*>(malloc(required_size_ * sizeof(T)));
-        do_free_ = true;
-      }
-    }
-
-    return ptr_;
-  }
-
-  ~MemStackAllocator() {
-    if (do_free_) free(ptr_);
-  }
-
-
- private:
-  T* ptr_ = nullptr;
-  bool do_free_ = false;
-  size_t required_size_;
-  T stack_mem_[MaxStackSize];
-};
 }  // namespace common
 }  // namespace xgboost
 #endif  // XGBOOST_COMMON_HIST_UTIL_H_