// this will be invoked by nvcc and compile GPU version

#include <dmlc/logging.h>

#include "./ndarray_function.h"

#include "./ndarray_function-inl.h"

namespace mxnet {

namespace ndarray {

template<>

void Copy<cpu, gpu>(const TBlob &from, TBlob *to,

Context from_ctx, Context to_ctx,

RunContext ctx) {

mshadow::Copy(to->FlatTo2D<gpu, real_t>(),

from.FlatTo2D<cpu, real_t>(),

static_cast<mshadow::Stream<gpu>*>(ctx.stream));

}

template<>

void Copy<gpu, cpu>(const TBlob &from, TBlob *to,

Context from_ctx, Context to_ctx,

RunContext ctx) {

mshadow::Copy(to->FlatTo2D<cpu, real_t>(),

from.FlatTo2D<gpu, real_t>(),

static_cast<mshadow::Stream<gpu>*>(ctx.stream));

}

template<>

void Copy<gpu, gpu>(const TBlob &from, TBlob *to,

Context from_ctx, Context to_ctx,

RunContext ctx) {

if (from_ctx.dev_id == to_ctx.dev_id) {

mshadow::Copy(to->FlatTo2D<gpu, real_t>(),

from.FlatTo2D<gpu, real_t>(),

static_cast<mshadow::Stream<gpu>*>(ctx.stream));

} else {

CHECK(from.CheckContiguous() && to->CheckContiguous())

<< "copy across only support continugous memory";

mshadow::Stream<gpu> *s = static_cast<mshadow::Stream<gpu>*>(ctx.stream);

CHECK(s != NULL) << "need stream in GPU context";

cudaMemcpyPeerAsync(to->dptr_,

to_ctx.dev_id,

from.dptr_,

from_ctx.dev_id,

from.shape_.Size() * sizeof(real_t),

s->stream_);

}

}

} // namespace ndarray

} // namespace mxnet