@@ -76,9 +76,7 @@ optimizer& optimizer::operator=(const optimizer& other) {
   m_gradient_allreduce_needed = other.m_gradient_allreduce_needed;
   m_gradient_allreduce_started = other.m_gradient_allreduce_started;
   m_gradient_allreduce_finished = other.m_gradient_allreduce_finished;
-  #ifdef LBANN_NBALLREDUCE_GRADIENT
   m_gradient_allreduce_started = other.m_gradient_allreduce_started;
-  #endif  // LBANN_NBALLREDUCE_GRADIENT
 
   // Deep copy matrices
   if (m_gradient != nullptr) { delete m_gradient; }
////////////////////////////////////////////////////////////////////////////////
// Copyright (c) 2014-2016, Lawrence Livermore National Security, LLC.
// Produced at the Lawrence Livermore National Laboratory.
// Written by the LBANN Research Team (B. Van Essen, et al.) listed in
// the CONTRIBUTORS file. <lbann-dev@llnl.gov>
//
// LLNL-CODE-697807.
// All rights reserved.
//
// This file is part of LBANN: Livermore Big Artificial Neural Network
// Toolkit. For details, see http://software.llnl.gov/LBANN or
// https://github.com/LLNL/LBANN.
//
// Licensed under the Apache License, Version 2.0 (the "Licensee"); you
// may not use this file except in compliance with the License.  You may
// obtain a copy of the License at:
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the license.
////////////////////////////////////////////////////////////////////////////////

#include "lbann/optimizers/optimizer.hpp"
#include "lbann/utils/cublas_wrapper.hpp"
#include "lbann/utils/timer.hpp"

namespace lbann {

optimizer::optimizer(lbann_comm *comm, DataType learning_rate)
  : m_comm(comm),
    m_cudnn(nullptr),
    m_weights(nullptr),
    m_learning_rate(learning_rate),
    m_gradient(nullptr),
    m_gradient_staging(nullptr),
    m_gradient_allreduce_needed(false),
    m_gradient_allreduce_started(false),
    m_gradient_allreduce_finished(false) {}

optimizer::optimizer(const optimizer& other)
  : m_comm(other.m_comm),
    m_cudnn(other.m_cudnn),
    m_weights(other.m_weights),
    m_learning_rate(other.m_learning_rate),
    m_gradient(other.m_gradient),
    #ifdef LBANN_HAS_CUDNN
    m_gradient_d(other.m_gradient_d),
    #endif // LBANN_HAS_CUDNN
    m_gradient_staging(other.m_gradient_staging),
    #ifdef LBANN_HAS_CUDNN
    m_gradient_staging_d(other.m_gradient_staging_d),
    #endif // LBANN_HAS_CUDNN
    m_gradient_allreduce_needed(other.m_gradient_allreduce_needed),
    m_gradient_allreduce_started(other.m_gradient_allreduce_started),
    m_gradient_allreduce_finished(other.m_gradient_allreduce_finished),
    m_step_time(other.m_step_time)
{
  if (m_gradient != nullptr) {
    m_gradient = m_gradient->Copy();
  }
  if (m_gradient_staging != nullptr) {
    m_gradient_staging = m_gradient_staging->Copy();
  }
}

optimizer& optimizer::operator=(const optimizer& other) {
  m_comm = other.m_comm;
  m_cudnn = other.m_cudnn;
  m_weights = other.m_weights;
  m_learning_rate = other.m_learning_rate;
  m_step_time = other.m_step_time;
  m_gradient_allreduce_needed = other.m_gradient_allreduce_needed;
  m_gradient_allreduce_started = other.m_gradient_allreduce_started;
  m_gradient_allreduce_finished = other.m_gradient_allreduce_finished;
  #ifdef LBANN_NBALLREDUCE_GRADIENT
  m_gradient_allreduce_started = other.m_gradient_allreduce_started;
  #endif  // LBANN_NBALLREDUCE_GRADIENT

  // Deep copy matrices
  if (m_gradient != nullptr) { delete m_gradient; }
  if (m_gradient_staging != nullptr) { delete m_gradient_staging; }
  m_gradient = other.m_gradient;
  m_gradient_staging = other.m_gradient_staging;
  if (m_gradient != nullptr) {
    m_gradient = m_gradient->Copy();
  }
  if (m_gradient_staging != nullptr) {
    m_gradient_staging = m_gradient_staging->Copy();
  }

  // Copy GPU data
  #ifdef LBANN_HAS_CUDNN
  m_gradient_d = other.m_gradient_d;
  m_gradient_staging_d = other.m_gradient_staging_d;
  #endif // LBANN_HAS_CUDNN

  return *this;
}

optimizer::~optimizer() {
  if (m_gradient != nullptr) { delete m_gradient; }
  if (m_gradient_staging != nullptr)  { delete m_gradient_staging; }
}

std::string optimizer::get_description() const {
  std::stringstream ss;
  ss << get_type();
  if (m_weights != nullptr) {
    ss << " (optimizing " << m_weights->get_name() << ")";
  }
  ss << "; learning_rate=" << m_learning_rate;
  return ss.str();
}

weights& optimizer::get_weights() {
  if (!is_initialized()) {
    std::stringstream err;
    err << __FILE__ << " " << __LINE__ << " :: "
        << "attempted to access the weights being optimized before they are set";
    throw lbann_exception(err.str());
  }
  return *m_weights;
}

const AbsDistMat& optimizer::get_gradient() {

  // Check if gradient is initialized
  if (!is_initialized()) {
    std::stringstream err;
    err << __FILE__ << " " << __LINE__ << " :: "
        << "attempted to access gradients before they are set up";
    throw lbann_exception(err.str());
  }

  // Perform allreduce on staging matrix if needed
  if (m_gradient_allreduce_needed && !m_gradient_allreduce_started) {
    start_gradient_staging_allreduce();
  }
  if (m_gradient_allreduce_started && !m_gradient_allreduce_finished) {
    #ifdef LBANN_NBALLREDUCE_GRADIENT
    m_comm->wait(m_gradient_allreduce_req);
    #endif // LBANN_NBALLREDUCE_GRADIENT
    m_gradient_allreduce_finished = true;
  }
  if (m_gradient_allreduce_needed) {
    if (m_cudnn == nullptr) {
      add_to_gradient(*m_gradient_staging);
    } else {
      #ifdef LBANN_HAS_CUDNN
      add_to_gradient(m_gradient_staging_d);
      #endif // LBANN_HAS_CUDNN
    }
  }
  m_gradient_allreduce_needed = false;
  m_gradient_allreduce_started = false;
  m_gradient_allreduce_finished = false;

  // Return CPU gradient matrix
  if (m_cudnn != nullptr) {
    #ifdef LBANN_HAS_CUDNN
    m_cudnn->copy_from_gpu(0,
                           m_gradient->Matrix(),
                           m_gradient_d.get_locked_data(0),
                           m_gradient_d.get_leading_dim());
    m_cudnn->synchronize();
    #endif // LBANN_HAS_CUDNN
  }
  return *m_gradient;

}

#ifdef LBANN_HAS_CUDNN
const cudnn::matrix& optimizer::get_gradient_gpu() {

  // Check if gradient is initialized
  if (!is_initialized()) {
    std::stringstream err;
    err << __FILE__ << " " << __LINE__ << " :: "
        << "attempted to access gradients before they are set up";
    throw lbann_exception(err.str());
  }
  if (m_cudnn == nullptr) {
    std::stringstream err;
    err << __FILE__ << " " << __LINE__ << " :: "
        << "attempted to get GPU gradient, but GPU is not set up";
    throw lbann_exception(err.str());
  }

  // Check if all gradient sources have made contributions
  m_gradient_sources.erase(nullptr);
  if (!m_gradient_sources.empty()) {
    std::stringstream err;
    err << __FILE__ << " " << __LINE__ << " :: "
        << "attempted to access gradient before all gradient sources "
        << "have made contributions "
        << "(missing " << m_gradient_sources.size() << " sources)";
    throw lbann_exception(err.str());
  }

  // Perform allreduce on staging matrix if needed
  if (m_gradient_allreduce_needed && !m_gradient_allreduce_started) {
    start_gradient_staging_allreduce();
  }
  if (m_gradient_allreduce_started && !m_gradient_allreduce_finished) {
    #ifdef LBANN_NBALLREDUCE_GRADIENT
    m_comm->wait(m_gradient_allreduce_req);
    #endif // LBANN_NBALLREDUCE_GRADIENT
    m_gradient_allreduce_finished = true;
  }
  if (m_gradient_allreduce_needed) {
    add_to_gradient(m_gradient_staging_d);
  }
  m_gradient_allreduce_needed = false;
  m_gradient_allreduce_started = false;
  m_gradient_allreduce_finished = false;

  // Return gradient
  return m_gradient_d;

}
#endif // LBANN_HAS_CUDNN

void optimizer::start_gradient_staging_allreduce() {
  if (!m_gradient_allreduce_needed || m_gradient_allreduce_started) {
    return;
  }

  m_gradient_allreduce_started = true;
  if (m_cudnn == nullptr) {
    #ifdef LBANN_NBALLREDUCE_GRADIENT
    m_comm->nb_allreduce(*m_gradient_staging,
                         m_gradient_staging->RedundantComm(),
                         m_gradient_allreduce_req);
    m_gradient_allreduce_finished = false;
    #else
    m_comm->allreduce(*m_gradient_staging,
                      m_gradient_staging->RedundantComm());
    m_gradient_allreduce_finished = true;
    #endif // LBANN_NBALLREDUCE_GRADIENT
  } else {
    #ifndef LBANN_HAS_CUDNN
    throw lbann_exception("optimizer: cuDNN not detected");
    #else
    m_cudnn->global_allreduce_on_gpus(m_gradient_staging_d.get_data(),
                                      m_gradient_staging_d.get_height(),
                                      m_gradient_staging_d.get_width_per_gpu(),
                                      m_gradient->RedundantComm());
    m_gradient_allreduce_finished = true;
    #endif // LBANN_HAS_CUDNN
  }

}

void optimizer::clear_gradient() {

  // Clear matrices
  if (m_cudnn == nullptr) {
    El::Zero(*m_gradient);
  } else {
    #ifdef LBANN_HAS_CUDNN
    m_gradient_d.zero();
    #endif // LBANN_HAS_CUDNN
  }

  // Reset gradient allreduce flags
  m_gradient_allreduce_needed = false;
  m_gradient_allreduce_started = false;
  m_gradient_allreduce_finished = false;

}

void optimizer::add_to_gradient(const AbsDistMat& gradient,
                                DataType scale) {
  if (!is_initialized()) {
    std::stringstream err;
    err << __FILE__ << " " << __LINE__ << " :: "
        << "attempted to access gradients before they are set up";
    throw lbann_exception(err.str());
  }
  if (scale != DataType(0)) {
    if (m_cudnn == nullptr) {
      El::Axpy(scale, gradient, *m_gradient);
    } else {
      #ifndef LBANN_HAS_CUDNN
      throw lbann_exception("optimizer: cuDNN not detected");
      #else
      cudnn::matrix gradient_d(m_cudnn,
                               gradient.LocalHeight(),
                               gradient.LocalWidth());
      m_cudnn->broadcast_to_gpus(gradient_d.get_data(),
                                 gradient.LockedMatrix());
      add_to_gradient(gradient_d, scale);
      #endif // LBANN_HAS_CUDNN
    }
  }
}
#ifdef LBANN_HAS_CUDNN
void optimizer::add_to_gradient(const cudnn::matrix& gradient_d,
                                DataType scale) {
  if (!is_initialized()) {
    std::stringstream err;
    err << __FILE__ << " " << __LINE__ << " :: "
        << "attempted to access gradients before they are set up";
    throw lbann_exception(err.str());
  }
  if (m_cudnn == nullptr) {
    std::stringstream err;
    err << __FILE__ << " " << __LINE__ << " :: "
        << "attempted to add to GPU gradient, but GPU is not set up";
    throw lbann_exception(err.str());
  }
  if (scale != DataType(0)) {
    for(int i = 0; i < m_cudnn->get_num_gpus(); ++i) {
      CHECK_CUDA(cudaSetDevice(m_cudnn->get_gpu(i)));
      CHECK_CUBLAS(cublas::axpy(m_cudnn->get_cublas_handle(i),
                                m_weights->get_size(),
                                scale, gradient_d.get_locked_data(i), 1,
                                m_gradient_d.get_data(i), 1));
    }
  }
}
#endif // LBANN_HAS_CUDNN

void optimizer::add_to_gradient_staging(const AbsDistMat& gradient,
                                        DataType scale) {
  if (!is_initialized()) {
    std::stringstream err;
    err << __FILE__ << " " << __LINE__ << " :: "
        << "attempted to access gradients before they are set up";
    throw lbann_exception(err.str());
  }
  if (m_gradient_allreduce_started) {
    std::stringstream err;
    err << __FILE__ << " " << __LINE__ << " :: "
        << "attempted to add to staging matrix after gradient accumulation has started";
    throw lbann_exception(err.str());
  }
  if (scale != DataType(0)) {

    // Clear staging matrix if needed
    if (!m_gradient_allreduce_needed) {
      if (m_cudnn == nullptr) {
        El::Zero(*m_gradient_staging);
      } else {
        #ifndef LBANN_HAS_CUDNN
        throw lbann_exception("optimizer: cuDNN not detected");
        #else
        m_gradient_staging_d.zero();
        #endif // LBANN_HAS_CUDNN
      }
    }
    m_gradient_allreduce_needed = true;

    // Add to staging matrix
    if (m_cudnn == nullptr) {
      El::Axpy(scale, gradient, *m_gradient_staging);
    } else {
      #ifndef LBANN_HAS_CUDNN
      throw lbann_exception("optimizer: cuDNN not detected");
      #else
      cudnn::matrix gradient_d(m_cudnn,
                               gradient.LocalHeight(),
                               gradient.LocalWidth());
      gradient_d.zero();
      m_cudnn->copy_to_gpu(0,
                           gradient_d.get_data(0),
                           gradient.LockedMatrix(),
                           gradient_d.get_leading_dim());
      add_to_gradient_staging(gradient_d, scale);
      #endif // LBANN_HAS_CUDNN
    }

  }
}
#ifdef LBANN_HAS_CUDNN
void optimizer::add_to_gradient_staging(const cudnn::matrix& gradient_d,
                                        DataType scale) {
  if (!is_initialized()) {
    std::stringstream err;
    err << __FILE__ << " " << __LINE__ << " :: "
        << "attempted to access gradients before they are set up";
    throw lbann_exception(err.str());
  }
  if (m_gradient_allreduce_started) {
    std::stringstream err;
    err << __FILE__ << " " << __LINE__ << " :: "
        << "attempted to add to staging matrix after gradient accumulation has started";
    throw lbann_exception(err.str());
  }
  if (m_cudnn == nullptr) {
    std::stringstream err;
    err << __FILE__ << " " << __LINE__ << " :: "
        << "attempted to add to GPU gradient, but GPU is not set up";
    throw lbann_exception(err.str());
  }
  if (scale != DataType(0)) {

    // Clear staging matrix if needed
    if (!m_gradient_allreduce_needed) {
      m_gradient_staging_d.zero();
    }
    m_gradient_allreduce_needed = true;

    // Add to staging matrix
    for(int i = 0; i < m_cudnn->get_num_gpus(); ++i) {
      CHECK_CUDA(cudaSetDevice(m_cudnn->get_gpu(i)));
      CHECK_CUBLAS(cublas::axpy(m_cudnn->get_cublas_handle(i),
                                m_weights->get_size(),
                                scale, gradient_d.get_locked_data(i), 1,
                                m_gradient_staging_d.get_data(i), 1));
    }

  }
}
#endif // LBANN_HAS_CUDNN

void optimizer::add_gradient_source(const void* source) {
  if (source != nullptr) {
    m_gradient_sources.insert(source);
  }
}

void optimizer::remove_gradient_source(const void* source) {
  m_gradient_sources.erase(nullptr);
  m_gradient_sources.erase(source);
  if (m_gradient_sources.empty()) {
    start_gradient_staging_allreduce();
  }
}

void optimizer::setup(weights& w) {
  if (is_initialized()) {
    std::stringstream err;
    err << __FILE__ << " " << __LINE__ << " :: "
        << "attempted to setup an optimizer that is already set up";
    throw lbann_exception(err.str());
  }
  set_weights(w);

  // Initialize matrices
  const int height = m_weights->get_matrix_height();
  const int width = m_weights->get_matrix_width();
  const AbsDistMat& values = m_weights->get_values();

  m_gradient = values.Construct(values.Grid(), values.Root());
  m_gradient_staging = values.Construct(values.Grid(), values.Root());
  m_gradient->Resize(height, width);
  m_gradient_staging->Resize(height, width);

  // Initialize GPU
  m_cudnn = m_weights->m_cudnn;
  if (m_cudnn != nullptr) {
#ifdef LBANN_HAS_CUDNN
    m_gradient_d = cudnn::matrix(m_cudnn, height, width);
    m_gradient_staging_d = cudnn::matrix(m_cudnn, height, width);
#endif // LBANN_HAS_CUDNN
  }

  // Initialize with zero gradient
  clear_gradient();

}

void optimizer::step() {
  if (!is_initialized()) {
    std::stringstream err;
    err << __FILE__ << " " << __LINE__ << " :: "
        << "optimizer must be set up before performing optimization step";
    throw lbann_exception(err.str());
  }

  double step_start = get_time();
  // Apply optimization step
  if (m_cudnn != nullptr) {
  #ifdef LBANN_HAS_CUDNN
    cudnn::matrix values_d(m_cudnn);
    values_d.attach(m_weights->m_values_d, m_weights->get_size());
    const auto& gradient_d = get_gradient_gpu();
    step_compute_gpu(values_d, gradient_d);
  #endif // LBANN_HAS_CUDNN
  } else {
    m_weights->get_values(); // Move data to CPU
    auto& values = *m_weights->m_values;
    const auto& gradient = get_gradient();
    step_compute(values, gradient);
  }

  // Clear gradients
  clear_gradient();

  m_step_time += get_time() - step_start;

}

#ifdef LBANN_HAS_CUDNN
void optimizer::step_compute_gpu(cudnn::matrix& values_d,
                                 const cudnn::matrix& gradient_d) {
  m_cudnn->copy_from_gpu(0, m_weights->m_values->Matrix(), values_d.get_locked_data(0));
  m_cudnn->copy_from_gpu(0, m_gradient->Matrix(), gradient_d.get_locked_data(0));
  m_cudnn->synchronize();
  step_compute(*m_weights->m_values, *m_gradient);
  m_cudnn->broadcast_to_gpus(values_d.get_data(), m_weights->m_values->LockedMatrix());
}
#endif // LBANN_HAS_CUDNN

//************************************************************************
// Checkpointing
//************************************************************************

bool optimizer::save_to_checkpoint_shared(persist& p, std::string m_name) {
  //  m_learning_rate;
  /** Running count of the time spent in step(). */
  //  double m_step_time = 0.0;
  p.write_datatype(persist_type::train, "learning_rate", m_learning_rate);
  return true;
}

bool optimizer::load_from_checkpoint_shared(persist& p, std::string m_name) {
  p.read_datatype(persist_type::train, "learning_rate", &m_learning_rate);
  MPI_Bcast(&m_learning_rate, 1, MPI_FLOAT, 0, MPI_COMM_WORLD);
  return true;
}
}  // namespace lbann

We no longer check for the `LBANN_NBALLREDUCE_GRADIENT` macro. LBANN will always use non-blocking allreduces if they are available.