Google Git
Sign in
apache / mxnet / HEAD / . / src / ndarray / ndarray.cc
blob: 902880fb1d52d9f18d3a5e5ca6b994686567b186 [file] [log] [blame]
/*
* 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.
*/
/*!
* \file ndarray.cc
* \brief ndarry module of mxnet
*/
#include <dmlc/io.h>
#include <dmlc/logging.h>
#include <dmlc/memory_io.h>
#include <dmlc/registry.h>
#include <mshadow/tensor.h>
#include <mxnet/base.h>
#include <mxnet/imperative.h>
#include <mxnet/ndarray.h>
#include <mxnet/resource.h>
#include "../common/utils.h"
#include "../operator/nn/dnnl/dnnl_base-inl.h"
#include "../operator/tensor/init_op.h"
#include "../operator/tensor/matrix_op-inl.h"
#include "../profiler/storage_profiler.h"
#include "./ndarray_function.h"
#if MXNET_USE_ONEDNN == 1
#include <dnnl.hpp>
#endif
#if MXNET_USE_OPENCV
#include <opencv2/opencv.hpp>
#endif // MXNET_USE_OPENCV
namespace dmlc {
DMLC_REGISTRY_ENABLE(::mxnet::NDArrayFunctionReg);
} // namespace dmlc
namespace mxnet {
void NDArray::ReInit(const NDArrayStorageType stype,
const mxnet::TShape& shape,
Context ctx,
int dtype,
bool delay_alloc,
const std::vector<int>* pAux_types,
const mxnet::ShapeVector* pAux_shapes,
const mxnet::TShape* pStorage_shapes) {
Init(stype, shape, dtype);
if (stype != kDefaultStorage) {
const auto sparseStorage = stype == kRowSparseStorage;
if (!sparseStorage && stype != kCSRStorage)
LOG(FATAL) << "Unknown storage type " << stype;
const auto& aux_types = (pAux_types && pAux_types->size()) ?
*pAux_types :
std::vector<int>(sparseStorage ? 1 : 2, mshadow::kInt64);
const auto& aux_shapes = (pAux_shapes && pAux_shapes->size()) ?
*pAux_shapes :
ShapeVector(sparseStorage ? 1 : 2, TShape(mshadow::Shape1(0)));
mxnet::TShape storage_shape;
if (!pStorage_shapes || !pStorage_shapes->Size()) {
if (sparseStorage) {
storage_shape = shape;
storage_shape[0] = aux_shapes[rowsparse::kIdx][0];
} else {
storage_shape = aux_shapes[csr::kIdx];
}
} else {
storage_shape = *pStorage_shapes;
}
ptr_ = std::make_shared<Chunk>(
stype, storage_shape, ctx, delay_alloc, dtype, aux_types, aux_shapes);
} else {
ptr_ = std::make_shared<Chunk>(shape, ctx, delay_alloc, dtype);
}
}
void NDArray::AssignStorageInfo(const std::string& profiler_scope, const std::string& name) {
if (is_none()) {
return;
}
ptr_->shandle.profiler_scope = profiler_scope;
ptr_->shandle.name = name;
#if MXNET_USE_CUDA
profiler::GpuDeviceStorageProfiler::Get()->UpdateStorageInfo(ptr_->shandle);
#endif // MXNET_USE_CUDA
for (Storage::Handle& aux_handle : ptr_->aux_handles) {
aux_handle.profiler_scope = profiler_scope;
aux_handle.name = name + "_aux_data";
#if MXNET_USE_CUDA
profiler::GpuDeviceStorageProfiler::Get()->UpdateStorageInfo(aux_handle);
#endif // MXNET_USE_CUDA
}
}
void NDArray::SetShapeFromChunk() const {
if (Imperative::Get()->is_np_shape() ||
!(ptr_->storage_shape.ndim() == 1 && ptr_->storage_shape[0] == 0)) {
shape_ = ptr_->storage_shape;
}
}
struct ChunkMem {
Storage::Handle h;
std::vector<Storage::Handle> aux_h;
#if MXNET_USE_ONEDNN == 1
std::shared_ptr<DNNLMemory> mem;
#endif
};
NDArray::Chunk::~Chunk() {
bool skip_free = static_data || delay_alloc;
ChunkMem mem;
mem.h = this->shandle;
mem.aux_h = this->aux_handles;
#if MXNET_USE_ONEDNN == 1
// We want to delete dnnl memory after deleting the variable.
mem.mem = this->dnnl_mem_;
#endif
if (auto engine = engine_ref_.lock()) {
engine->DeleteVariable(
[mem, skip_free, var = this->var](RunContext s) mutable {
#if MXNET_USE_CUDA
auto& sync_obj = var->sync_object;
Storage::SyncObj storage_sync_obj;
{
std::lock_guard<std::mutex> l(sync_obj.mutex);
for (auto& ev : sync_obj.reader_events) {
storage_sync_obj.events.push_back(ev.event);
}
if (!sync_obj.writer_event.empty()) {
auto ev = sync_obj.writer_event[0];
storage_sync_obj.events.push_back(ev.event);
}
}
mem.h.sync_obj = storage_sync_obj;
#endif
if (skip_free == false) {
#if MXNET_USE_ONEDNN == 1
if (mem.mem) {
CHECK_LE(mem.mem->GetSize(), mem.h.size);
CHECK_EQ(mem.mem->GetDataHandle(), mem.h.dptr);
}
#endif
Storage::Get()->Free(mem.h);
for (const auto& aux : mem.aux_h) {
Storage::Get()->Free(aux);
}
}
},
shandle.ctx,
var);
}
}
void NDArray::Chunk::CheckAndAllocData(const mxnet::TShape& shape, int dtype) {
CHECK_NE(aux_shapes.size(), 0) << "data is expected to be allocated after aux_data";
auto dbytes = shape.Size() * mshadow::mshadow_sizeof(dtype);
if (!features::is_enabled(features::INT64_TENSOR_SIZE)) {
CHECK_LT(shape.Size(), (int64_t{1} << 31) - 1)
<< "[CheckAndAllocData] Size of tensor you are trying to allocate is larger than "
"2^31 elements. Please build with flag USE_INT64_TENSOR_SIZE=1";
}
if (shandle.size < dbytes) {
// free storage
Storage::Get()->Free(shandle);
// init storage
shandle.size = dbytes;
Storage::Get()->Alloc(&shandle);
#if MXNET_USE_ONEDNN == 1
dnnl_mem_ = nullptr;
#endif
}
// init shape
storage_shape = shape;
// delay_alloc is only set when data storage handle is present
delay_alloc = false;
}
NDArray NDArray::grad() const {
if (Imperative::AGInfo::IsNone(*this))
return NDArray();
Imperative::AGInfo& info = Imperative::AGInfo::Get(autograd_entry_.node);
if (info.out_grads.size()) {
CHECK_EQ(info.out_grads.size(), 1);
return info.out_grads[0];
}
return NDArray();
}
nnvm::Symbol NDArray::get_autograd_symbol() const {
CHECK(!Imperative::AGInfo::IsNone(*this))
<< "NDArray is not part of a computation graph. Did you forget to turn on recording?";
nnvm::Symbol ret;
ret.outputs.emplace_back(autograd_entry_);
return ret;
}
#if MXNET_USE_ONEDNN == 1
NDArray::NDArray(const void* md_desc) : storage_type_(kDefaultStorage), autograd_entry_(nullptr) {
dnnl::memory::desc md = *static_cast<const dnnl::memory::desc*>(md_desc);
shape_ = mxnet::TShape(md.data.dims, md.data.dims + md.data.ndims);
dtype_ = get_mxnet_type(md.data.data_type);
ptr_ = std::make_shared<Chunk>(shape_, Context::CPU(), true, dtype_);
ptr_->CheckAndAlloc(md.get_size());
ptr_->dnnl_mem_ = std::make_shared<DNNLMemory>(md, ptr_->shandle.dptr);
}
NDArray::NDArray(const std::shared_ptr<dnnl::memory>& dnnl_mem)
: storage_type_(kDefaultStorage), autograd_entry_(nullptr) {
auto mem_desc = dnnl_mem->get_desc();
shape_ = mxnet::TShape(mem_desc.data.dims, mem_desc.data.dims + mem_desc.data.ndims);
dtype_ = get_mxnet_type(mem_desc.data.data_type);
ptr_ = std::make_shared<Chunk>(shape_, Context::CPU(), true, dtype_);
ptr_->shandle.dptr = dnnl_mem->get_data_handle();
ptr_->shandle.size = mem_desc.get_size();
ptr_->delay_alloc = false;
ptr_->dnnl_mem_ = std::make_shared<DNNLMemory>(dnnl_mem);
ptr_->static_data = true;
}
NDArray NDArray::DNNLDataReshape(const mxnet::TShape& shape) const {
CHECK(!is_none()) << "NDArray is not initialized";
CHECK_GE(shape_.Size(), shape.Size())
<< "NDArray.Reshape: target shape size is larger current shape";
CHECK_EQ(storage_type(), kDefaultStorage);
if (!IsDNNLData()) {
NDArray ret = this->Detach();
ret.shape_ = shape;
return ret;
} else {
NDArray ret(shape, ctx(), true, dtype());
// We shouldn't submit the reorder primitive here because submit will
// be called in operators.
dnnl_format_tag_t format = ptr_->dnnl_mem_->GetDefaultFormat();
CHECK(ptr_->IsDNNL());
dnnl::memory::desc def_desc = ptr_->dnnl_mem_->GetDesc(format);
dnnl::memory* def_mem = TmpMemMgr::Get()->Alloc(def_desc);
DNNLStream* stream = DNNLStream::Get();
std::shared_ptr<dnnl::memory> curr_mem = ptr_->dnnl_mem_->GetMem();
stream->RegisterMem(curr_mem);
std::unordered_map<int, dnnl::memory> args(
{{DNNL_ARG_FROM, *curr_mem}, {DNNL_ARG_TO, *def_mem}});
stream->RegisterPrimArgs(dnnl::reorder(*curr_mem, *def_mem), args);
// def_mem points to a memory region in the temp space. It's only valid
// inside an operator. As such, the returned NDArray can only be valid
// inside an operator and the shared point doesn't need to do anything
// when it's destroyed.
auto tmp = std::shared_ptr<dnnl::memory>(def_mem, [](dnnl::memory* mem) {});
ret.ptr_->dnnl_mem_.reset(new DNNLMemory(tmp));
ret.ptr_->shandle.dptr = def_mem->get_data_handle();
ret.ptr_->shandle.size = def_mem->get_desc().get_size();
ret.ptr_->delay_alloc = false;
ret.ptr_->static_data = true;
ret.byte_offset_ = byte_offset_;
ret.reuse_ = false;
return ret;
}
}
#endif
NDArray NDArray::Reshape(const mxnet::TShape& shape) const {
CHECK(!is_none()) << "NDArray is not initialized";
if (Imperative::Get()->is_np_shape()) {
CHECK_EQ(shape_.Size(), shape.Size())
<< "NDArray.Reshape: target shape must have the same size as "
<< "current shape.";
} else {
CHECK_GE(shape_.Size(), shape.Size())
<< "NDArray.Reshape: target shape size is larger than the current shape";
}
NDArray ret = this->Detach();
// If the shape doesn't change, we can just return it now.
if (ret.shape_ == shape)
return ret;
// Otherwise, reshape only works on the default layout.
CHECK_EQ(storage_type(), kDefaultStorage);
ret.shape_ = shape;
ret.reuse_ = false;
return ret;
}
NDArray NDArray::ReshapeWithRecord(const mxnet::TShape& shape) {
bool is_recording = Imperative::Get()->is_recording();
bool is_deferred_compute = Imperative::Get()->is_deferred_compute();
NDArray ret;
if (!is_deferred_compute) {
// The new array shares memory with this array, thus make sure this array
// has been computed already computed. (noop if this array is not deferred)
Imperative::DCInfo::Compute(*this);
ret = this->Reshape(shape);
if (!is_recording) {
return ret;
}
} else {
if (shape_is_known(this->shape())) {
// Imperative reshape only works if shape is already known.
ret = this->Reshape(shape);
} else {
// Reshape called on after dynamic shape operator.
ret = this->Detach();
}
}
if (!is_deferred_compute || shape_is_known(this->shape())) {
CHECK_EQ(shape_.Size(), shape.Size())
<< "NDArray.Reshape: target shape must have the same size as "
<< "current shape when recording with autograd "
<< "or in deferred compute mode.";
}
nnvm::NodeAttrs attrs;
std::ostringstream os;
os << shape;
if (!Imperative::Get()->is_np_shape()) {
attrs.op = nnvm::Op::Get("Reshape");
attrs.dict.insert({"shape", os.str()});
} else {
attrs.op = nnvm::Op::Get("_np_reshape");
attrs.dict.insert({"newshape", os.str()});
}
attrs.op->attr_parser(&attrs);
std::vector<NDArray*> inputs(1, this), outputs(1, &ret);
if (is_recording) {
Imperative::Get()->RecordOp(std::move(attrs), inputs, outputs);
} else if (is_deferred_compute) {
Imperative::Get()->RecordDeferredCompute(std::move(attrs), inputs, outputs);
}
return ret;
}
NDArray NDArray::Slice(index_t begin, index_t end) const {
CHECK(!is_none()) << "NDArray is empty";
CHECK_LE(begin, end) << "Invalid slicing range [" << begin << ", " << end << ")";
CHECK_GE(shape_[0], end) << "Slice end index out of range";
CHECK_EQ(storage_type(), kDefaultStorage);
NDArray ret = this->Detach();
size_t length = shape_.ProdShape(1, shape_.ndim());
MSHADOW_TYPE_SWITCH_EXT_WITH_BOOL(
ret.dtype(), DType, { ret.byte_offset_ += begin * length * sizeof(DType); });
ret.reuse_ = false;
ret.shape_[0] = end - begin;
return ret;
}
NDArray NDArray::SliceWithRecord(index_t begin, index_t end) {
bool is_recording = Imperative::Get()->is_recording();
bool is_deferred_compute = Imperative::Get()->is_deferred_compute();
NDArray ret;
if (!is_deferred_compute) {
// The new array shares memory with this array, thus make sure this array
// has been computed already computed. (noop if this array is not deferred)
Imperative::DCInfo::Compute(*this);
ret = this->Slice(begin, end);
if (!is_recording) {
return ret;
}
} else {
if (shape_is_known(this->shape())) {
// Imperative slice only works if shape is already known.
ret = this->Slice(begin, end);
} else {
// Slice called on after dynamic shape operator.
ret = this->Detach();
}
}
// fake a slice op
nnvm::NodeAttrs attrs;
attrs.op = nnvm::Op::Get("slice");
attrs.dict.insert({"begin", std::to_string(begin)});
attrs.dict.insert({"end", std::to_string(end)});
attrs.op->attr_parser(&attrs);
std::vector<NDArray*> inputs(1, this), outputs(1, &ret);
if (is_recording) {
Imperative::Get()->RecordOp(std::move(attrs), inputs, outputs);
} else if (is_deferred_compute) {
Imperative::Get()->RecordDeferredCompute(std::move(attrs), inputs, outputs);
}
return ret;
}
NDArray NDArray::At(index_t idx) const {
CHECK(storage_type() == kDefaultStorage)
<< "Storage type " << storage_type() << " doesn't support At()";
NDArray ret = this->Slice(idx, idx + 1);
if (shape_.ndim() > 1) {
return ret.Reshape(mxnet::TShape(shape_.data() + 1, shape_.data() + shape_.ndim()));
} else {
return ret;
}
}
NDArray NDArray::AtWithRecord(index_t idx) {
CHECK(storage_type() == kDefaultStorage)
<< "Storage type " << storage_type() << " doesn't support At()";
NDArray sliced = this->SliceWithRecord(idx, idx + 1);
if (shape_.ndim() > 1 || Imperative::Get()->is_np_shape()) {
// Imperative reshape with concrete shape
NDArray reshaped =
sliced.Reshape(mxnet::TShape(shape_.data() + 1, shape_.data() + shape_.ndim()));
// Record reshape with magic numbers
nnvm::NodeAttrs attrs;
std::ostringstream os;
if (!Imperative::Get()->is_np_shape()) {
os << mxnet::TShape({-3, -2}); // See ndarray.py reshape for definition of magic numbers
attrs.op = nnvm::Op::Get("Reshape");
attrs.dict.insert({"shape", os.str()});
} else {
// See NumpyXReshapeInferShape for definition of magic numbers
os << mxnet::TShape({-3, -4});
attrs.op = nnvm::Op::Get("_npx_reshape");
attrs.dict.insert({"newshape", os.str()});
}
attrs.op->attr_parser(&attrs);
std::vector<NDArray*> inputs(1, &sliced), outputs(1, &reshaped);
bool is_recording = Imperative::Get()->is_recording();
bool is_deferred_compute = Imperative::Get()->is_deferred_compute();
if (is_recording) {
Imperative::Get()->RecordOp(std::move(attrs), inputs, outputs);
} else if (is_deferred_compute) {
Imperative::Get()->RecordDeferredCompute(std::move(attrs), inputs, outputs);
}
return reshaped;
} else {
return sliced;
}
}
/*!
* \brief Return deep copy of the current ndarry's aux_data(i)
* as an NDArray of default storage type. This function blocks.
*/
NDArray NDArray::aux_ndarray(size_t i) const {
CHECK_NE(storage_type(), kDefaultStorage);
CHECK(i < ptr_->aux_shapes.size());
// create a delay_alloc default ndarray as output
NDArray ret(mxnet::TShape(), ctx(), true, aux_type(i));
ret.SyncCopyFromNDArray(*this, i);
return ret;
}
NDArray NDArray::data_ndarray() const {
NDArray ret(mxnet::TShape(), ctx(), true, dtype_);
ret.SyncCopyFromNDArray(*this);
return ret;
}
struct NDArrayDLManager {
NDArray handle; // ref NDArray
DLManagedTensor tensor;
};
DLManagedTensor* NDArray::ToDLPack() const {
CHECK(!is_none()) << "NDArray is not initialized";
NDArrayDLManager* dlmanager(new NDArrayDLManager);
dlmanager->handle = *this;
dlmanager->tensor.dl_tensor = dlmanager->handle.data().dltensor();
dlmanager->tensor.manager_ctx = dlmanager;
dlmanager->tensor.deleter = [](DLManagedTensor* dlmanager) {
delete static_cast<NDArrayDLManager*>(dlmanager->manager_ctx);
};
return &(dlmanager->tensor);
}
NDArray NDArray::FromDLPack(const DLManagedTensor* tensor, bool transient_handle) {
DLManagedTensor* tensor_copy =
transient_handle ? new DLManagedTensor(*tensor) : const_cast<DLManagedTensor*>(tensor);
auto deleter = [tensor_copy, transient_handle]() {
if (tensor_copy->deleter != nullptr) {
tensor_copy->deleter(tensor_copy);
}
if (transient_handle) {
delete tensor_copy;
}
};
return NDArray(TBlob(tensor_copy->dl_tensor), tensor_copy->dl_tensor.ctx.device_id, deleter);
}
bool NDArray::fresh_out_grad() const {
if (Imperative::AGInfo::IsNone(*this))
return false;
Imperative::AGInfo& info = Imperative::AGInfo::Get(autograd_entry_.node);
return info.fresh_out_grad;
}
void NDArray::set_fresh_out_grad(bool state) const {
CHECK(!Imperative::AGInfo::IsNone(*this))
<< "NDArray has not been marked as a variable and does not have gradient state";
Imperative::AGInfo& info = Imperative::AGInfo::Get(autograd_entry_.node);
info.fresh_out_grad = state;
}
#if MXNET_USE_ONEDNN == 1
bool NDArray::Chunk::IsDNNL() const {
if (storage_type != kDefaultStorage)
return false;
if (dnnl_mem_ == nullptr)
return false;
return dnnl_mem_->IsDNNL();
}
bool NDArray::Chunk::IsDefault() const {
if (storage_type != kDefaultStorage)
return false;
// If we don't have dnnl memory yet, we just assume it's not the default
// format.
if (dnnl_mem_ == nullptr)
return true;
return !dnnl_mem_->IsDNNL();
}
void NDArray::Chunk::Reorder2Default() {
if (dnnl_mem_ == nullptr)
return;
if (IsDefault())
return;
dnnl_format_tag_t format = dnnl_mem_->GetDefaultFormat();
dnnl::memory::desc def_desc = dnnl_mem_->GetDesc(format);
CHECK(shandle.size >= def_desc.get_size());
CheckAndAlloc(def_desc.get_size());
// oneDNN reorder can't be performed in-place
if (shandle.dptr == dnnl_mem_->GetDataHandle()) {
dnnl_mem_ptr def_mem(new dnnl::memory(def_desc, CpuEngine::Get()->get_engine()));
dnnl_mem_->ReorderTo(def_mem.get());
memcpy(shandle.dptr, def_mem->get_data_handle(), def_desc.get_size());
} else {
dnnl_mem_ptr def_mem(new dnnl::memory(def_desc, CpuEngine::Get()->get_engine(), shandle.dptr));
dnnl_mem_->ReorderTo(def_mem.get());
}
dnnl_mem_ = nullptr;
}
void NDArray::Chunk::DNNLDataReorder(const void* mem_desc) {
const dnnl::memory::desc md = *static_cast<const dnnl::memory::desc*>(mem_desc);
// If the memory already uses the specified layout, don't do anything.
if (dnnl_mem_ != nullptr && dnnl_mem_->SameFormat(md))
return;
// If the memory is default, don't do anything.
if (!mxnet::IsDNNL(md) && IsDefault())
return;
if (!mxnet::IsDNNL(md)) {
// If the specified layout is default, we should use Reorder2Default.
Reorder2Default();
return;
}
auto engine = CpuEngine::Get()->get_engine();
dnnl::stream s(engine);
std::shared_ptr<dnnl::memory> new_mem(new dnnl::memory(md, engine));
std::shared_ptr<dnnl::memory> old_mem;
if (IsDefault()) {
dnnl_format_tag_t def_format = GetDefaultFormat(md);
dnnl::memory::desc def_desc = GetDesc(md, def_format);
old_mem.reset(new dnnl::memory(def_desc, engine, shandle.dptr));
} else {
old_mem = this->dnnl_mem_->GetMem();
}
CHECK(old_mem->get_desc().data.ndims == md.data.ndims);
// This may be called in DNNL operators. We can't use DNNLStream here.
dnnl::reorder(*old_mem, *new_mem).execute(s, *old_mem, *new_mem);
CHECK(shandle.size >= md.get_size());
CheckAndAlloc(md.get_size());
// TODO(zhengda) We need to avoid memory copy here.
memcpy(shandle.dptr, new_mem->get_data_handle(), md.get_size());
dnnl_mem_.reset(new DNNLMemory(md, shandle.dptr));
}
void NDArray::Chunk::SetDNNLMem(const mxnet::TShape& shape, int dtype) {
// The shape of the array and the one of the DNNL memory may mismatch.
// For example, if the array stores parameters, the DNNL memory may store data
// in 5 dimensions while the NDArray stores data in 4 dimensions.
if (dnnl_mem_ && dnnl_mem_->GetDataHandle() == shandle.dptr &&
dnnl_mem_->SameFormat(shape, dtype)) {
return;
}
dnnl::memory::dims dims;
// These are shapes supprted by DNNL.
const int MAX_ONEDNN_DIMS = 12;
if (shape.ndim() >= 1 && shape.ndim() <= MAX_ONEDNN_DIMS) {
dims.resize(shape.ndim());
for (size_t i = 0; i < dims.size(); i++)
dims[i] = shape[i];
} else {
LOG(FATAL) << "oneDNN doesn't support " << shape.ndim() << " dimensions";
}
auto layout = static_cast<dnnl::memory::format_tag>(GetDefaultFormat(dims.size()));
dnnl::memory::desc data_md{dims, get_dnnl_type(dtype), layout};
if (shandle.dptr == nullptr) {
CHECK(delay_alloc);
CheckAndAlloc();
}
CHECK(shandle.size >= data_md.get_size());
dnnl_mem_.reset(new DNNLMemory(data_md, shandle.dptr));
}
const dnnl::memory* NDArray::GetDNNLData(const void* mem_desc) const {
const dnnl::memory::desc desc = *static_cast<const dnnl::memory::desc*>(mem_desc);
if (desc.get_size() != shape().Size() * GetTypeSize(dtype_)) {
LOG(FATAL) << "The size of NDArray doesn't match the requested oneDNN memory desc";
return nullptr;
}
const dnnl::memory* mem = GetDNNLData();
dnnl::memory::desc desc1 = mem->get_desc();
// The DNNL memory has the same format and shape as required,
// or both use the default format, we can return the DNNL memory.
if (desc1 == desc || ((!mxnet::IsDNNL(desc1)) && (!mxnet::IsDNNL(desc)))) {
return GetDNNLExact(mem, desc);
} else {
return nullptr;
}
}
const dnnl::memory* NDArray::GetDNNLDataReorder(const void* mem_desc) const {
dnnl::memory::desc new_desc = *static_cast<const dnnl::memory::desc*>(mem_desc);
CHECK(storage_type() == kDefaultStorage);
const dnnl::memory* mem = GetDNNLData();
// If the memory descriptor matches, it's easy.
DNNLStream* stream = DNNLStream::Get();
if (mem->get_desc() == new_desc) {
return GetDNNLExact(mem, new_desc);
}
dnnl::memory::desc old_desc = mem->get_desc();
// Now we need to determine if we should reorder the memory.
// If both use the default formats, we think we don't need to reorder.
if ((!mxnet::IsDNNL(old_desc)) && (!mxnet::IsDNNL(new_desc))) {
dnnl_mem_ptr ret(
new dnnl::memory(new_desc, CpuEngine::Get()->get_engine(), mem->get_data_handle()));
stream->RegisterMem(ret);
return ret.get();
} else if (same_shape(old_desc, new_desc)) {
// If they have the same shape, we can reorder data directly.
dnnl::memory* ret = TmpMemMgr::Get()->Alloc(new_desc);
std::unordered_map<int, dnnl::memory> args({{DNNL_ARG_FROM, *mem}, {DNNL_ARG_TO, *ret}});
stream->RegisterPrimArgs(dnnl::reorder(*mem, *ret), args);
return ret;
} else {
// If they have different shapes, we need to reshape the array first.
// Since this method will only be used inside an operator, we can call
// DNNLDataReshape to reshape an array.
mxnet::TShape required_shape(new_desc.data.ndims, -1);
for (int i = 0; i < new_desc.data.ndims; i++)
required_shape[i] = new_desc.data.dims[i];
NDArray reshaped = DNNLDataReshape(required_shape);
const dnnl::memory* ret = reshaped.GetDNNLData();
if (ret->get_desc() == new_desc) {
return GetDNNLExact(ret, new_desc);
} else {
dnnl::memory* ret2 = TmpMemMgr::Get()->Alloc(new_desc);
std::unordered_map<int, dnnl::memory> args({{DNNL_ARG_FROM, *ret}, {DNNL_ARG_TO, *ret2}});
stream->RegisterPrimArgs(dnnl::reorder(*ret, *ret2), args);
return ret2;
}
}
}
NDArray NDArray::Reorder2Default() const {
CHECK(storage_type() == kDefaultStorage);
if (ptr_->dnnl_mem_ == nullptr)
return *this;
if (!ptr_->dnnl_mem_->IsDNNL())
return *this;
// create new ndarray from dnnl layout
dnnl::memory::desc from_desc = ptr_->dnnl_mem_->GetDesc();
mxnet::TShape tshape(from_desc.data.ndims, -1);
for (int i = 0; i < from_desc.data.ndims; i++)
tshape[i] = from_desc.data.dims[i];
NDArray ret(tshape, ctx(), false, dtype());
dnnl_format_tag_t format = ptr_->dnnl_mem_->GetDefaultFormat();
dnnl::memory::desc def_desc = ptr_->dnnl_mem_->GetDesc(format);
CHECK(ret.ptr_->shandle.size >= def_desc.get_size());
dnnl::memory def_mem(def_desc, CpuEngine::Get()->get_engine(), ret.ptr_->shandle.dptr);
ptr_->dnnl_mem_->ReorderTo(&def_mem);
// reshape as needed
ret.shape_ = shape_;
ret.byte_offset_ = byte_offset_;
ret.reuse_ = false;
return ret;
}
void NDArray::SelfReorder2Default() {
if (!IsDNNLData())
return;
CHECK(storage_type() == kDefaultStorage);
const auto dnnl_mem = ptr_->dnnl_mem_;
if (dnnl_mem == nullptr || !dnnl_mem->IsDNNL())
return;
// create new ndarray from dnnl layout
dnnl::memory::desc from_desc = dnnl_mem->GetDesc();
mxnet::TShape tshape(from_desc.data.ndims, -1);
for (int i = 0; i < from_desc.data.ndims; i++)
tshape[i] = from_desc.data.dims[i];
const auto saved_shape = shape_;
const auto saved_byte_offset = byte_offset_;
this->ReInit(kDefaultStorage, tshape, ctx(), dtype(), false);
dnnl_format_tag_t format = dnnl_mem->GetDefaultFormat();
dnnl::memory::desc def_desc = dnnl_mem->GetDesc(format);
CHECK(ptr_->shandle.size >= def_desc.get_size());
dnnl::memory def_mem(def_desc, CpuEngine::Get()->get_engine(), ptr_->shandle.dptr);
dnnl_mem->ReorderTo(&def_mem);
// reshape as needed
shape_ = saved_shape;
byte_offset_ = saved_byte_offset;
reuse_ = false;
}
void NDArray::Reorder2DefaultAsync() const {
std::vector<Engine::VarHandle> const_vars;
std::vector<Engine::VarHandle> mutable_vars(1, this->var());
NDArray tmp = *this;
Engine::Get()->PushAsync(
[tmp](RunContext ctx,
Engine::CallbackOnStart on_start,
Engine::CallbackOnComplete on_complete) {
on_start();
tmp.ptr_->Reorder2Default();
on_complete();
},
ctx(),
const_vars,
mutable_vars,
FnProperty::kNormal,
0,
"Reorder2Default");
}
// now just support bf16->fp32
NDArray NDArray::Reorder2DefaultFloatFormat() const {
CHECK(storage_type() == kDefaultStorage && IsView() == false);
if (dtype() != mshadow::kBfloat16) {
return Reorder2Default();
}
NDArray ret(shape(), ctx(), false, mshadow::DataType<float>::kFlag);
auto src_mem = GetDNNLData();
auto dst_mem = ret.GetDNNLData();
ReorderTo(src_mem, dst_mem);
return ret;
}
void NDArray::DNNLDataReorderAsync(const void* mem_desc) const {
dnnl::memory::desc desc = *static_cast<const dnnl::memory::desc*>(mem_desc);
std::vector<Engine::VarHandle> const_vars;
std::vector<Engine::VarHandle> mutable_vars(1, this->var());
NDArray tmp = *this;
const auto version = this->version();
Engine::Get()->PushAsync(
[tmp, version, desc](RunContext ctx,
Engine::CallbackOnStart on_start,
Engine::CallbackOnComplete on_complete) {
on_start();
// MXNet will try to reuse NDArray from memory planning, so we need to ensure
// the NDArray is still holding the original trunk data.
if (tmp.version() == version) {
tmp.ptr_->DNNLDataReorder(&desc);
}
on_complete();
},
ctx(),
const_vars,
mutable_vars,
FnProperty::kNormal,
0,
"Reorder");
}
const dnnl::memory* NDArray::GetDNNLData() const {
CHECK(storage_type() == kDefaultStorage);
const auto is_view = IsView();
if (IsDNNLData()) {
// If this array uses DNNL layout, we have to make sure it's not a view.
// Otherwise, we'll have to change the layout inside the array.
CHECK(!is_view);
DNNLStream::Get()->RegisterMem(ptr_->dnnl_mem_->GetMem());
// If this array uses DNNL format, we should return now. Otherwise,
// SetDNNLMem may mess up dnnl_mem_.
return ptr_->dnnl_mem_->GetRaw();
}
CheckAndAlloc();
if (is_view) {
// If this is a view, we can't create a DNNL memory for the chunk
// because we don't have the complete data type and shape information for
// the chunk.
void* off_addr = static_cast<char*>(ptr_->shandle.dptr) + byte_offset_;
// Create the primitive desc for the new dnnl memory.
dnnl::memory::dims dims(shape().ndim());
for (size_t i = 0; i < dims.size(); i++)
dims[i] = shape()[i];
const auto cpp_format = static_cast<dnnl::memory::format_tag>(GetDefaultFormat(shape().ndim()));
dnnl::memory::desc data_md(dims, get_dnnl_type(dtype_), cpp_format);
std::shared_ptr<dnnl::memory> ret(
new dnnl::memory(data_md, CpuEngine::Get()->get_engine(), off_addr));
DNNLStream::Get()->RegisterMem(ret);
return ret.get();
}
// If this isn't a view, we can create a DNNL memory and store it in the chunk
ptr_->SetDNNLMem(shape_, dtype_);
DNNLStream::Get()->RegisterMem(ptr_->dnnl_mem_->GetMem());
return ptr_->dnnl_mem_->GetRaw();
}
void NDArray::InvalidateDNNLData() {
// Removing dnnl_mem_ means the NDArray will store data in the default format.
if (ptr_->dnnl_mem_ && ptr_->dnnl_mem_->IsDNNL())
ptr_->dnnl_mem_ = nullptr;
}
void NDArray::CopyFrom(const dnnl::memory& mem) {
CHECK(ptr_ != nullptr) << "The NDArray hasn't been initialized";
if (ptr_->dnnl_mem_ && ptr_->dnnl_mem_->GetRaw() == &mem)
return;
CHECK(mem.get_desc().get_size() == shape().Size() * GetTypeSize(dtype_))
<< "The size of NDArray doesn't match the requested oneDNN memory desc";
// If this array uses DNNL layout, we have to make sure it's not a view.
// Otherwise, we'll have to change the layout inside the array.
if (IsDNNLData() && IsView())
ptr_->Reorder2Default();
const dnnl::memory* this_mem = GetDNNLData();
DNNLMemoryCopy(mem, this_mem);
}
dnnl::memory* NDArray::CreateDNNLData(const void* mem_desc) {
dnnl::memory::desc desc = *static_cast<const dnnl::memory::desc*>(mem_desc);
if (desc.get_size() != shape().Size() * GetTypeSize(dtype_)) {
LOG(FATAL) << "The size of NDArray doesn't match the requested oneDNN memory desc. "
<< "oneDNN memory requests for " << desc.get_size() << " bytes, but got "
<< shape().Size() * GetTypeSize(dtype_) << " bytes from NDArray";
return nullptr;
}
bool isDefaultFormat = IsDefaultFormat(desc);
if (isDefaultFormat && !IsView()) {
ptr_->SetDNNLMem(shape_, dtype_);
DNNLStream::Get()->RegisterMem(ptr_->dnnl_mem_->GetMem());
return GetDNNLExact(ptr_->dnnl_mem_->GetRaw(), desc);
} else if (isDefaultFormat) {
ptr_->CheckAndAlloc();
CHECK(ptr_->shandle.dptr);
// When this is a view and a user wants the default layout, we can simply
// create a new dnnl memory that points to the right memory.
std::shared_ptr<dnnl::memory> mem(
new dnnl::memory(desc,
CpuEngine::Get()->get_engine(),
static_cast<char*>(ptr_->shandle.dptr) + byte_offset_));
DNNLStream::Get()->RegisterMem(mem);
return mem.get();
} else if (IsView()) {
// If this is a view and a user wants to write data to it with special
// a DNNL format, we should reorder the data in the array and return NULL.
// In this way, the user will create a new NDArray for the special format
// and copy data back.
ptr_->Reorder2Default();
return nullptr;
}
if (ptr_->dnnl_mem_)
CHECK(ptr_->dnnl_mem_->GetDataHandle() == ptr_->shandle.dptr);
if (ptr_->dnnl_mem_ && ptr_->dnnl_mem_->GetDesc() == desc) {
DNNLStream::Get()->RegisterMem(ptr_->dnnl_mem_->GetMem());
return GetDNNLExact(ptr_->dnnl_mem_->GetRaw(), desc);
}
CHECK(ptr_->shandle.size >= desc.get_size());
ptr_->CheckAndAlloc(desc.get_size());
ptr_->dnnl_mem_.reset(new DNNLMemory(desc, ptr_->shandle.dptr));
DNNLStream::Get()->RegisterMem(ptr_->dnnl_mem_->GetMem());
return ptr_->dnnl_mem_->GetRaw();
}
void NDArray::UpdateDNNLMemDesc(const void* mem_desc) {
dnnl::memory::desc desc = *static_cast<const dnnl::memory::desc*>(mem_desc);
auto new_desc = desc;
auto this_dtype = get_dnnl_type(dtype());
new_desc.data.data_type = static_cast<dnnl_data_type_t>(this_dtype);
ptr_->dnnl_mem_.reset(new DNNLMemory(new_desc, ptr_->shandle.dptr));
DNNLStream::Get()->RegisterMem(ptr_->dnnl_mem_->GetMem());
}
#endif
void NDArray::SetTBlob() const {
CHECK(ptr_ != nullptr);
mxnet::TShape shape = shape_;
char* dptr = static_cast<char*>(ptr_->shandle.dptr);
auto stype = storage_type();
if (stype == kDefaultStorage) {
#if MXNET_USE_ONEDNN == 1
CHECK(!IsDNNLData()) << "We can't generate TBlob for oneDNN data. "
<< "Please use Reorder2Default() to generate a new NDArray first";
#endif
dptr += byte_offset_;
} else if (stype == kCSRStorage || stype == kRowSparseStorage) {
CHECK_EQ(byte_offset_, 0);
shape = storage_shape();
} else {
LOG(FATAL) << "unknown storage type " << stype;
}
tblob_.dptr_ = dptr;
tblob_.shape_ = shape;
tblob_.type_flag_ = dtype_;
tblob_.SetDLTensor(ptr_->shandle.ctx.dev_mask(), ptr_->shandle.ctx.dev_id);
}
/*!
* \brief run a ternary operation
* \param lhs left operand
* \param mhs middle operand
* \param rhs right operand
* \param out the output ndarray
*/
template <typename OP>
void TernaryOp(const NDArray& lhs, const NDArray& mhs, const NDArray& rhs, NDArray* out) {
// no check if all of them are on cpu
if (lhs.ctx().dev_mask() != cpu::kDevMask || mhs.ctx().dev_mask() != cpu::kDevMask ||
rhs.ctx().dev_mask() != cpu::kDevMask) {
CHECK((lhs.ctx() == mhs.ctx()) && (mhs.ctx() == rhs.ctx())) << "operands context mismatch";
}
// if out is none, allocate space
if (out->is_none()) {
*out = NDArray(OP::GetShape(lhs.shape(), mhs.shape(), rhs.shape()), lhs.ctx(), true);
} else {
// no check if both of them are on cpu
if (lhs.ctx().dev_mask() != cpu::kDevMask || out->ctx().dev_mask() != cpu::kDevMask) {
CHECK(out->ctx() == lhs.ctx()) << "target context mismatch";
}
CHECK(out->shape() == OP::GetShape(lhs.shape(), mhs.shape(), rhs.shape()))
<< "target shape mismatch";
}
// important: callback must always capture by value
NDArray ret = *out;
// get the const variables
std::vector<Engine::VarHandle> const_vars;
if (lhs.var() != ret.var())
const_vars.push_back(lhs.var());
if (mhs.var() != ret.var())
const_vars.push_back(mhs.var());
if (rhs.var() != ret.var())
const_vars.push_back(rhs.var());
// redirect everything to mshadow operations
switch (lhs.ctx().dev_mask()) {
case cpu::kDevMask: {
Engine::Get()->PushSync(
[lhs, mhs, rhs, ret](RunContext ctx) {
TBlob tmp = ret.data();
ndarray::Eval<cpu, OP>(lhs.data(), mhs.data(), rhs.data(), &tmp, ctx);
},
lhs.ctx(),
const_vars,
{ret.var()},
FnProperty::kNormal,
0,
PROFILER_MESSAGE_FUNCNAME);
break;
}
#if MXNET_USE_CUDA
case gpu::kDevMask: {
Engine::Get()->PushSync(
[lhs, mhs, rhs, ret](RunContext ctx) {
TBlob tmp = ret.data();
ndarray::Eval<gpu, OP>(lhs.data(), mhs.data(), rhs.data(), &tmp, ctx);
},
lhs.ctx(),
const_vars,
{ret.var()},
FnProperty::kNormal,
0,
PROFILER_MESSAGE_FUNCNAME);
break;
}
#endif
default:
LOG(FATAL) << MXNET_GPU_NOT_ENABLED_ERROR;
}
}
/*!
* \brief Performs some preparation required to apply binary operators.
* Checks context and shape of ndarrays, allocates space for output
* and prepares const variables for engine
* \param lhs left operand
* \param rhs right operand
* \param out the output ndarray
* \param binary_op the real operation
*/
template <typename OP>
std::vector<Engine::VarHandle> BinaryOpPrepare(const NDArray& lhs,
const NDArray& rhs,
NDArray* out) {
// no check if both of them are on cpu
if (lhs.ctx().dev_mask() != cpu::kDevMask || rhs.ctx().dev_mask() != cpu::kDevMask) {
CHECK(lhs.ctx() == rhs.ctx()) << "operands context mismatch";
}
// if out is none, allocate space
if (out->is_none()) {
*out = NDArray(OP::GetShape(lhs.shape(), rhs.shape()), lhs.ctx(), true, lhs.dtype());
} else {
// no check if both of them are on cpu
if (lhs.ctx().dev_mask() != cpu::kDevMask || out->ctx().dev_mask() != cpu::kDevMask) {
CHECK(out->ctx() == lhs.ctx()) << "target context mismatch";
}
CHECK(out->shape() == OP::GetShape(lhs.shape(), rhs.shape())) << "target shape mismatch";
}
std::vector<Engine::VarHandle> const_vars;
// prepare const variables for engine
if (lhs.var() != out->var())
const_vars.push_back(lhs.var());
if (rhs.var() != out->var())
const_vars.push_back(rhs.var());
return const_vars;
}
/*!
* \brief run a binary operation using the kernel launch method
* \param lhs left operand
* \param rhs right operand
* \param out the output ndarray
* \param binary_op the real operation
*/
template <typename OP>
void BinaryOpKernel(const NDArray& lhs, const NDArray& rhs, NDArray* out) {
std::vector<Engine::VarHandle> const_vars = BinaryOpPrepare<OP>(lhs, rhs, out);
// important: callback must always capture by value
NDArray ret = *out;
switch (lhs.ctx().dev_mask()) {
case cpu::kDevMask: {
Engine::Get()->PushSync(
[lhs, rhs, ret](RunContext ctx) {
TBlob tmp = ret.data();
mshadow::Stream<cpu>* s = ctx.get_stream<cpu>();
ndarray::BinaryOpKernelImpl<OP>(s, lhs.data(), rhs.data(), &tmp);
},
lhs.ctx(),
const_vars,
{ret.var()},
FnProperty::kNormal,
0,
PROFILER_MESSAGE_FUNCNAME);
break;
}
#if MXNET_USE_CUDA
case gpu::kDevMask: {
Engine::Get()->PushSync(
[lhs, rhs, ret](RunContext ctx) {
TBlob tmp = ret.data();
mshadow::Stream<gpu>* s = ctx.get_stream<gpu>();
ndarray::BinaryOpKernelImpl<OP>(s, lhs.data(), rhs.data(), &tmp);
},
lhs.ctx(),
const_vars,
{ret.var()},
FnProperty::kNormal,
0,
PROFILER_MESSAGE_FUNCNAME);
break;
}
#endif
default:
LOG(FATAL) << MXNET_GPU_NOT_ENABLED_ERROR;
}
}
/*!
* \brief run a binary operation using mshadow operations
* \param lhs left operand
* \param rhs right operand
* \param out the output ndarray
* \param binary_op the real operation
*/
template <typename OP>
void BinaryOp(const NDArray& lhs, const NDArray& rhs, NDArray* out) {
std::vector<Engine::VarHandle> const_vars = BinaryOpPrepare<OP>(lhs, rhs, out);
// important: callback must always capture by value
NDArray ret = *out;
// redirect everything to mshadow operations
switch (lhs.ctx().dev_mask()) {
case cpu::kDevMask: {
Engine::Get()->PushSync(
[lhs, rhs, ret](RunContext ctx) {
TBlob tmp = ret.data();
ndarray::Eval<cpu, OP>(lhs.data(), rhs.data(), &tmp, ctx);
},
lhs.ctx(),
const_vars,
{ret.var()},
FnProperty::kNormal,
0,
PROFILER_MESSAGE_FUNCNAME);
break;
}
#if MXNET_USE_CUDA
case gpu::kDevMask: {
Engine::Get()->PushSync(
[lhs, rhs, ret](RunContext ctx) {
TBlob tmp = ret.data();
ndarray::Eval<gpu, OP>(lhs.data(), rhs.data(), &tmp, ctx);
},
lhs.ctx(),
const_vars,
{ret.var()},
FnProperty::kNormal,
0,
PROFILER_MESSAGE_FUNCNAME);
break;
}
#endif
default:
LOG(FATAL) << MXNET_GPU_NOT_ENABLED_ERROR;
}
}
void SetValueOp(const real_t& rhs, NDArray* out) {
CHECK_NE(out->is_none(), true) << "Set value target must not be empty";
// important: callback must always capture by value
NDArray ret = *out;
const NDArrayStorageType stype = ret.storage_type();
Engine::Get()->PushSync(
[rhs, ret, stype](RunContext ctx) {
TBlob tmp = ret.data();
switch (ret.ctx().dev_mask()) {
case cpu::kDevMask: {
if (stype == kDefaultStorage) {
ndarray::Eval<cpu>(rhs, &tmp, ctx);
} else {
ndarray::Eval(ctx.get_stream<cpu>(), rhs, ret);
}
break;
}
#if MXNET_USE_CUDA
case gpu::kDevMask: {
if (stype == kDefaultStorage) {
ndarray::Eval<gpu>(rhs, &tmp, ctx);
} else {
ndarray::Eval(ctx.get_stream<gpu>(), rhs, ret);
}
break;
}
#endif
default:
LOG(FATAL) << MXNET_GPU_NOT_ENABLED_ERROR;
}
},
ret.ctx(),
{},
{ret.var()},
FnProperty::kNormal,
0,
PROFILER_MESSAGE_FUNCNAME);
}
/*!
* \brief run a binary operation
* \param lhs left operand
* \param rhs right operand
* \param out the output ndarray
* \param binary_op the real
*/
template <typename OP, bool reverse>
void ScalarOp(const NDArray& lhs, const real_t& rhs, NDArray* out) {
if (out->is_none()) {
*out = NDArray(lhs.shape(), lhs.ctx(), true, lhs.dtype());
} else {
CHECK(out->ctx() == lhs.ctx()) << "target context mismatch";
CHECK(out->shape() == lhs.shape()) << "target shape mismatch";
}
// important: callback must always capture by value
NDArray ret = *out;
// get the const variables
std::vector<Engine::VarHandle> const_vars;
if (lhs.var() != ret.var())
const_vars.push_back(lhs.var());
// redirect everything to mshadow operations
switch (lhs.ctx().dev_mask()) {
case cpu::kDevMask: {
Engine::Get()->PushSync(
[lhs, rhs, ret](RunContext ctx) {
TBlob tmp = ret.data();
ndarray::Eval<cpu, OP, reverse>(lhs.data(), rhs, &tmp, ctx);
},
lhs.ctx(),
const_vars,
{ret.var()},
FnProperty::kNormal,
0,
PROFILER_MESSAGE_FUNCNAME);
break;
}
#if MXNET_USE_CUDA
case gpu::kDevMask: {
Engine::Get()->PushSync(
[lhs, rhs, ret](RunContext ctx) {
TBlob tmp = ret.data();
ndarray::Eval<gpu, OP, reverse>(lhs.data(), rhs, &tmp, ctx);
},
lhs.ctx(),
const_vars,
{ret.var()},
FnProperty::kNormal,
0,
PROFILER_MESSAGE_FUNCNAME);
break;
}
#endif
default:
LOG(FATAL) << MXNET_GPU_NOT_ENABLED_ERROR;
}
}
size_t num_aux_data(NDArrayStorageType stype) {
size_t num = 0;
switch (stype) {
case kDefaultStorage:
num = 0;
break;
case kCSRStorage:
num = 2;
break;
case kRowSparseStorage:
num = 1;
break;
default:
LOG(FATAL) << "Unknown storage type" << stype;
break;
}
return num;
}
// Make a copy of a CSR NDArray
template <typename from_xpu, typename to_xpu>
inline void CopyFromToCsrImpl(const NDArray& from, const NDArray& to, RunContext ctx) {
using namespace mshadow;
CHECK_EQ(from.storage_type(), to.storage_type()) << "Copying with different storage type";
// if source storage is not initialized, fill destination with zeros
auto s = ctx.get_stream<to_xpu>();
if (!from.storage_initialized()) {
op::FillZerosCsrImpl(s, to);
return;
}
// Allocate storage
to.CheckAndAllocAuxData(csr::kIndPtr, from.aux_shape(csr::kIndPtr));
to.CheckAndAllocAuxData(csr::kIdx, from.aux_shape(csr::kIdx));
to.CheckAndAllocData(from.aux_shape(csr::kIdx));
TBlob val = to.data();
TBlob indptr = to.aux_data(csr::kIndPtr);
TBlob idx = to.aux_data(csr::kIdx);
ndarray::Copy<from_xpu, to_xpu>(from.data(), &val, from.ctx(), to.ctx(), ctx);
ndarray::Copy<from_xpu, to_xpu>(from.aux_data(csr::kIndPtr), &indptr, from.ctx(), to.ctx(), ctx);
ndarray::Copy<from_xpu, to_xpu>(from.aux_data(csr::kIdx), &idx, from.ctx(), to.ctx(), ctx);
}
// Make a copy of a row-sparse NDArray
template <typename from_xpu, typename to_xpu>
inline void CopyFromToRspImpl(const NDArray& from, const NDArray& to, RunContext ctx) {
using namespace mshadow;
CHECK_EQ(from.storage_type(), to.storage_type()) << "Copying with different storage type";
// if source is zeros, fill destination with zeros, too
auto s = ctx.get_stream<to_xpu>();
if (!from.storage_initialized()) {
op::FillZerosRspImpl(s, to);
return;
}
const auto& aux_shape = from.aux_shape(rowsparse::kIdx);
to.CheckAndAlloc({aux_shape});
TBlob val = to.data();
TBlob idx = to.aux_data(rowsparse::kIdx);
ndarray::Copy<from_xpu, to_xpu>(from.data(), &val, from.ctx(), to.ctx(), ctx);
ndarray::Copy<from_xpu, to_xpu>(from.aux_data(rowsparse::kIdx), &idx, from.ctx(), to.ctx(), ctx);
}
// Make a copy of a dense NDArray
template <typename from_xpu, typename to_xpu>
inline void CopyFromToDnsImpl(const NDArray& from, const NDArray& to, RunContext ctx) {
#if MXNET_USE_ONEDNN == 1
// If neither is DNNL, we can copy data normally.
if (!from.IsDNNLData() && !to.IsDNNLData()) {
#endif
using namespace mshadow;
CHECK_EQ(from.storage_type(), to.storage_type()) << "Copying with different storage type";
TBlob tmp = to.data();
ndarray::Copy<from_xpu, to_xpu>(from.data(), &tmp, from.ctx(), to.ctx(), ctx);
#if MXNET_USE_ONEDNN == 1
} else if (SupportDNNL(from) && SupportDNNL(to) && from.ctx().dev_mask() == cpu::kDevMask &&
to.ctx().dev_mask() == cpu::kDevMask) {
// If we copy data directly, we need to make sure both NDArrays are supported
// by DNNL.
auto from_mem = from.GetDNNLData();
auto to_mem = to.GetDNNLData();
if (from_mem->get_desc() == to_mem->get_desc()) {
size_t size = std::min(from_mem->get_desc().get_size(), to_mem->get_desc().get_size());
memcpy(to_mem->get_data_handle(), from_mem->get_data_handle(), size);
} else {
const_cast<NDArray&>(to).CopyFrom(*from_mem);
DNNLStream::Get()->Submit();
}
} else {
// In this case, one of the NDArray isn't supported by DNNL, we need
// to convert the DNNL array to the default format first and copy data
// with Copy().
NDArray tmp_from = from;
if (tmp_from.IsDNNLData()) {
// TODO(zhengda) tmp_from should be cached.
tmp_from = NDArray(from.shape(), from.ctx(), false, from.dtype());
auto tmp_mem = from.GetDNNLData();
tmp_from.CopyFrom(*tmp_mem);
DNNLStream::Get()->Submit();
}
CHECK(tmp_from.IsDefaultData());
CHECK(to.IsDefaultData());
TBlob tmp = to.data();
ndarray::Copy<from_xpu, to_xpu>(tmp_from.data(), &tmp, from.ctx(), to.ctx(), ctx);
}
#endif
}
// Make a copy of an NDArray based on storage type
template <typename from_xpu, typename to_xpu>
void CopyFromToImpl(const NDArray& from,
const NDArray& to,
RunContext rctx,
const std::vector<Resource>& requested) {
using namespace std;
using namespace mshadow;
// if storage type doesn't match, cast the storage first
const NDArrayStorageType from_stype = from.storage_type();
const NDArrayStorageType to_stype = to.storage_type();
CHECK(from_stype == kDefaultStorage || to_stype == kDefaultStorage || from_stype == to_stype)
<< "Copying ndarray of stype = " << from_stype << " to stype = " << to_stype
<< " is not supported";
const Context from_ctx = from.ctx();
const Context to_ctx = to.ctx();
bool is_train = Imperative::Get()->is_training();
OpContext opctx{
Imperative::Get()->is_recording(), is_train, rctx, engine::CallbackOnComplete(), requested};
if (from_ctx == to_ctx && from_stype != to_stype) {
// same ctx, different stypes, use cast op directly without copying
common::CastStorageDispatch<from_xpu>(opctx, from, to);
} else {
NDArray casted_nd; // an intermediate result before copying from to to
if (from_stype == to_stype) {
casted_nd = from; // same stype, no need to cast from
} else { // different stypes on different ctx needs an temporary casted_nd
const mxnet::TShape& shape = from.shape();
if (to_stype == kDefaultStorage) {
casted_nd = NDArray(shape, from_ctx);
} else {
casted_nd = NDArray(to_stype, shape, from_ctx);
}
// convert from_nd to the same stype as to_nd
common::CastStorageDispatch<from_xpu>(opctx, from, casted_nd);
}
if (to_stype == kDefaultStorage) {
CopyFromToDnsImpl<from_xpu, to_xpu>(casted_nd, to, rctx);
} else if (to_stype == kRowSparseStorage) {
CopyFromToRspImpl<from_xpu, to_xpu>(casted_nd, to, rctx);
} else if (to_stype == kCSRStorage) {
CopyFromToCsrImpl<from_xpu, to_xpu>(casted_nd, to, rctx);
} else {
LOG(FATAL) << "unknown storage type" << to_stype;
}
}
}
void CopyFromTo(const NDArray& from, const NDArray& to, int priority, bool is_opr) {
if (from.var() == to.var() && from.byte_offset() == to.byte_offset()) {
// skip to copy to itself
return;
}
CHECK(from.shape() == to.shape())
<< "operands shape mismatch "
<< "from.shape = " << from.shape() << " to.shape=" << to.shape();
CHECK(!mxnet::op::shape_is_none(from.shape())) << "source operands have undefined shape";
// zero-size array, no need to copy
if (from.shape().Size() == 0U) {
return;
}
// important: callback must always capture by value
const Context from_ctx = from.ctx();
const int a = from_ctx.dev_mask();
const int b = to.ctx().dev_mask();
std::vector<Engine::VarHandle> const_vars;
if (from.var() != to.var())
const_vars.push_back(from.var());
const NDArrayStorageType from_stype = from.storage_type();
const NDArrayStorageType to_stype = to.storage_type();
std::vector<Engine::VarHandle> mutable_vars(1, to.var());
std::vector<Resource> requested;
if (from_stype != to_stype) {
using namespace common;
static bool log = dmlc::GetEnv("MXNET_STORAGE_FALLBACK_LOG_VERBOSE", true);
if (log) {
std::ostringstream os;
os << "\nStorage fallback detected:\n"
<< "Copy from " << stype_string(from_stype) << " storage type on " << dev_type_string(a)
<< " to " << stype_string(to_stype) << " storage type on " << dev_type_string(b)
<< ".\nA temporary ndarray with " << stype_string(to_stype)
<< " storage type will be generated in order to perform the copy. "
"This does not affect the correctness of the programme. "
"You can set environment variable "
"MXNET_STORAGE_FALLBACK_LOG_VERBOSE to 0 to suppress this warning.";
LogOnce(os.str());
}
// request temp resource if cast_storage performs on GPU
if (a == gpu::kDevMask) {
Resource rsc =
ResourceManager::Get()->Request(from_ctx, ResourceRequest(ResourceRequest::kTempSpace));
requested.push_back(rsc);
mutable_vars.push_back(rsc.var);
}
}
if (a == cpu::kDevMask && b == cpu::kDevMask) {
Engine::Get()->PushAsync(
[from, to, requested](RunContext ctx,
Engine::CallbackOnStart on_start,
Engine::CallbackOnComplete on_complete) {
on_start();
CopyFromToImpl<cpu, cpu>(from, to, ctx, requested);
on_complete();
},
from.ctx(),
const_vars,
mutable_vars,
FnProperty::kNormal,
priority,
"CopyCPU2CPU");
} else {
#if MXNET_USE_CUDA
if (a == cpu::kDevMask && b == gpu::kDevMask) {
Engine::Get()->PushAsync(
[from, to, requested](RunContext ctx,
Engine::CallbackOnStart on_start,
Engine::CallbackOnComplete on_complete) {
on_start();
CopyFromToImpl<cpu, gpu>(from, to, ctx, requested);
on_complete();
},
to.ctx(),
const_vars,
mutable_vars,
FnProperty::kCopyToGPU,
priority,
"CopyCPU2GPU");
} else if (a == gpu::kDevMask && b == cpu::kDevMask) {
Engine::Get()->PushAsync(
[from, to, requested](RunContext ctx,
Engine::CallbackOnStart on_start,
Engine::CallbackOnComplete on_complete) {
on_start();
CopyFromToImpl<gpu, cpu>(from, to, ctx, requested);
on_complete();
},
from.ctx(),
const_vars,
mutable_vars,
FnProperty::kCopyFromGPU,
priority,
"CopyGPU2CPU");
} else if (a == gpu::kDevMask && b == gpu::kDevMask) {
Engine::Get()->PushAsync(
[from, to, requested](RunContext ctx,
Engine::CallbackOnStart on_start,
Engine::CallbackOnComplete on_complete) {
on_start();
CopyFromToImpl<gpu, gpu>(from, to, ctx, requested);
on_complete();
},
from.ctx(),
const_vars,
mutable_vars,
from.dtype() != to.dtype() ? FnProperty::kNormal : FnProperty::kCopyFromGPU,
priority,
is_opr ? "_copyto_GPU2GPU" : "CopyGPU2GPU");
} else {
LOG(FATAL) << "unknown device mask";
}
#else
LOG(FATAL) << MXNET_GPU_NOT_ENABLED_ERROR;
#endif
}
}
void CopyFromTo(const NDArray& from, const NDArray* to, int priority) {
CopyFromTo(from, *to, priority);
}
void ElementwiseSum(const std::vector<NDArray>& source, NDArray* out, int priority) {
std::vector<Engine::VarHandle> const_vars;
const_vars.reserve(source.size());
for (const auto& source_array : source) {
if (source_array.var() != out->var()) {
const_vars.push_back(source_array.var());
}
CHECK_EQ(source_array.shape(), out->shape()) << "operands shape mismatch";
if (out->ctx().dev_mask() == Context::kCPU) {
CHECK_EQ(source_array.ctx().dev_mask(), Context::kCPU) << "operands context mismatch";
} else {
CHECK_EQ(source_array.ctx(), out->ctx()) << "operands context mismatch";
}
}
// important: callback must always capture by value
NDArray ret = *out;
const NDArrayStorageType stype = ret.storage_type();
if (stype == kDefaultStorage) {
switch (out->ctx().dev_mask()) {
case cpu::kDevMask: {
Engine::Get()->PushSync(
[source, ret](RunContext ctx) {
std::vector<TBlob> source_tblob(source.size());
for (size_t i = 0; i < source.size(); ++i) {
source_tblob[i] = source[i].data();
}
TBlob tmp = ret.data();
ndarray::ElementwiseSum<cpu>(source_tblob, &tmp, ctx);
},
out->ctx(),
const_vars,
{ret.var()},
FnProperty::kNormal,
priority,
PROFILER_MESSAGE_FUNCNAME);
break;
}
#if MXNET_USE_CUDA
case gpu::kDevMask: {
Engine::Get()->PushSync(
[source, ret](RunContext ctx) {
std::vector<TBlob> source_tblob(source.size());
for (size_t i = 0; i < source.size(); ++i) {
source_tblob[i] = source[i].data();
}
TBlob tmp = ret.data();
ndarray::ElementwiseSum<gpu>(source_tblob, &tmp, ctx);
},
out->ctx(),
const_vars,
{ret.var()},
FnProperty::kNormal,
priority,
"DenseElementwiseSum");
break;
}
#endif
default:
LOG(FATAL) << MXNET_GPU_NOT_ENABLED_ERROR;
}
} else if (stype == kRowSparseStorage) {
Resource rsc =
ResourceManager::Get()->Request(ret.ctx(), ResourceRequest(ResourceRequest::kTempSpace));
Engine::Get()->PushSync(
[source, ret, rsc](RunContext rctx) {
NDArray result = ret;
switch (ret.ctx().dev_mask()) {
case cpu::kDevMask: {
mxnet::ndarray::ElementwiseSum(rctx.get_stream<cpu>(), rsc, source, &result);
break;
}
#if MXNET_USE_CUDA
case gpu::kDevMask: {
mxnet::ndarray::ElementwiseSum(rctx.get_stream<gpu>(), rsc, source, &result);
break;
}
#endif
default:
LOG(FATAL) << MXNET_GPU_NOT_ENABLED_ERROR;
}
},
ret.ctx(),
const_vars,
{ret.var(), rsc.var},
FnProperty::kNormal,
priority,
"RowSparseElementwiseSum");
} else {
LOG(FATAL) << "Not implemented for storage_type " << common::stype_string(stype);
}
}
void ClipOp(const NDArray& src, const real_t& a_min, const real_t& a_max, NDArray* out) {
if (out->is_none()) {
*out = NDArray(src.shape(), src.ctx(), true, src.dtype());
} else {
CHECK(out->ctx() == src.ctx()) << "target context mismatch";
CHECK(out->shape() == src.shape()) << "target shape mismatch";
}
NDArray ret = *out;
std::vector<Engine::VarHandle> const_vars;
if (src.var() != ret.var())
const_vars.push_back(src.var());
switch (src.ctx().dev_mask()) {
case cpu::kDevMask: {
Engine::Get()->PushSync(
[src, a_min, a_max, ret](RunContext ctx) {
TBlob tmp = ret.data();
ndarray::EvalClip<cpu>(src.data(), a_min, a_max, &tmp, ctx);
},
src.ctx(),
const_vars,
{ret.var()},
FnProperty::kNormal,
0,
PROFILER_MESSAGE_FUNCNAME);
break;
}
#if MXNET_USE_CUDA
case gpu::kDevMask: {
Engine::Get()->PushSync(
[src, a_min, a_max, ret](RunContext ctx) {
TBlob tmp = ret.data();
ndarray::EvalClip<gpu>(src.data(), a_min, a_max, &tmp, ctx);
},
src.ctx(),
const_vars,
{ret.var()},
FnProperty::kNormal,
0,
PROFILER_MESSAGE_FUNCNAME);
break;
}
#endif
default:
LOG(FATAL) << MXNET_GPU_NOT_ENABLED_ERROR;
}
}
template <typename Distribution>
void SampleOP(const real_t& a, const real_t& b, NDArray* out) {
CHECK(!out->is_none());
Resource resource = ResourceManager::Get()->Request(out->ctx(), ResourceRequest::kRandom);
// important: callback must always capture by value
NDArray ret = *out;
// redirect everything to mshadow operations
switch (out->ctx().dev_mask()) {
case cpu::kDevMask: {
Engine::Get()->PushSync(
[a, b, resource, ret](RunContext ctx) {
TBlob tmp = ret.data();
ndarray::EvalRandom<cpu, Distribution>(a, b, resource, &tmp, ctx);
},
out->ctx(),
{},
{ret.var(), resource.var},
FnProperty::kNormal,
0,
PROFILER_MESSAGE_FUNCNAME);
break;
}
#if MXNET_USE_CUDA
case gpu::kDevMask: {
Engine::Get()->PushSync(
[a, b, resource, ret](RunContext ctx) {
TBlob tmp = ret.data();
ndarray::EvalRandom<gpu, Distribution>(a, b, resource, &tmp, ctx);
},
out->ctx(),
{},
{ret.var(), resource.var},
FnProperty::kNormal,
0,
PROFILER_MESSAGE_FUNCNAME);
break;
}
#endif
default:
LOG(FATAL) << MXNET_GPU_NOT_ENABLED_ERROR;
}
}
void SampleUniform(real_t begin, real_t end, NDArray* out) {
SampleOP<ndarray::UniformDistribution>(begin, end, out);
}
void SampleGaussian(real_t mu, real_t sigma, NDArray* out) {
SampleOP<ndarray::GaussianDistribution>(mu, sigma, out);
}
void SampleExponential(real_t lambda, NDArray* out) {
if (out->ctx().dev_mask() != cpu::kDevMask) {
LOG(FATAL) << "exponential sampling only valid on cpu";
}
real_t dummy;
SampleOP<ndarray::ExponentialDistribution>(lambda, dummy, out);
}
void SamplePoisson(real_t lambda, NDArray* out) {
if (out->ctx().dev_mask() != cpu::kDevMask) {
LOG(FATAL) << "poisson sampling only valid on cpu";
}
real_t dummy;
SampleOP<ndarray::PoissonDistribution>(lambda, dummy, out);
}
void SampleNegBinomial(int32_t k, real_t p, NDArray* out) {
if (out->ctx().dev_mask() != cpu::kDevMask) {
LOG(FATAL) << "negative binomial sampling only valid on cpu";
}
SampleOP<ndarray::NegBinomialDistribution>(k, p, out);
}
void SampleGenNegBinomial(real_t mu, real_t alpha, NDArray* out) {
if (out->ctx().dev_mask() != cpu::kDevMask) {
LOG(FATAL) << "negative binomial sampling only valid on cpu";
}
SampleOP<ndarray::GenNegBinomialDistribution>(mu, alpha, out);
}
void RandomSeed(uint32_t seed) {
ResourceManager::Get()->SeedRandom(seed);
}
void RandomSeed(Context ctx, uint32_t seed) {
ResourceManager::Get()->SeedRandom(ctx, seed);
}
template <typename OP>
inline NDArray BinaryOpRet(const NDArray& lhs, const NDArray& rhs) {
NDArray ret;
BinaryOpKernel<OP>(lhs, rhs, &ret);
return ret;
}
template <typename OP, bool reverse>
inline NDArray ScalarOpRet(const NDArray& lhs, const real_t& rhs) {
NDArray ret;
ScalarOp<OP, reverse>(lhs, rhs, &ret);
return ret;
}
template <typename OP>
inline NDArray& BinaryOpApply(NDArray* dst, const NDArray& src) {
BinaryOpKernel<OP>(*dst, src, dst);
return *dst;
}
template <typename OP>
inline NDArray& ScalarOpApply(NDArray* dst, const real_t& src) {
ScalarOp<OP, false>(*dst, src, dst);
return *dst;
}
// Binary
NDArray operator+(const NDArray& lhs, const NDArray& rhs) {
return BinaryOpRet<ndarray::Plus>(lhs, rhs);
}
NDArray operator-(const NDArray& lhs, const NDArray& rhs) {
return BinaryOpRet<ndarray::Minus>(lhs, rhs);
}
NDArray operator*(const NDArray& lhs, const NDArray& rhs) {
return BinaryOpRet<ndarray::Mul>(lhs, rhs);
}
NDArray operator/(const NDArray& lhs, const NDArray& rhs) {
return BinaryOpRet<ndarray::Div>(lhs, rhs);
}
// Scalar
NDArray operator+(const NDArray& lhs, const real_t& rhs) {
return ScalarOpRet<ndarray::Plus, false>(lhs, rhs);
}
NDArray operator-(const NDArray& lhs, const real_t& rhs) {
return ScalarOpRet<ndarray::Minus, false>(lhs, rhs);
}
NDArray operator*(const NDArray& lhs, const real_t& rhs) {
return ScalarOpRet<ndarray::Mul, false>(lhs, rhs);
}
NDArray operator/(const NDArray& lhs, const real_t& rhs) {
return ScalarOpRet<ndarray::Div, false>(lhs, rhs);
}
// Binary
NDArray& NDArray::operator=(real_t scalar) {
SetValueOp(scalar, this);
return *this;
}
NDArray& NDArray::operator+=(const NDArray& src) {
return BinaryOpApply<ndarray::Plus>(this, src);
}
NDArray& NDArray::operator-=(const NDArray& src) {
return BinaryOpApply<ndarray::Minus>(this, src);
}
NDArray& NDArray::operator*=(const NDArray& src) {
return BinaryOpApply<ndarray::Mul>(this, src);
}
NDArray& NDArray::operator/=(const NDArray& src) {
return BinaryOpApply<ndarray::Div>(this, src);
}
// Scalar
NDArray& NDArray::operator+=(const real_t& src) {
return ScalarOpApply<ndarray::Plus>(this, src);
}
NDArray& NDArray::operator-=(const real_t& src) {
return ScalarOpApply<ndarray::Minus>(this, src);
}
NDArray& NDArray::operator*=(const real_t& src) {
return ScalarOpApply<ndarray::Mul>(this, src);
}
NDArray& NDArray::operator/=(const real_t& src) {
return ScalarOpApply<ndarray::Div>(this, src);
}
/* magic number for ndarray version 1, with int64_t mxnet::TShape */
static const uint32_t NDARRAY_V1_MAGIC = 0xF993fac8;
/* magic number for ndarray version 2, with storage type */
static const uint32_t NDARRAY_V2_MAGIC = 0xF993fac9;
// magic number for ndarray version 3, with np shape semantics.
// The ndarray must be saved and loaded within np shape semantics.
static const uint32_t NDARRAY_V3_MAGIC = 0xF993faca;
void NDArray::Save(dmlc::Stream* strm) const {
if (Imperative::Get()->is_np_shape()) {
CHECK_EQ(storage_type(), kDefaultStorage)
<< "only allow serializing ndarray of default storage type in np shape semantics";
strm->Write(NDARRAY_V3_MAGIC);
} else {
// write magic number to mark this version
// for storage type
strm->Write(NDARRAY_V2_MAGIC);
}
// save storage type
int32_t stype = storage_type();
strm->Write(&stype, sizeof(stype));
const int32_t nad = num_aux_data(storage_type());
// save storage shape if ndarray is sparse
if (nad > 0) {
storage_shape().Save(strm);
}
// save shape
shape_.Save(strm);
if (is_none())
return;
// save context
Context ctx = this->ctx();
ctx.Save(strm);
TBlob save_data;
NDArray nd_cpu; // a copy of *this on cpu
if (ctx.dev_mask() != cpu::kDevMask) {
nd_cpu = this->Copy(Context::CPU());
nd_cpu.WaitToRead();
save_data = nd_cpu.data();
} else {
this->WaitToRead();
nd_cpu = *this;
#if MXNET_USE_ONEDNN == 1
if (nd_cpu.IsDNNLData())
nd_cpu = nd_cpu.Reorder2Default();
#endif
save_data = nd_cpu.data();
}
// save type flag
int32_t type_flag = save_data.type_flag_;
strm->Write(&type_flag, sizeof(type_flag));
// save aux_types and aux_shapes
if (nad > 0) {
for (int i = 0; i < nad; ++i) {
int32_t aux_type_flag = aux_type(i);
strm->Write(&aux_type_flag, sizeof(aux_type_flag));
aux_shape(i).Save(strm);
}
}
// save data
CHECK(save_data.CheckContiguous());
size_t type_size = mshadow::mshadow_sizeof(type_flag);
// save data could be values of sparse tensors
// must use save_data.shape_ instead of this->shape_
strm->Write(save_data.dptr_, type_size * save_data.shape_.Size());
// save aux data
if (nad > 0) {
for (int i = 0; i < nad; ++i) {
TBlob save_data = nd_cpu.aux_data(i);
// save aux_data
CHECK(save_data.CheckContiguous());
size_t aux_type_size = mshadow::mshadow_sizeof(aux_type(i));
strm->Write(save_data.dptr_, aux_type_size * save_data.Size());
}
}
}
bool LegacyTShapeLoad(dmlc::Stream* strm, mxnet::TShape* shape, const uint32_t magic) {
switch (magic) {
case NDARRAY_V1_MAGIC:
return shape->Load(strm);
default:
// meet legacy mxnet::TShape, magic is ndim here
uint32_t ndim = magic;
*shape = mxnet::TShape(ndim, -1);
std::vector<uint32_t> buffer(ndim);
size_t nread = ndim * sizeof(uint32_t);
if (strm->Read(buffer.data(), nread) != nread)
return false;
nnvm::ShapeTypeCast(buffer.begin(), buffer.end(), shape->begin());
return true;
}
}
bool NDArray::LegacyLoad(dmlc::Stream* strm, const uint32_t magic) {
// load shape
mxnet::TShape shape;
if (!LegacyTShapeLoad(strm, &shape, magic))
return false;
if (mxnet::op::shape_is_none(shape)) {
*this = NDArray();
return true;
}
// load context
Context ctx;
if (!ctx.Load(strm))
return false;
// load type flag
int32_t type_flag;
if (strm->Read(&type_flag, sizeof(type_flag)) != sizeof(type_flag))
return false;
// load data into CPU
NDArray temp(shape, Context::CPU(), false, type_flag);
TBlob load_data = temp.data();
size_t type_size = mshadow::mshadow_sizeof(type_flag);
size_t nread = type_size * shape.Size();
if (strm->Read(load_data.dptr_, nread) != nread)
return false;
if (ctx.dev_mask() == cpu::kDevMask) {
*this = std::move(temp);
return true;
} else {
#if MXNET_USE_CUDA
*this = temp.Copy(ctx);
return true;
#else
*this = std::move(temp);
return true;
#endif
}
}
bool NDArray::Load(dmlc::Stream* strm) {
uint32_t magic;
if (strm->Read(&magic, sizeof(uint32_t)) != sizeof(uint32_t))
return false;
if (magic == NDARRAY_V3_MAGIC) {
CHECK(Imperative::Get()->is_np_shape())
<< "ndarray was saved in np shape semantics, must be loaded in the same semantics."
" Please turn on np shape semantics in Python using `with np_shape(True)`"
" or decorator `use_np_shape` to scope the code of loading the ndarray.";
} else {
// when the flag is global on, skip the check since it would be always global on.
CHECK(Imperative::Get()->is_np_shape() == GlobalOn || !Imperative::Get()->is_np_shape())
<< "ndarray was not saved in np shape semantics, but being loaded in np shape semantics."
" Please turn off np shape semantics in Python using `with np_shape(False)`"
" to scope the code of loading the ndarray.";
}
if (magic != NDARRAY_V2_MAGIC && magic != NDARRAY_V3_MAGIC) {
return LegacyLoad(strm, magic);
}
// load storage type
int32_t stype;
if (strm->Read(&stype, sizeof(stype)) != sizeof(stype))
return false;
if (Imperative::Get()->is_np_shape()) {
CHECK_EQ(stype, kDefaultStorage)
<< "only allow deserializing ndarray of default storage type in np shape semantics";
}
const int32_t nad = num_aux_data(static_cast<NDArrayStorageType>(stype));
// load storage shape
mxnet::TShape sshape;
if (nad > 0) {
if (!sshape.Load(strm))
return false;
}
// load shape
mxnet::TShape shape;
if (!shape.Load(strm))
return false;
if (Imperative::Get()->is_np_shape()) {
if (!shape_is_known(shape)) {
*this = NDArray();
return true;
}
} else if (shape.ndim() == 0) {
*this = NDArray();
return true;
}
// load context
Context ctx;
if (!ctx.Load(strm))
return false;
// load type flag
int32_t type_flag;
if (strm->Read(&type_flag, sizeof(type_flag)) != sizeof(type_flag))
return false;
// load aux_types and aux_shapes
std::vector<int32_t> aux_types;
mxnet::ShapeVector aux_shapes;
if (nad > 0) {
aux_types.resize(nad);
aux_shapes.resize(nad);
for (int i = 0; i < nad; ++i) {
// load aux_type(i)
if (strm->Read(&aux_types[i], sizeof(aux_types[i])) != sizeof(aux_types[i]))
return false;
// load aux_shapes(i)
if (!aux_shapes[i].Load(strm))
return false;
}
}
// load data into CPU
NDArray temp;
if (0 == nad) {
temp = NDArray(shape, Context::CPU(), false, type_flag);
} else {
temp = NDArray(static_cast<NDArrayStorageType>(stype),
shape,
Context::CPU(),
false,
type_flag,
aux_types,
aux_shapes,
sshape);
}
// load data
TBlob load_data = temp.data();
size_t type_size = mshadow::mshadow_sizeof(type_flag);
size_t nread = type_size * load_data.Size();
if (strm->Read(load_data.dptr_, nread) != nread)
return false;
// load aux_data
if (nad > 0) {
for (int i = 0; i < nad; ++i) {
load_data = temp.aux_data(i);
type_size = mshadow::mshadow_sizeof(load_data.type_flag_);
nread = type_size * load_data.Size();
if (strm->Read(load_data.dptr_, nread) != nread)
return false;
}
}
if (ctx.dev_mask() == cpu::kDevMask) {
*this = std::move(temp);
return true;
} else {
#if MXNET_USE_CUDA
int device_count = -1;
cudaGetDeviceCount(&device_count);
if (device_count > 0) {
*this = temp.Copy(ctx);
return true;
} else {
*this = std::move(temp);
return true;
}
#else
*this = std::move(temp);
return true;
#endif
}
}
const uint64_t kMXAPINDArrayListMagic = 0x112;
void NDArray::Save(dmlc::Stream* fo,
const std::vector<NDArray>& data,
const std::vector<std::string>& names) {
uint64_t header = kMXAPINDArrayListMagic, reserved = 0;
fo->Write(&header, sizeof(header));
fo->Write(&reserved, sizeof(reserved));
fo->Write(data);
fo->Write(names);
}
void NDArray::Load(dmlc::Stream* fi, std::vector<NDArray>* data, std::vector<std::string>* keys) {
uint64_t header, reserved;
CHECK(fi->Read(&header)) << "Invalid NDArray file format";
CHECK(fi->Read(&reserved)) << "Invalid NDArray file format";
CHECK(header == kMXAPINDArrayListMagic) << "Invalid NDArray file format";
CHECK(fi->Read(data)) << "Invalid NDArray file format";
CHECK(fi->Read(keys)) << "Invalid NDArray file format";
CHECK(keys->size() == 0 || keys->size() == data->size()) << "Invalid NDArray file format";
}
NDArray NDArray::Copy(Context ctx) const {
NDArray ret;
if (kDefaultStorage == storage_type()) {
ret = NDArray(shape(), ctx, false, dtype_);
} else if (kUndefinedStorage != storage_type()) {
ret = NDArray(storage_type(),
shape(),
ctx,
false,
dtype_,
ptr_->aux_types,
ptr_->aux_shapes,
storage_shape());
} else {
LOG(FATAL) << "NDArray::Copy cannot copy undefined storage-type ndarray to ctx.dev_type="
<< ctx.dev_type << ", ctx.dev_id=" << ctx.dev_id;
}
CopyFromTo(*this, ret);
return ret;
}
void NDArray::SyncCopyFromCPU(const void* data, size_t size) const {
mxnet::TShape dshape = this->shape();
if (!features::is_enabled(features::INT64_TENSOR_SIZE)) {
CHECK_LT(size, (int64_t{1} << 31) - 1)
<< "[SyncCopyFromCPU] Size of tensor you are trying to allocate is larger than "
"2^31 elements. Please build with flag USE_INT64_TENSOR_SIZE=1";
}
CHECK_EQ(dshape.Size(), size) << "Memory size do not match";
// zero-size array, no need to copy
if (size == 0U) {
return;
}
TBlob src((void*)data, dshape, cpu::kDevMask, this->dtype_, 0); // NOLINT(*)
if (this->ctx().dev_mask() == cpu::kDevMask) {
this->WaitToWrite();
RunContext rctx{this->ctx(), nullptr, nullptr};
TBlob dst = this->data();
ndarray::Copy<cpu, cpu>(src, &dst, Context::CPU(), Context::CPU(), rctx);
} else {
#if MXNET_USE_CUDA
Engine::Get()->PushAsync(
[&](RunContext rctx,
Engine::CallbackOnStart on_start,
Engine::CallbackOnComplete on_complete) {
on_start();
TBlob dst = this->data();
ndarray::Copy<cpu, gpu>(src, &dst, Context::CPU(), this->ctx(), rctx);
on_complete();
},
this->ctx(),
{},
{this->var()},
FnProperty::kCopyToGPU,
0,
"SyncCopyCPU2GPU");
this->WaitToRead();
#else
LOG(FATAL) << "GPU is not enabled";
#endif
}
}
/*!
* \brief Copy src.data()/aux_data(i) to dst->data()/aux_data(j).
*/
void NDArray::SyncCopyFromNDArray(const NDArray& src, int i, int j) {
if (i >= 0) {
CHECK_NE(src.storage_type(), kDefaultStorage);
} else {
CHECK(!src.is_none()) << "src dense ndarray must have been initialized";
}
if (j >= 0) {
CHECK_NE(storage_type(), kDefaultStorage);
} else {
CHECK(!this->is_none()) << "dst dense ndarray must have been initialized";
}
if (src.var() == var()) {
// skip to copy to itself
LOG(WARNING) << "SyncCopyFromNDArray does not support copying to self";
return;
}
const int src_dev_mask = src.ctx().dev_mask();
const int dst_dev_mask = ctx().dev_mask();
std::vector<Engine::VarHandle> const_vars;
const_vars.push_back(src.var());
// get or create a dst tblob for copying src to it
// if dst is a dense format and has not been allocated, allocate memory for it
// else if dst is not initialized, allocate corresponding data blob for it
auto get_dst_data = [&](const mxnet::TShape& src_shape) {
if (this->storage_type() == kDefaultStorage) {
this->ReshapeAndAlloc(src_shape);
} else if (!this->storage_initialized()) {
if (j < 0) {
this->CheckAndAllocData(src_shape);
} else {
this->CheckAndAllocAuxData(j, src_shape);
}
}
TBlob dst_data = (j >= 0 ? this->aux_data(j) : this->data());
CHECK_LE(src_shape.Size(), dst_data.shape_.Size());
return dst_data;
};
if (src_dev_mask == cpu::kDevMask && dst_dev_mask == cpu::kDevMask) {
Engine::Get()->PushSync(
[&](RunContext rctx) {
const TBlob src_data = (i >= 0 ? src.aux_data(i) : src.data());
TBlob dst_data = get_dst_data(src_data.shape_);
ndarray::Copy<cpu, cpu>(src_data, &dst_data, src.ctx(), this->ctx(), rctx);
},
this->ctx(),
const_vars,
{this->var()},
FnProperty::kNormal,
0,
"SyncCopyFromNDArrayCPU2CPU");
} else {
#if MXNET_USE_CUDA
if (src_dev_mask == cpu::kDevMask && dst_dev_mask == gpu::kDevMask) {
Engine::Get()->PushAsync(
[&](RunContext rctx,
Engine::CallbackOnStart on_start,
Engine::CallbackOnComplete on_complete) {
on_start();
const TBlob src_data = (i >= 0 ? src.aux_data(i) : src.data());
TBlob dst_data = get_dst_data(src_data.shape_);
ndarray::Copy<cpu, gpu>(src_data, &dst_data, src.ctx(), this->ctx(), rctx);
on_complete();
},
this->ctx(),
const_vars,
{this->var()},
FnProperty::kCopyToGPU,
0,
"SyncCopyFromNDArrayCPU2GPU");
} else if (src_dev_mask == gpu::kDevMask && dst_dev_mask == cpu::kDevMask) {
Engine::Get()->PushAsync(
[&](RunContext rctx,
Engine::CallbackOnStart on_start,
Engine::CallbackOnComplete on_complete) {
on_start();
const TBlob src_data = (i >= 0 ? src.aux_data(i) : src.data());
TBlob dst_data = get_dst_data(src_data.shape_);
ndarray::Copy<gpu, cpu>(src_data, &dst_data, src.ctx(), this->ctx(), rctx);
on_complete();
},
src.ctx(),
const_vars,
{this->var()},
FnProperty::kCopyFromGPU,
0,
"SyncCopyFromNDArrayGPU2CPU");
} else if (src_dev_mask == gpu::kDevMask && dst_dev_mask == gpu::kDevMask) {
Engine::Get()->PushAsync(
[&](RunContext rctx,
Engine::CallbackOnStart on_start,
Engine::CallbackOnComplete on_complete) {
on_start();
const TBlob src_data = (i >= 0 ? src.aux_data(i) : src.data());
TBlob dst_data = get_dst_data(src_data.shape_);
ndarray::Copy<gpu, gpu>(src_data, &dst_data, src.ctx(), this->ctx(), rctx);
on_complete();
},
this->ctx(),
const_vars,
{this->var()},
src.dtype() != this->dtype() ? FnProperty::kNormal : FnProperty::kCopyFromGPU,
0,
"SyncCopyFromNDArrayGPU2GPU");
} else {
LOG(FATAL) << "unknown device mask";
}
#else
LOG(FATAL) << MXNET_GPU_NOT_ENABLED_ERROR;
#endif
}
// The copy operation was pushed to engine to execute.
// Need to wait here for it being completed.
// The reason for pushing the copy operation to engine
// is because when copying data from a sparse tensor
// to the current one, that sparse ndarray's storage_shape/aux_shape
// may not be ready or changed and we need to ensure
// thread safty for reading the correct shape info to allocate
// memory for the current ndarray.
WaitToRead();
}
void NDArray::SyncCopyToCPU(void* data, size_t size) const {
mxnet::TShape dshape = this->shape();
if (!features::is_enabled(features::INT64_TENSOR_SIZE)) {
CHECK_LT(size, (int64_t{1} << 31) - 1)
<< "[SyncCopyToCPU] Size of tensor you are trying to allocate is larger than "
"2^31 elements. Please build with flag USE_INT64_TENSOR_SIZE=1";
}
CHECK_EQ(dshape.Size(), size) << "Memory size do not match";
// zero-size array, no need to copy
if (size == 0U) {
return;
}
TBlob dst(data, dshape, cpu::kDevMask, this->dtype_, 0); // NOLINT(*)
this->WaitToRead();
if (this->ctx().dev_mask() == cpu::kDevMask) {
RunContext rctx{this->ctx(), nullptr, nullptr};
NDArray src = *this;
#if MXNET_USE_ONEDNN == 1
if (src.IsDNNLData())
src = this->Reorder2Default();
#endif
ndarray::Copy<cpu, cpu>(src.data(), &dst, Context::CPU(), Context::CPU(), rctx);
} else {
#if MXNET_USE_CUDA
Engine::Get()->PushAsync(
[&](RunContext rctx,
Engine::CallbackOnStart on_start,
Engine::CallbackOnComplete on_complete) {
on_start();
{
auto var = this->var();
auto& sync_obj = var->sync_object;
std::lock_guard<std::mutex> lock{sync_obj.mutex};
bool has_writer = false;
std::shared_ptr<cudaEvent_t> w_ev_ptr;
if (!sync_obj.writer_event.empty()) {
w_ev_ptr = sync_obj.writer_event[0].event.lock();
has_writer = w_ev_ptr ? true : false;
}
for (auto ev : sync_obj.reader_events) {
auto event_ptr = ev.event.lock();
if (!event_ptr) {
continue;
}
cudaEvent_t event = *event_ptr;
if (has_writer) {
auto w_ev = sync_obj.writer_event[0];
if (w_ev.stream == ev.stream) {
event = w_ev.pool_index > ev.pool_index ? *w_ev_ptr : *event_ptr;
has_writer = false;
}
}
CUDA_CALL(cudaEventSynchronize(event));
}
if (has_writer) {
CUDA_CALL(cudaEventSynchronize(*w_ev_ptr));
}
}
ndarray::Copy<gpu, cpu>(this->data(), &dst, this->ctx(), Context::CPU(), rctx);
on_complete();
},
this->ctx(),
{this->var()},
{},
FnProperty::kCopyFromGPU,
0,
"SyncCopyGPU2CPU");
this->WaitToWrite();
#else
LOG(FATAL) << "GPU is not enabled";
#endif
}
}
void NDArray::SyncCheckFormat(const bool full_check) const {
int32_t err = kNormalErr;
TBlob err_cpu(&err, mshadow::Shape1(1), cpu::kDevMask, 0);
if (this->ctx().dev_mask() == cpu::kDevMask) {
Engine::Get()->PushSync(
[&](RunContext rctx) { common::CheckFormatWrapper<cpu>(rctx, *this, err_cpu, full_check); },
this->ctx(),
{this->var()},
{},
FnProperty::kNormal,
0,
"CheckFormat");
} else {
#if MXNET_USE_CUDA
Engine::Get()->PushSync(
[&](RunContext rctx) { common::CheckFormatWrapper<gpu>(rctx, *this, err_cpu, full_check); },
this->ctx(),
{this->var()},
{},
FnProperty::kNormal,
0,
"CheckFormat");
#else
LOG(FATAL) << "GPU is not enabled";
#endif
}
this->WaitToWrite();
CHECK_NE(err, kCSRShapeErr) << "Shape mismatch of this csr NDArray";
CHECK_NE(err, kCSRIndPtrErr)
<< "IndPtr of csr NDArray should be non-negative, in non-decreasing order, "
<< "start with 0, and end with value equal with size of indices.";
CHECK_NE(err, kCSRIdxErr)
<< "Indices of csr NDArray should be non-negative, in ascending order per row "
<< " and less than the number of columns.";
CHECK_NE(err, kRSPShapeErr) << "Shape mismatch of this row_sparse NDArray";
CHECK_NE(err, kRSPIdxErr) << "Indices of row_sparse NDArray should be non-negative, "
<< "less than the size of first dimension and in ascending order";
CHECK_EQ(err, kNormalErr) << "Check the validity of this sparse NDArray";
}
void NDArray::WaitToRead() const {
if (is_none())
return;
Imperative::DCInfo::Compute(*this);
Engine::Get()->WaitForVar(ptr_->var);
}
void NDArray::WaitToWrite() const {
if (is_none())
return;
Imperative::DCInfo::Compute(*this);
// Push an empty mutable function to flush all preceding reads to the variable.
Engine::Get()->PushAsync(
[](RunContext, Engine::CallbackOnStart on_start, Engine::CallbackOnComplete on_complete) {
on_start();
on_complete();
},
Context{},
{},
{ptr_->var});
Engine::Get()->WaitForVar(ptr_->var);
}
void NDArray::StreamSync(int stream) const {
if (is_none())
return;
Imperative::DCInfo::Compute(*this);
#if MXNET_USE_CUDA
Engine::Get()->PushAsync(
[this, stream](RunContext ctx,
Engine::CallbackOnStart on_start,
Engine::CallbackOnComplete on_complete) {
on_start();
cudaStream_t consumer = reinterpret_cast<cudaStream_t>(stream);
std::unordered_map<cudaStream_t, engine::EventInfo> events_per_stream;
auto& sync_obj = this->var()->sync_object;
std::lock_guard<std::mutex> l(sync_obj.mutex);
auto& reader_events = sync_obj.reader_events;
reader_events.erase(
std::remove_if(reader_events.begin(),
reader_events.end(),
[&](const engine::EventInfo e_i) { return e_i.event.expired(); }),
reader_events.end());
for (auto& writer : sync_obj.writer_event) {
if (writer.event.expired()) {
sync_obj.writer_event.clear();
break;
}
if (writer.stream != consumer) {
bool found = false;
for (const auto& reader : reader_events) {
if (reader.stream == consumer) {
found = true;
break;
}
}
if (!found) {
auto event_stream = writer.stream;
if (events_per_stream.count(event_stream) > 0) {
if (events_per_stream[event_stream].pool_index < writer.pool_index) {
events_per_stream[event_stream] = writer;
}
} else {
events_per_stream.emplace(event_stream, writer);
}
}
}
}
for (auto event : events_per_stream) {
auto ev = event.second.event.lock();
MSHADOW_CUDA_CALL(cudaStreamWaitEvent(consumer, *ev, 0));
}
on_complete();
},
this->ctx(),
{},
{});
#else
LOG(FATAL) << "GPU is not enabled";
#endif
}
#if MXNET_PREDICT_ONLY == 0
// register API function
// those with underscore will be registered at NDArray
MXNET_REGISTER_NDARRAY_FUN(_set_value).set_function(SetValueOp);
MXNET_REGISTER_NDARRAY_FUN(_onehot_encode).set_function(BinaryOp<ndarray::OneHotEncode>);
MXNET_REGISTER_NDARRAY_FUN(fill_element_0index)
.set_function(TernaryOp<ndarray::MatFillRowElem>)
.describe(
"Fill one element of each line(row for python, column for R/Julia)"
" in lhs according to index indicated by rhs and values indicated by mhs."
" This function assume rhs uses 0-based index.");
// register API function
// those with underscore will be registered at NDArray
void CopyFromToSimple(const nnvm::NodeAttrs& attrs,
const OpContext& ctx,
const std::vector<NDArray>& inputs,
const std::vector<OpReqType>& req,
const std::vector<NDArray>& outputs) {
CopyFromTo(inputs[0], outputs[0], 0, true);
}
bool CopyToType(const nnvm::NodeAttrs& attrs,
std::vector<int>* in_attrs,
std::vector<int>* out_attrs) {
CHECK_EQ(in_attrs->size(), 1U);
CHECK_EQ(out_attrs->size(), 1U);
int in_type = in_attrs->at(0);
if (out_attrs->at(0) == -1) {
TYPE_ASSIGN_CHECK(*out_attrs, 0, in_type);
}
return out_attrs->at(0) != -1;
}
// copy function is special
// that we need to remove kAcceptEmptyMutateTarget from it
NNVM_REGISTER_OP(_copyto)
.add_alias("_npi_copyto")
.set_num_inputs(1)
.set_num_outputs(1)
.set_attr<mxnet::FInferShape>("FInferShape", op::ElemwiseShape<1, 1>)
.set_attr<nnvm::FInferType>("FInferType", CopyToType)
.set_attr<FInferStorageType>("FInferStorageType",
[](const NodeAttrs& attrs,
const int dev_mask,
DispatchMode* dispatch_mode,
std::vector<int>* in_attrs,
std::vector<int>* out_attrs) {
op::dispatch_mode_assign(dispatch_mode,
DispatchMode::kFComputeEx);
if (op::storage_type_is_none((*out_attrs)[0])) {
(*out_attrs)[0] = (*in_attrs)[0];
}
return true;
})
.set_attr<FExecType>("FExecType",
[](const NodeAttrs& attrs) { return ExecType::kCrossDeviceCopy; })
.set_attr<nnvm::FGradient>("FGradient", op::ElemwiseGradUseNone{"_copyto"})
.set_attr<bool>("TIsBackward", true)
.set_attr<FComputeEx>("FComputeEx<cpu>", CopyFromToSimple)
.set_attr<FComputeEx>("FComputeEx<gpu>", CopyFromToSimple)
.add_argument("data", "NDArray", "input data");
void Imdecode(NDArray* ret,
NDArray mean,
size_t index,
size_t x0,
size_t y0,
size_t x1,
size_t y1,
size_t n_channels,
size_t size,
char* str_img) {
#if MXNET_USE_OPENCV
cv::Mat buf(1, size, CV_8U, str_img);
cv::Mat res = cv::imdecode(buf, n_channels == 1 ? 0 : -1);
CHECK(res.data != nullptr) << "OpenCV Failed to decode image";
CHECK_LE(n_channels, static_cast<size_t>(res.channels()));
if (y1 - y0 == 0) {
x0 = 0;
x1 = res.cols;
y0 = 0;
y1 = res.rows;
}
CHECK(x1 <= static_cast<size_t>(res.cols) && y1 <= static_cast<size_t>(res.rows));
if (ret->is_none()) {
*ret = NDArray(mshadow::Shape3(n_channels, y1 - y0, x1 - x0),
Context::CPU(),
false,
mean.is_none() ? mshadow::default_type_flag : mean.dtype());
}
NDArray buff;
if (ret->shape().ndim() == 3) {
buff = ret->Reshape(mshadow::Shape4(1, ret->shape()[0], ret->shape()[1], ret->shape()[2]));
} else {
CHECK_EQ(ret->shape().ndim(), 4U);
buff = ret->Slice(index, index + 1);
}
CHECK_EQ(buff.ctx().dev_mask(), Context::kCPU);
CHECK_EQ(n_channels, buff.shape()[1]);
CHECK_EQ(y1 - y0, buff.shape()[2]);
CHECK_EQ(x1 - x0, buff.shape()[3]);
buff.WaitToWrite();
if (mean.is_none()) {
MSHADOW_TYPE_SWITCH(buff.dtype(), DType, {
mshadow::Tensor<cpu, 4, DType> tensor = buff.data().get<cpu, 4, DType>();
for (size_t i = 0; i < y1 - y0; i++) {
uchar* im_data = res.ptr<uchar>(y0 + i) + res.channels() * x0;
for (size_t j = 0; j < x1 - x0; j++) {
for (size_t k = 0; k < n_channels; k++) {
tensor[0][k][i][j] = DType(im_data[k]); // NOLINT(*)
}
im_data += res.channels();
}
}
})
} else {
CHECK_EQ(mean.dtype(), buff.dtype());
CHECK_EQ(mean.ctx().dev_mask(), Context::kCPU);
CHECK_EQ(mean.shape()[0], buff.shape()[1]);
CHECK_EQ(mean.shape()[1], buff.shape()[2]);
CHECK_EQ(mean.shape()[2], buff.shape()[3]);
mean.WaitToRead();
MSHADOW_TYPE_SWITCH(buff.dtype(), DType, {
mshadow::Tensor<cpu, 4, DType> tensor = buff.data().get<cpu, 4, DType>();
mshadow::Tensor<cpu, 3, DType> tmean = mean.data().get<cpu, 3, DType>();
for (size_t i = 0; i < y1 - y0; i++) {
uchar* im_data = res.ptr<uchar>(y0 + i) + res.channels() * x0;
for (size_t j = 0; j < x1 - x0; j++) {
for (size_t k = 0; k < n_channels; k++) {
tensor[0][k][i][j] = DType(im_data[k]) - tmean[k][i][j]; // NOLINT(*)
}
im_data += res.channels();
}
}
})
}
#else
LOG(FATAL) << "Compile with OpenCV for image decoding.";
#endif // MXNET_USE_OPENCV
}
MXNET_REGISTER_NDARRAY_FUN(_imdecode)
.set_type_mask(kAcceptEmptyMutateTarget | kNDArrayArgBeforeScalar)
.set_body([](NDArray** u,
real_t* s,
NDArray** out,
int num_params,
char** param_keys,
char** param_vals) {
CHECK_EQ(num_params, 1);
Imdecode(out[0],
*u[0],
static_cast<size_t>(s[0]),
static_cast<size_t>(s[1]),
static_cast<size_t>(s[2]),
static_cast<size_t>(s[3]),
static_cast<size_t>(s[4]),
static_cast<size_t>(s[5]),
static_cast<size_t>(s[6]),
param_vals[0]);
})
.set_num_use_vars(1)
.set_num_scalars(7)
.set_num_mutate_vars(1)
.describe("Decode an image, clip to (x0, y0, x1, y1), subtract mean, and write to buffer")
.add_argument("mean", "NDArray-or-Symbol", "image mean")
.add_argument("index", "int", "buffer position for output")
.add_argument("x0", "int", "x0")
.add_argument("y0", "int", "y0")
.add_argument("x1", "int", "x1")
.add_argument("y1", "int", "y1")
.add_argument("c", "int", "channel")
.add_argument("size", "int", "length of str_img");
#endif
} // namespace mxnet