text: /**
   * (C) Copyright 2020, 2021 IBM. All Rights Reserved.
   *
   * This code is licensed under the Apache License, Version 2.0. You may
   * obtain a copy of this license in the LICENSE.txt file in the root directory
   * of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
   *
   * Any modifications or derivative works of this code must retain this
   * copyright notice, and modified files need to carry a notice indicating
   * that they have been altered from the originals.
   */
  
  #include "rpu_weight_updater.h"
  #include "rpu_vector_device.h"
  #include "utility_functions.h"
  
  namespace RPU {
  
  /*FP update */
  template <typename T>
  void RPUWeightUpdater<T>::updateVector(
      T **weights,
      const T *x_input,
      const int x_inc,
      const T *d_input,
      const int d_inc,
      const T learning_rate) {
  
    RPU::math::ger<T>(
        CblasRowMajor, this->d_size_, this->x_size_, -learning_rate, d_input, d_inc, x_input, x_inc,
        weights[0], this->x_size_);
  }
  
  template class RPUWeightUpdater<float>;
  #ifdef RPU_USE_DOUBLE
  template class RPUWeightUpdater<double>;
  #endif
  
  /*********************************************************************/
  /* Pulsed update */
  
  template <typename T> void PulsedRPUWeightUpdater<T>::allocateContainers() {
  
    if (!containers_allocated_) {
      freeContainers();
  
      sblm_ = std::unique_ptr<SparseBitLineMaker<T>>(
          new SparseBitLineMaker<T>(this->x_size_, this->d_size_));
      dblm_ = std::unique_ptr<DenseBitLineMaker<T>>(
          new DenseBitLineMaker<T>(this->x_size_, this->d_size_));
  
      containers_allocated_ = true;
    }
  }
  
  template <typename T> void PulsedRPUWeightUpdater<T>::freeContainers() {
  
    if (containers_allocated_) {
  
      sblm_ = nullptr;
      dblm_ = nullptr;
  
      containers_allocated_ = false;
    }
  }
  
  // ctor
  template <typename T>
  PulsedRPUWeightUpdater<T>::PulsedRPUWeightUpdater(
      int x_size, int d_size, std::shared_ptr<RNG<T>> rng)
      : RPUWeightUpdater<T>(x_size, d_size), rng_(rng) {
    allocateContainers();
  }
  
  // dtor
  template <typename T> PulsedRPUWeightUpdater<T>::~PulsedRPUWeightUpdater() { freeContainers(); }
  
  // copy construcutor
  template <typename T>
  PulsedRPUWeightUpdater<T>::PulsedRPUWeightUpdater(const PulsedRPUWeightUpdater<T> &other)
      : RPUWeightUpdater<T>(other) {
    up_ = other.up_;
    rng_ = other.rng_;
  
    if (other.containers_allocated_) {
      allocateContainers();
      *sblm_ = *other.sblm_;
      *dblm_ = *other.dblm_;
    }
  }
  
  // copy assignment
  template <typename T>
  PulsedRPUWeightUpdater<T> &
  PulsedRPUWeightUpdater<T>::operator=(const PulsedRPUWeightUpdater<T> &other) {
  
    PulsedRPUWeightUpdater<T> tmp(other);
    swap(*this, tmp);
    return *this;
  }
  
  // move constructor
  template <typename T>
  PulsedRPUWeightUpdater<T>::PulsedRPUWeightUpdater(PulsedRPUWeightUpdater<T> &&other) {
  
    *this = std::move(other);
  }
  
  // move assignment
  template <typename T>
  PulsedRPUWeightUpdater<T> &PulsedRPUWeightUpdater<T>::operator=(PulsedRPUWeightUpdater<T> &&other) {
  
    RPUWeightUpdater<T>::operator=(std::move(other));
  
    up_ = other.up_;
  
    // pointers
    dblm_ = std::move(other.dblm_);
    sblm_ = std::move(other.sblm_);
    rng_ = std::move(other.rng_);
  
    containers_allocated_ = other.containers_allocated_;
  
    return *this;
  }
  
  template <typename T>
  void PulsedRPUWeightUpdater<T>::setUpPar(const PulsedUpdateMetaParameter<T> &up) {
    up_ = up;
    // check the parameters
    up_.initialize();
  }
  
  template <typename T>
  bool PulsedRPUWeightUpdater<T>::checkForFPUpdate(AbstractRPUDevice<T> *rpu_device_in) {
    return (up_.pulse_type == PulseType::None) || (rpu_device_in == nullptr) ||
           !rpu_device_in->isPulsedDevice() ||
           (rpu_device_in->implements() == DeviceUpdateType::FloatingPoint);
  }
  
  template <typename T>
  void PulsedRPUWeightUpdater<T>::updateVectorWithDevice(
      T **weights,
      const T *x_input,
      const int x_inc,
      const T *d_input,
      const int d_inc,
      const T learning_rate,
      const int m_batch_info,
      AbstractRPUDevice<T> *rpu_device_in) {
    if (!learning_rate) {
      return; // do nothing
    }
  
    // handle cases with no device or FP device
    if (rpu_device_in != nullptr && rpu_device_in->hasDirectUpdate()) {
      rpu_device_in->doDirectVectorUpdate(
          weights, x_input, x_inc, d_input, d_inc, learning_rate, m_batch_info, up_);
      return;
    } else if (up_.pulse_type == PulseType::NoneWithDevice || checkForFPUpdate(rpu_device_in)) {
  
      RPUWeightUpdater<T>::updateVector(weights, x_input, x_inc, d_input, d_inc, learning_rate);
  
      if (up_.pulse_type == PulseType::NoneWithDevice) {
        rpu_device_in->clipWeights(weights, (T)-1.0);
      }
      return;
    }
  
    auto *rpu_device = static_cast<PulsedRPUDeviceBase<T> *>(rpu_device_in);
  
    // check learning rate and update management
    T weight_granularity = rpu_device->getWeightGranularity();
  
    // pulsed device update
    rpu_device->initUpdateCycle(weights, up_, learning_rate, m_batch_info);
  
    if (sblm_->supports(up_.pulse_type)) {
      // envoke sparse bit line maker to get the counts and indices
      int BL = sblm_->makeCounts(
          x_input, x_inc, d_input, d_inc, &*rng_, learning_rate < 0 ? -learning_rate : learning_rate,
          weight_granularity, up_);
      // positive LR actually means that positive signs *decrease* the weight (as in SGD).
      int lr_sign = learning_rate < 0 ? -1 : 1;
  
      if (BL > 0) {
  
        int *x_counts_p;
        int *x_counts_n;
        int *d_counts;
        int **x_indices_p;
        int **x_indices_n;
        int **d_indices;
  
        bool do_negative_separatly = sblm_->getCountsAndIndices(
            x_counts_p, x_counts_n, d_counts, x_indices_p, x_indices_n, d_indices);
  
        for (int k = 0; k < BL; k++) {
          if (d_counts[k] > 0) {
            for (int ii = 0; ii < d_counts[k]; ii++) {
  
              int i_signed = d_indices[k][ii];
              int d_sign = i_signed < 0 ? -lr_sign : lr_sign;
              int i = i_signed < 0 ? -i_signed - 1 : i_signed - 1;
  
              // let rpu_device decide how to update w
              if (x_counts_p[k] > 0) {
                rpu_device->doSparseUpdate(weights, i, x_indices_p[k], x_counts_p[k], d_sign, &*rng_);
              }
              if (do_negative_separatly) {
                if (x_counts_n[k] > 0) {
                  rpu_device->doSparseUpdate(
                      weights, i, x_indices_n[k], x_counts_n[k], d_sign, &*rng_);
                }
              }
            }
          }
        }
      }
    } else {
      // use dense update
      int *coincidences = dblm_->makeCoincidences(
          x_input, x_inc, d_input, d_inc, &*rng_, learning_rate, weight_granularity, up_);
      rpu_device->doDenseUpdate(weights, coincidences, &*rng_);
    }
  
    rpu_device->finishUpdateCycle(weights, up_, learning_rate, m_batch_info);
  }
  
  namespace test_helper {
  void getSparseCountsFromCounts(
      int **&sparse_indices, int *&sparse_counts, uint32_t *&counts, int K, int size) {
  
    for (int k = 0; k < K; k++) { // BL
      sparse_counts[k] = 0;
    }
  
    for (int i = 0; i < size; i++) { // vector index
      int nK32 = (K + 1 + 31) / 32;
      uint32_t one = 1;
      uint32_t negative = counts[i] & one; // first bit is sign bit
      int s = 0;                           // overall BL
      for (int j = 0; j < nK32; j++) {     // BL in chunks of 32 bits
        uint32_t c = counts[i + j * size];
        for (int l = (j == 0 ? 1 : 0); l < 32; l++) {
  
          if ((c & (one << l)) != 0) {
            int iplus1 = negative > 0 ? -(i + 1) : (i + 1);
            sparse_indices[s][sparse_counts[s]++] = iplus1;
          };
          s++;
          if (s == K) {
            break;
          }
        }
      }
    }
  }
  } // namespace test_helper
  
  template <typename T>
  void PulsedRPUWeightUpdater<T>::updateVectorWithDeviceAndCounts(
      T **weights,
      const T *x_input,
      const int x_inc,
      const T *d_input,
      const int d_inc,
      const T learning_rate,
      const int m_batch_info,
      PulsedRPUDeviceBase<T> *rpu_device,
      uint32_t *x_counts32,
      uint32_t *d_counts32) {
  
    // for debugging: use cuda format of counts to update.
    // simply generate some fake bit lines (to setup the memory etc and to get the current BL)
    if (!sblm_->supports(up_.pulse_type)) {
      RPU_FATAL("Requested pulse type not supported.");
    }
  
    int BL = sblm_->makeCounts(
        x_input, x_inc, d_input, d_inc, &*rng_, fabs(learning_rate),
        rpu_device->getWeightGranularity(), up_);
  
    // translate to sparse format
    int *x_counts_p;
    int *x_counts_n;
    int *d_counts;
    int **x_indices_p;
    int **x_indices_n;
    int **d_indices;
  
    bool do_negative_separatly = sblm_->getCountsAndIndices(
        x_counts_p, x_counts_n, d_counts, x_indices_p, x_indices_n, d_indices);
    if (do_negative_separatly) {
      RPU_FATAL("no supported mode for debugging currently")
    }
    // could test for similarity of given counts with generated ones...
  
    // translate given counts
    test_helper::getSparseCountsFromCounts(x_indices_p, x_counts_p, x_counts32, BL, this->x_size_);
    test_helper::getSparseCountsFromCounts(d_indices, d_counts, d_counts32, BL, this->d_size_);
  
    // pulsed device update
    rpu_device->initUpdateCycle(weights, up_, learning_rate, m_batch_info);
  
    // sblm_->printCounts(BL);
    // for info: in BLM (cuda) makeCounts is additional info for debugging the bit line makers
    if (BL > 0) {
      int lr_sign = learning_rate < 0 ? -1 : 1;
      for (int k = 0; k < BL; k++) {
        if (d_counts[k] > 0) {
          for (int ii = 0; ii < d_counts[k]; ii++) {
  
            int i_signed = d_indices[k][ii];
            int d_sign = i_signed < 0 ? -lr_sign : lr_sign;
            int i = (i_signed < 0 ? -i_signed : i_signed) - 1;
            if (x_counts_p[k] > 0) {
              rpu_device->doSparseUpdate(weights, i, x_indices_p[k], x_counts_p[k], d_sign, &*rng_);
            }
          }
        }
      }
    }
    rpu_device->finishUpdateCycle(weights, up_, learning_rate, m_batch_info);
  }
  
  template class PulsedRPUWeightUpdater<float>;
  #ifdef RPU_USE_DOUBLE
  template class PulsedRPUWeightUpdater<double>;
  #endif
  
  } // namespace RPU
  