lazy_alloc_array.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); 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.
 */

#ifndef MXNET_COMMON_LAZY_ALLOC_ARRAY_H_
#define MXNET_COMMON_LAZY_ALLOC_ARRAY_H_

#include <dmlc/logging.h>
#include <memory>
#include <mutex>
#include <array>
#include <vector>
#include <atomic>

namespace mxnet {
namespace common {

template<typename TElem>
class LazyAllocArray {
 public:
  LazyAllocArray();
  template<typename FCreate>
  inline std::shared_ptr<TElem> Get(int index, FCreate creator);
  template<typename FVisit>
  inline void ForEach(FVisit fvisit);
  inline void Clear();

  void SignalForKill();

 private:
  template<typename SyncObject>
  class unique_unlock {
   public:
    explicit unique_unlock(std::unique_lock<SyncObject> *lock)
      : lock_(lock) {
      if (lock_) {
        lock_->unlock();
      }
    }
    ~unique_unlock() {
      if (lock_) {
        lock_->lock();
      }
    }
   private:
    std::unique_lock<SyncObject> *lock_;
  };

  static constexpr std::size_t kInitSize = 16;
  std::mutex create_mutex_;
  std::array<std::shared_ptr<TElem>, kInitSize> head_;
  std::vector<std::shared_ptr<TElem> > more_;
  std::atomic<bool> exit_now_;
};

template<typename TElem>
inline LazyAllocArray<TElem>::LazyAllocArray()
  : exit_now_(false) {
}

// implementations
template<typename TElem>
template<typename FCreate>
inline std::shared_ptr<TElem> LazyAllocArray<TElem>::Get(int index, FCreate creator) {
  CHECK_GE(index, 0);
  size_t idx = static_cast<size_t>(index);
  if (idx < kInitSize) {
    std::shared_ptr<TElem> ptr = head_[idx];
    if (ptr) {
      return ptr;
    } else {
      std::lock_guard<std::mutex> lock(create_mutex_);
      if (!exit_now_.load()) {
        std::shared_ptr<TElem> ptr = head_[idx];
        if (ptr) {
          return ptr;
        }
        ptr = head_[idx] = std::shared_ptr<TElem>(creator());
        return ptr;
      }
    }
  } else {
    std::lock_guard<std::mutex> lock(create_mutex_);
    if (!exit_now_.load()) {
      idx -= kInitSize;
      if (more_.size() <= idx) {
        more_.reserve(idx + 1);
        while (more_.size() <= idx) {
          more_.push_back(std::shared_ptr<TElem>(nullptr));
        }
      }
      std::shared_ptr<TElem> ptr = more_[idx];
      if (ptr) {
        return ptr;
      }
      ptr = more_[idx] = std::shared_ptr<TElem>(creator());
      return ptr;
    }
  }
  return nullptr;
}

template<typename TElem>
inline void LazyAllocArray<TElem>::Clear() {
  std::unique_lock<std::mutex> lock(create_mutex_);
  exit_now_.store(true);
  // Currently, head_ and more_ never get smaller, so it's safe to
  // iterate them outside of the lock.  The loops should catch
  // any growth which might happen when create_mutex_ is unlocked
  for (size_t i = 0; i < head_.size(); ++i) {
    std::shared_ptr<TElem> p = head_[i];
    head_[i] = std::shared_ptr<TElem>(nullptr);
    unique_unlock<std::mutex> unlocker(&lock);
    p = std::shared_ptr<TElem>(nullptr);
  }
  for (size_t i = 0; i < more_.size(); ++i) {
    std::shared_ptr<TElem> p = more_[i];
    more_[i] = std::shared_ptr<TElem>(nullptr);
    unique_unlock<std::mutex> unlocker(&lock);
    p = std::shared_ptr<TElem>(nullptr);
  }
}

template<typename TElem>
template<typename FVisit>
inline void LazyAllocArray<TElem>::ForEach(FVisit fvisit) {
  std::lock_guard<std::mutex> lock(create_mutex_);
  for (size_t i = 0; i < head_.size(); ++i) {
    if (head_[i].get() != nullptr) {
      fvisit(i, head_[i].get());
    }
  }
  for (size_t i = 0; i < more_.size(); ++i) {
    if (more_[i].get() != nullptr) {
      fvisit(i + kInitSize, more_[i].get());
    }
  }
}

template<typename TElem>
inline void LazyAllocArray<TElem>::SignalForKill() {
  std::lock_guard<std::mutex> lock(create_mutex_);
  exit_now_.store(true);
}

}  // namespace common
}  // namespace mxnet
#endif  // MXNET_COMMON_LAZY_ALLOC_ARRAY_H_