workingset.h

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2019-2024, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include <cuml/common/logger.hpp>
#include <cuml/svm/svm_parameter.h>
 
#include <raft/core/handle.hpp>
 
#include <rmm/device_scalar.hpp>
#include <rmm/device_uvector.hpp>
 
#include <algorithm>
#include <cstddef>
#include <limits>
 
namespace ML {
namespace SVM {
 
template <typename math_t>
class WorkingSet {
 public:
  bool FIFO_strategy = true;
 
  WorkingSet(const raft::handle_t& handle,
             cudaStream_t stream,
             int n_rows      = 0,
             int n_ws        = 0,
             SvmType svmType = C_SVC)
    : handle(handle),
      stream(stream),
      svmType(svmType),
      n_rows(n_rows),
      available(0, stream),
      available_sorted(0, stream),
      cub_storage(0, stream),
      f_idx(0, stream),
      f_idx_sorted(0, stream),
      f_sorted(0, stream),
      idx_tmp(0, stream),
      idx(0, stream),
      ws_idx_sorted(0, stream),
      ws_idx_selected(0, stream),
      ws_idx_save(0, stream),
      ws_priority(0, stream),
      ws_priority_sorted(0, stream),
      d_num_selected(stream)
  {
    n_train = (svmType == EPSILON_SVR) ? n_rows * 2 : n_rows;
    SetSize(n_train, n_ws);
  }
 
  ~WorkingSet() {}
 
  void SetSize(int n_train, int n_ws = 0)
  {
    if (n_ws == 0 || n_ws > n_train) { n_ws = n_train; }
    n_ws       = std::min(1024, n_ws);
    this->n_ws = n_ws;
    CUML_LOG_DEBUG("Creating working set with %d elements", n_ws);
    AllocateBuffers();
  }
 
  int GetSize() { return n_ws; }
 
  int* GetIndices() { return idx.data(); }
 
  void SimpleSelect(
    math_t* f, math_t* alpha, math_t* y, const math_t* C, int n_already_selected = 0);
 
  void Select(math_t* f, math_t* alpha, math_t* y, const math_t* C)
  {
    if (n_ws >= n_train) {
      // All elements are selected, we have initialized idx to cover this case
      return;
    }
    int nc         = n_ws / 4;
    int n_selected = 0;
    if (firstcall) {
      if (nc >= 1) {
        firstcall = false;
      } else {
        // This can only happen for n_ws < 4.
        // We keep the calculation always in firstcall mode (only SimpleSelect
        // is used, no advanced strategies because we do not have enough elements)
        //
        // Nothing to do, firstcall is already true
      }
    } else {
      // keep 1/2 of the old working set
      if (FIFO_strategy) {
        // FIFO selection following ThunderSVM
        raft::copy(idx.data(), ws_idx_save.data() + 2 * nc, 2 * nc, stream);
        n_selected = nc * 2;
      } else {
        // priority based selection preferring to keep newer elements in ws
        n_selected = PrioritySelect(alpha, C, nc);
      }
    }
    SimpleSelect(f, alpha, y, C, n_selected);
    raft::copy(ws_idx_save.data(), idx.data(), n_ws, stream);
  }
 
  int PrioritySelect(math_t* alpha, const math_t* C, int nc);
 
 private:
  const raft::handle_t& handle;
  cudaStream_t stream;
 
  bool firstcall = true;
  int n_train    = 0;  
  int n_rows     = 0;  
  int n_ws       = 0;
 
  SvmType svmType;
 
  int TPB = 256;  
 
  // Buffers for the domain size [n_train]
  rmm::device_uvector<int> f_idx;  
  rmm::device_uvector<int> f_idx_sorted;
  rmm::device_uvector<int> idx_tmp;
  rmm::device_uvector<math_t> f_sorted;
  rmm::device_uvector<bool> available;
  rmm::device_uvector<bool> available_sorted;
 
  // working set buffers size [n_ws]
  rmm::device_uvector<int> idx;  
  rmm::device_uvector<int> ws_idx_sorted;
  rmm::device_uvector<int> ws_idx_selected;
  rmm::device_uvector<int> ws_idx_save;
 
  rmm::device_uvector<int> ws_priority;
  rmm::device_uvector<int> ws_priority_sorted;
 
  rmm::device_scalar<int> d_num_selected;
  std::size_t cub_bytes = 0;
  rmm::device_uvector<char> cub_storage;
 
  void AllocateBuffers();
 
  int GatherAvailable(int n_already_selected, int n_needed, bool copy_front);
 
  void Initialize();
 
  template <typename select_op>
  int SelectPrevWs(int n_needed, int n_already_selected, select_op op);
};
 
};  // end namespace SVM
};  // end namespace ML