MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / get_unary_ws

Function get_unary_ws

dnn/src/cambricon/elemwise/opr_impl.cpp:676–699  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

674}
675
676WorkspaceBundle get_unary_ws(
677 HandleImpl* handle, const TensorLayoutArray& src, const TensorLayout& dst,
678 const param::Elemwise::Mode& mode) {
679 auto cnnl_handler = handle->cnnl_handle();
680 bool mode_need_ws = mode == Mode::NOT;
681 if (!mode_need_ws)
682 return {nullptr, {}, handle->alignment_requirement()};
683 CnnlTensorDescriptor input_desc, output_desc;
684 input_desc.set(src[0]);
685 output_desc.set(dst);
686 // 1st workspace is cnnl workspace
687 SmallVector<size_t> sizes_in_bytes(1, 0);
688 switch (mode) {
689 case Mode::NOT:
690 cnnl_check(cnnlGetLogicOpWorkspaceSize(
691 cnnl_handler, input_desc.desc(), input_desc.desc(),
692 output_desc.desc(), &sizes_in_bytes[0]));
693 break;
694 default:
695 break;
696 }
697
698 return make_bundle(handle, sizes_in_bytes);
699}
700WorkspaceBundle get_binary_ws(
701 HandleImpl* handle, const TensorLayoutArray& src, const TensorLayout& dst,
702 const param::Elemwise::Mode& mode) {

Callers 1

alloc_cnnl_workspaceMethod · 0.85

Calls 5

make_bundleFunction · 0.85
cnnl_handleMethod · 0.80
alignment_requirementMethod · 0.45
setMethod · 0.45
descMethod · 0.45

Tested by

no test coverage detected