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

Method get_workspace_in_bytes

dnn/src/cuda/group_local/forward/opr_impl.cpp:92–104  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

90 : GroupLocalForward(handle) {}
91
92size_t GroupLocalForwardImpl::get_workspace_in_bytes(
93 const TensorLayout& src, const TensorLayout& filter, const TensorLayout& dst) {
94 if (prefer_inference_kernel(src, filter, dst)) {
95 return 0;
96 } else {
97 auto G = filter[0];
98 TensorLayout src_g = prepare_src_dst(src, G);
99 TensorLayout dst_g = prepare_src_dst(dst, G);
100 TensorLayout filter_g = prepare_filter(filter);
101 return get_opr(handle(), param())
102 ->get_workspace_in_bytes(src_g, filter_g, dst_g);
103 }
104}
105
106bool GroupLocalForwardImpl::prefer_inference_kernel(
107 const TensorLayout& src, const TensorLayout& filter, const TensorLayout& dst) {

Callers

nothing calls this directly

Calls 4

prepare_src_dstFunction · 0.85
prepare_filterFunction · 0.85
get_oprFunction · 0.85
paramFunction · 0.50

Tested by

no test coverage detected