#include <cstdlib>

#include <string>

#include <stdio.h>

#include <iostream>

#include <cstring>

#include <sstream>

#include <pthread.h>

#include <glog/logging.h>

#include <boost/shared_ptr.hpp>

#include <boost/algorithm/string.hpp>

#include <sys/mman.h>

#include <sys/socket.h>

#include <sys/stat.h>

#include <sys/ioctl.h>

#include <fcntl.h>

#include <netdb.h>

#include <net/if.h>

#include <unistd.h>

#include <iomanip>

#include <caffe/caffe.hpp>

#include "caffe/filler.hpp"

#include "caffe/parallel.hpp"

using namespace std;

namespace caffe {

void Meter::show(std::ostream& s) const {

}

template<typename Dtype>

static size_t len(const vector<shared_ptr<Blob<Dtype> > >& params) {

}

template<typename Dtype>

static size_t align(const size_t len) {

}

template<typename Dtype>

Params<Dtype>::Params(const vector<shared_ptr<Blob<Dtype> > >& blobs,

const string& file_map)

: len_used_(len<Dtype>(blobs)),

len_buff_(align<Dtype>(len_used_)) {

bool exists = false;

if (file_map.empty()) {

CaffeMallocHost((void**) &cpu_, len_buff_ * sizeof(Dtype));

memset(cpu_, 0, len_buff_ * sizeof(Dtype));

} else {

struct stat st_buf;

exists = stat(file_map.c_str(), &st_buf) == 0;

int fd = open(file_map.c_str(), O_RDWR | O_CREAT, //

S_IRWXU | S_IRWXG | S_IRWXO);

CHECK(!ftruncate(fd, len_buff_ * sizeof(Dtype)));

cpu_ = (Dtype*) mmap(NULL, //

len_buff_ * sizeof(Dtype),

PROT_READ | PROT_WRITE,

MAP_SHARED, fd, 0);

close(fd);

}

Dtype* cpu = cpu_;

for (int i = 0; i < blobs.size(); ++i) {

int size = blobs[i]->data()->size();

// Init to current values of blobs if file doesn't already exists

if (!exists)

memcpy(cpu, blobs[i]->data()->cpu_data(), size);

cpu += size / sizeof(Dtype);

CHECK(cpu <= cpu_ + len_used_);

}

size_t check = 0;

for (int i = 0; i < blobs.size(); ++i)

check += blobs[i]->count();

Dtype* expect = cpu_ + check;

CHECK_EQ(expect, cpu);

iterations_ = 0;

}

template<typename Dtype>

Params<Dtype>::~Params() {

CaffeFreeHost((void*) cpu_);

}

template<typename Dtype>

void Params<Dtype>::configure(Solver<Dtype>* solver) const {

// Replace weights

vector<shared_ptr<Blob<Dtype> > > &blobs = solver->net()->params();

Dtype* cpu = cpu_;

for (int i = 0; i < blobs.size(); ++i) {

blobs[i]->data()->set_cpu_data(cpu);

cpu += blobs[i]->data()->size() / sizeof(Dtype);

CHECK(cpu <= cpu_ + len_used_);

}

// Check sizes

size_t check = 0;

for (int i = 0; i < blobs.size(); ++i)

check += blobs[i]->count();

Dtype* expect = cpu_ + check;

CHECK_EQ(expect, cpu);

solver->iter_total(&iterations_);

}

#ifndef CPU_ONLY

#include <cuda_runtime.h>

template<typename Dtype>

GPUParams<Dtype>::GPUParams(const Params<Dtype>& params, int device)

: params_(params),

device_(device) {

int current;

CUDA_CHECK(cudaGetDevice(&current));

CUDA_CHECK(cudaSetDevice(device));

const size_t size = params.len_buff() * sizeof(Dtype);

CUDA_CHECK(cudaMalloc((void** ) &gpu_, size));

CUDA_CHECK(cudaMemcpy(gpu_, params.cpu(), size, cudaMemcpyHostToDevice));

CUDA_CHECK(cudaSetDevice(current));

}

template<typename Dtype>

GPUParams<Dtype>::~GPUParams() {

CUDA_CHECK(cudaFree((void* ) gpu_));

}

template<typename Dtype>

void GPUParams<Dtype>::configure(Solver<Dtype>* solver) const {

// Replace GPU weights

vector<shared_ptr<Blob<Dtype> > > &blobs = solver->net()->params();

Dtype* gpu = gpu_;

for (int i = 0; i < blobs.size(); ++i) {

blobs[i]->data()->set_gpu_data(gpu);

gpu += blobs[i]->data()->size() / sizeof(Dtype);

CHECK(gpu <= gpu_ + params_.len_used());

}

size_t check = 0;

for (int i = 0; i < blobs.size(); ++i)

check += blobs[i]->count();

Dtype* expect = gpu_ + check;

CHECK_EQ(expect, gpu);

solver->iter_total(&params_.iterations_);

}

template<typename Dtype>

GPUStream<Dtype>::GPUStream() {

int least, greatest;

cudaDeviceGetStreamPriorityRange(&least, &greatest);

cudaStreamCreateWithPriority(&stream_, cudaStreamNonBlocking, least);

}

template<typename Dtype>

GPUStream<Dtype>::~GPUStream() {

cudaStreamDestroy(stream_);

}

template<typename Dtype>

GPUSync<Dtype>::GPUSync(const GPUParams<Dtype>& params)

: params_(params) {

size_t size = params.params().len_buff() * sizeof(Dtype);

Dtype* gpu = params.gpu();

CUDA_CHECK(cudaMalloc((void** ) &gpu_last_, size));

CUDA_CHECK(cudaMemcpy(gpu_last_, gpu, size, cudaMemcpyDeviceToDevice));

}

template<typename Dtype>

GPUSync<Dtype>::~GPUSync() {

CUDA_CHECK(cudaFree((void* ) gpu_last_));

}

template<typename Dtype>

CPUGPUSync<Dtype>::CPUGPUSync(const GPUParams<Dtype>& params)

: GPUSync<Dtype>(params),

chunks_(chunks(params.params().len_used())),

calls_("calls", CHUNK * sizeof(Dtype)),

cycles_("cycles") {

}

template<typename Dtype>

CPUGPUSync<Dtype>::~CPUGPUSync() {

stop();

}

template<typename Dtype>

void CPUGPUSync<Dtype>::run() {

CUDA_CHECK(cudaSetDevice(this->params_.device()));

GPUStream<Dtype> gpu_stream;

const cudaStream_t& stream = gpu_stream.stream();

// Current cpu values when invoking kernel, gradients on the way back

Dtype* buf;

Dtype* tmp;

CUDA_CHECK(cudaMalloc((void** ) &buf, CHUNK * sizeof(Dtype)));

CaffeMallocHost((void**) &tmp, CHUNK * sizeof(Dtype));

const size_t len = CHUNK * sizeof(Dtype);

// Explicit directions for readability

const cudaMemcpyKind put = cudaMemcpyHostToDevice;

const cudaMemcpyKind get = cudaMemcpyDeviceToHost;

uint32_t index = 0;

Dtype* cpu = this->params_.params().cpu();

Dtype* gpu = this->params_.gpu();

Dtype* last = this->gpu_last_;

uint8_t get_grads = true;

while (!must_stop()) {

size_t off = index * CHUNK;

CUDA_CHECK(cudaMemcpyAsync(buf, &cpu[off], len, put, stream));

// TODO simpler kernel

sync_worker_kernel<Dtype>(gpu, last, &buf, &off, &buf, &get_grads, //

0, 1, stream, CHUNK);

CUDA_CHECK(cudaMemcpyAsync(tmp, buf, len, get, stream));

cudaStreamSynchronize(stream);

for (size_t i = 0; i < CHUNK; ++i)

cpu[off + i] += tmp[i];

if (++index == chunks_) {

index = 0;

cycles_++;

}

calls_++;

}

CaffeFreeHost((void*) tmp);

CUDA_CHECK(cudaFree((void* ) buf));

}

#endif

template<typename Dtype>

DistSync<Dtype>::DistSync(uint32_t nodes, uint32_t chunks)

: nodes_(nodes),

chunks_(chunks),

received_(chunks),

remaining_(chunks),

cycles_("cycles") {

own_start_ = own_until_ = chunk_ = 0;

}

template<typename Dtype>

void DistSync<Dtype>::dist_init(int rank) {

own_start_ = (rank + 0) * chunks_ / nodes_;

own_until_ = (rank + 1) * chunks_ / nodes_;

LOG(INFO)<< "range: " << own_start_ << " " << own_until_;

chunk_ = own_start_;

for (uint32_t chunk = own_start_; chunk < own_until_; ++chunk) {

received_[chunk] = true;

remaining_--;

}

}

template<typename Dtype>

inline int DistSync<Dtype>::chunk_master(uint32_t chunk) {

// TODO find range without loop?

for (int i = nodes_ - 1; i >= 0; --i) {

uint32_t start = i * chunks_ / nodes_;

if (start <= chunk)

return i;

}

CHECK(false);

return -1;

}

INSTANTIATE_CLASS(Params);

#ifndef CPU_ONLY

INSTANTIATE_CLASS(GPUParams);

INSTANTIATE_CLASS(GPUSync);

INSTANTIATE_CLASS(CPUGPUSync);

#endif

INSTANTIATE_CLASS(DistSync);

#ifdef RDMA

ibv_context* IBChannel::open_device(ibv_device* ib_dev) {

ibv_context* context = ibv_open_device(ib_dev);

CHECK(context) << "Open context failed for " << ibv_get_device_name(ib_dev);

return context;

}

ibv_pd* IBChannel::alloc_pd(ibv_context* context) {

ibv_pd* pd = ibv_alloc_pd(context);

CHECK(pd) << "Failed to allocate protection domain";

return pd;

}

IBChannel::IBChannel(ibv_device* ib_dev)

}

static ib_addr mcast_init(ibv_context* context, int port, const ibv_gid* mgid) {

}

ib_addr IBChannel::mcast_create() const {

}

void IBChannel::mcast_join(const ib_addr& addr) const {

mcast_init(context_, PORT, &addr.gid);

}

void IBChannel::mcast_attach_qp(const ib_addr& addr) const {

CHECK(!ibv_attach_mcast(qp_, &addr.gid, addr.lid))

<< "Failed to attach to the multicast group";

}

void IBChannel::start(uint8_t* buf_send, size_t buf_size, bool gpu) const {

}

IBChannel::~IBChannel() {

CHECK(!ibv_destroy_qp(qp_)) << "Failed to destroy QP";

CHECK(!ibv_destroy_cq(cq_)) << "Failed to destroy CQ";

CHECK(!ibv_dereg_mr(mr_send_)) << "Failed to deregister MR";

CHECK(!ibv_dereg_mr(mr_recv_)) << "Failed to deregister MR";

CHECK(!ibv_dealloc_pd(pd_)) << "Failed to deallocate PD";

CHECK(!ibv_close_device(context_)) << "Failed to release context";

free(buf_send_);

free(buf_recv_);

}

bool IBChannel::can_send() const {

return !send_queue_.empty();

}

int IBChannel::send_init(uint8_t*& buf) const {

}

void IBChannel::send(int id, const ib_addr& addr, uint8_t* buf,

}

bool IBChannel::can_recv() const {

return !recv_queue_.empty();

}

int IBChannel::recv(uint8_t*& buf, uint32_t& imm_data) const {

}

void IBChannel::recv_done(int id) const {

}

void IBChannel::poll() const {

}

template<typename Dtype>

IBSync<Dtype>::IBSync(const Params<Dtype>& params, int rank,

const IBChannel& ucast, const IBChannel& mcast,

const vector<ib_addr>& ucast_addrs,

const vector<ib_addr>& mcast_addrs)

: DistSync<Dtype>(ucast_addrs.size(), chunks(params.len_used())),

rank_(rank),

ucast_(ucast),

mcast_(mcast),

ucast_addrs_(ucast_addrs),

mcast_addr_(mcast_addrs[rank]) {

for (int i = 0; i < ucast_addrs_.size(); ++i) {

CHECK(ucast_addrs_[i].ah == NULL);

if (i != rank) {

struct ibv_ah_attr ah_attr;

memset(&ah_attr, 0, sizeof ah_attr);

ah_attr.dlid = (uint16_t) ucast_addrs[i].lid;

ah_attr.sl = (uint8_t) 0; // Service level

ah_attr.src_path_bits = 0;

ah_attr.is_global = 0;

ah_attr.port_num = IBChannel::PORT;

ucast_addrs_[i].ah = ibv_create_ah(ucast.pd_, &ah_attr);

CHECK(ucast_addrs_[i].ah) << "Failed to create address handle";

}

}

struct ibv_ah_attr ah_attr;

memset(&ah_attr, 0, sizeof ah_attr);

ah_attr.grh.dgid = mcast_addr_.gid;

ah_attr.dlid = (uint16_t) mcast_addr_.lid;

ah_attr.sl = (uint8_t) 0; // Service level

ah_attr.src_path_bits = 0;

ah_attr.is_global = 1;

ah_attr.port_num = IBChannel::PORT;

mcast_addr_.ah = ibv_create_ah(mcast.pd_, &ah_attr);

CHECK(mcast_addr_.ah) << "Failed to create address handle";

for (int i = 0; i < mcast_addrs.size(); ++i) {

if (i != rank) {

mcast_.mcast_join(mcast_addrs[i]);

mcast_.mcast_attach_qp(mcast_addrs[i]);

}

}

this->dist_init(rank);

}

template<typename Dtype>

IBSync<Dtype>::~IBSync() {

for (int i = 0; i < this->ucast_addrs_.size(); ++i) {

if (i == rank_) {

CHECK(!ibv_destroy_ah(this->ucast_addrs_[i].ah))

<< "Failed to destroy ucast AH";

}

}

CHECK(!ibv_destroy_ah(this->mcast_addr_.ah)) << "Failed to destroy mcast AH";

}

template<typename Dtype>

CPUIBSync<Dtype>::CPUIBSync(const Params<Dtype>& params, int rank,

const IBChannel& ucast, const IBChannel& mcast,

const vector<ib_addr>& ucast_addrs,

const vector<ib_addr>& mcast_addrs)

: IBSync<Dtype>(params, rank, ucast, mcast, ucast_addrs, mcast_addrs) {

cpu_ = params.cpu();

CaffeMallocHost((void**) &cpu_last_, params.len_buff() * sizeof(Dtype));

memcpy(cpu_last_, cpu_, params.len_used() * sizeof(Dtype));

}

template<typename Dtype>

CPUIBSync<Dtype>::~CPUIBSync() {

CaffeFreeHost((void*) cpu_last_);

}

template<typename Dtype>

void CPUIBSync<Dtype>::run() {

// TODO

}

template<typename Dtype>

GPUIBSync<Dtype>::GPUIBSync(const GPUParams<Dtype>& params, int rank,

const IBChannel& ucast, const IBChannel& mcast,

const vector<ib_addr>& ucast_addrs,

const vector<ib_addr>& mcast_addrs)

: GPUSync<Dtype>(params),

IBSync<Dtype>(params.params(), rank, //

ucast, mcast, //

ucast_addrs, mcast_addrs) {

gpu_ = params.gpu();

int device;

CUDA_CHECK(cudaGetDevice(&device));

CUDA_CHECK(cudaSetDevice(params.device()));

size_t size = params.params().len_buff() * sizeof(Dtype);

CUDA_CHECK(cudaMalloc((void** ) &gpu_last_, size));

CUDA_CHECK(cudaMemcpy(gpu_last_, gpu_, size, cudaMemcpyDeviceToDevice));

CUDA_CHECK(cudaSetDevice(device));

}

template<typename Dtype>

GPUIBSync<Dtype>::~GPUIBSync() {

CUDA_CHECK(cudaFree((void* ) gpu_last_));

}

class Queue {

};

class EventQueue : Queue {

};

template<typename Dtype>

void GPUIBSync<Dtype>::run() {

CUDA_CHECK(cudaSetDevice(this->params_.device()));

const IBChannel& ucast = this->ucast_;

const IBChannel& mcast = this->mcast_;

ucast.start(NULL, 0, true);

mcast.start((uint8_t*) gpu_, (size_t) this->chunks_ * IBChannel::MTU, true);

GPUStream<Dtype> master_stream;

Queue master_queue;

uint16_t master_ids_[FRAMES];

EventQueue master_events(master_stream.stream());

GPUStream<Dtype> worker_stream;

Queue worker_queue;

struct worker_item {

};

worker_item worker_items[FRAMES];

EventQueue worker_events(worker_stream.stream());

const size_t real_size = FRAMES * sizeof(Dtype*);

const size_t size_size = FRAMES * sizeof(size_t);

const size_t bool_size = FRAMES * sizeof(size_t);

Dtype** master_gpu_grds;

size_t* master_gpu_offs;

CUDA_CHECK(cudaMalloc((void** ) &master_gpu_grds, real_size));

CUDA_CHECK(cudaMalloc((void** ) &master_gpu_offs, size_size));

Dtype** master_cpu_grds;

size_t* master_cpu_offs;

CUDA_CHECK(cudaMallocHost((void** ) &master_cpu_grds, real_size));

CUDA_CHECK(cudaMallocHost((void** ) &master_cpu_offs, size_size));

Dtype** worker_gpu_pos;

size_t* worker_gpu_offs;

Dtype** worker_gpu_grds;

uint8_t* worker_gpu_gets;

CUDA_CHECK(cudaMalloc((void** ) &worker_gpu_pos, real_size));

CUDA_CHECK(cudaMalloc((void** ) &worker_gpu_offs, size_size));

CUDA_CHECK(cudaMalloc((void** ) &worker_gpu_grds, real_size));

CUDA_CHECK(cudaMalloc((void** ) &worker_gpu_gets, bool_size));

Dtype** worker_cpu_pos;

size_t* worker_cpu_offs;

Dtype** worker_cpu_grds;

uint8_t* worker_cpu_gets;

CUDA_CHECK(cudaMallocHost((void** ) &worker_cpu_pos, real_size));

CUDA_CHECK(cudaMallocHost((void** ) &worker_cpu_offs, size_size));

CUDA_CHECK(cudaMallocHost((void** ) &worker_cpu_grds, real_size));

CUDA_CHECK(cudaMallocHost((void** ) &worker_cpu_gets, bool_size));

int master_batch_start = 0;

int master_batch_count = 0;

int worker_batch_start = 0;

int worker_batch_count = 0;

const int batch = 128; // TODO bench

while (!this->must_stop()) {

ucast.poll();

mcast.poll();

// Receive gradients for chunks for which we are master

{

while (ucast.can_recv()) {

uint8_t* buf;

uint32_t chunk;

int id = ucast.recv(buf, chunk);

Dtype* grd = (Dtype*) buf;

CHECK(this->chunk_master(chunk) == this->rank_);

size_t off = ((size_t) chunk) * IBSync<Dtype>::CHUNK;

int index = master_queue.back_;

master_ids_[index] = id;

master_cpu_grds[index] = grd;

master_cpu_offs[index] = off;

master_queue.push();

master_batch_count++;

}

// Add gradients to our weights

if (master_batch_count >= batch) {

CUDA_CHECK(

cudaMemcpyAsync(master_gpu_grds, master_cpu_grds, real_size,

cudaMemcpyHostToDevice, master_stream.stream()));

CUDA_CHECK(

cudaMemcpyAsync(master_gpu_offs, master_cpu_offs, size_size,

cudaMemcpyHostToDevice, master_stream.stream()));

sync_master_kernel<Dtype>(gpu_, master_gpu_grds, master_gpu_offs,

master_batch_start, master_batch_count,

master_stream.stream(), IBSync<Dtype>::CHUNK);

master_events.record(master_batch_count);

master_batch_start = master_queue.back_;

master_batch_count = 0;

}

}

// Start receiving again once kernels are done with buffers

for (;;) {

int batch;

if (!master_events.query(batch)) {

break;

}

for (int i = 0; i < batch; ++i) {

int index = master_queue.front_;

master_queue.pop();

ucast.recv_done(master_ids_[index]);

}

}

// Send absolute positions for chunks for which we are master

while (mcast.can_send()) {

uint8_t* buf;

int id = mcast.send_init(buf); // buf ignored

size_t off = (size_t) this->chunk_ * IBSync<Dtype>::CHUNK;

buf = (uint8_t*) (gpu_ + off);

CHECK(id >= 0 && id < FRAMES);

mcast.send(id, this->mcast_addr_, buf, this->chunk_);

if (++this->chunk_ == this->own_until_) {

this->chunk_ = this->own_start_;

this->cycles_++;

}

}

// Receive absolute positions for other chunks

{

while (mcast.can_recv()) {

Dtype* pos;

uint32_t chunk;

int recv_id, send_id;

size_t off;

{

uint8_t* buf;

recv_id = mcast.recv(buf, chunk);

pos = (Dtype*) buf;

off = ((size_t) chunk) * IBSync<Dtype>::CHUNK;

}

// Send back the gradients if frame is available

Dtype* grd = NULL;

if (ucast.can_send()) {

uint8_t* buf;

send_id = ucast.send_init(buf);

grd = (Dtype*) buf;

}

int index = worker_queue.back_;

worker_items[index].recv_id = recv_id;

worker_items[index].send_id = send_id;

worker_items[index].grd = grd;

worker_items[index].chunk = chunk;

worker_cpu_pos[index] = pos;

worker_cpu_offs[index] = off;

worker_cpu_grds[index] = grd;

worker_cpu_gets[index] = grd != NULL;

worker_queue.push();

worker_batch_count++;

}

if (worker_batch_count >= batch) {

CUDA_CHECK(

cudaMemcpyAsync(worker_gpu_pos, worker_cpu_pos, real_size,

cudaMemcpyHostToDevice, worker_stream.stream()));

CUDA_CHECK(

cudaMemcpyAsync(worker_gpu_offs, worker_cpu_offs, size_size,

cudaMemcpyHostToDevice, worker_stream.stream()));

CUDA_CHECK(

cudaMemcpyAsync(worker_gpu_grds, worker_cpu_grds, real_size,

cudaMemcpyHostToDevice, worker_stream.stream()));

CUDA_CHECK(

cudaMemcpyAsync(worker_gpu_gets, worker_cpu_gets, bool_size,

cudaMemcpyHostToDevice, worker_stream.stream()));

sync_worker_kernel<Dtype>(gpu_, gpu_last_, worker_gpu_pos,

worker_gpu_offs, worker_gpu_grds,

worker_gpu_gets, worker_batch_start,

worker_batch_count, worker_stream.stream(),

IBSync<Dtype>::CHUNK);

worker_events.record(worker_batch_count);

worker_batch_start = worker_queue.back_;

worker_batch_count = 0;

}

}

for (;;) {

int batch;

if (!worker_events.query(batch)) {

break;

}

for (int i = 0; i < batch; ++i) {

int index = worker_queue.front_;

worker_queue.pop();

int recv_id = worker_items[index].recv_id;

int send_id = worker_items[index].send_id;

Dtype* grd = worker_items[index].grd;

uint32_t chunk = worker_items[index].chunk;

mcast.recv_done(recv_id);

if (grd) {

int master = this->chunk_master(chunk);

CHECK(master != this->rank_);

ib_addr& a = this->ucast_addrs_[master];

ucast.send(send_id, a, (uint8_t*) grd, chunk);

}

if (this->remaining_ > 0 && !this->received_[chunk]) {

this->received_[chunk] = true;

this->remaining_--;

}

}

}

}

}

INSTANTIATE_CLASS(IBSync);

INSTANTIATE_CLASS(CPUIBSync);

INSTANTIATE_CLASS(GPUIBSync);

#endif

#ifdef __linux__

static uint8_t* parse_mac(const char* str) {

uint8_t* bytes = (uint8_t*) malloc(ETH_ALEN);

for (int i = 0; i < ETH_ALEN; ++i) {

int value;

sscanf(str + 2 * i, "%02x", &value);

bytes[i] = value;

}

return bytes;

}

static vector<uint8_t*> parse_macs(const vector<string>& macs) {

}

static string adapter(const uint8_t* mac) {