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

Function apply_on_physical_tensor

imperative/src/impl/ops/convolution.cpp:57–144  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

55}
56
57SmallVector<TensorPtr> apply_on_physical_tensor(
58 const OpDef& def, const SmallVector<TensorPtr>& inputs,
59 SmallVector<LogicalTensorDesc>& output_descs, const bool& validated) {
60 // create megdnn opr
61 auto&& conv = def.cast_final_safe<Convolution>();
62 CompNode cn = inputs[0]->comp_node();
63
64 auto check_empty_inputs = [&]() -> void {
65 using Format = megdnn::Convolution::Param::Format;
66 auto format = conv.param().format;
67 mgb_assert(
68 format == Format::NCHW || format == Format::NHWC,
69 "Convolution support empty input only when the input format is "
70 "NHWC or NCHW");
71
72 auto filter_ndim = inputs[1]->layout().ndim;
73 size_t h = 2, w = 3, kh = filter_ndim - 2, kw = filter_ndim - 1;
74 if (format == Format::NHWC) {
75 h = 1;
76 w = 2;
77 }
78 mgb_assert(
79 inputs[0]->layout()[h] != 0 && inputs[0]->layout()[w] != 0 &&
80 inputs[1]->layout()[kh] != 0 && inputs[1]->layout()[kw] != 0,
81 "Convolution expect input to have non-zero size for non-batch and "
82 "non-channel "
83 "dimensions, but the inputs has layout %s and %s",
84 inputs[0]->layout().to_string().c_str(),
85 inputs[1]->layout().to_string().c_str());
86 };
87
88 // calling dnn ConvolutionForward when device is rocm
89 // because there is no dnn ConvBiasForward on rocm
90 if (cn.device_type() == CompNode::DeviceType::ROCM) {
91 DnnOprCaller<megdnn::ConvolutionForward> dnn_opr(
92 cn, conv.param(), conv.policy());
93 auto out_layout = [&] {
94 if (validated) {
95 return output_descs[0].layout;
96 } else {
97 return dnn_opr.deduce_layout(inputs[0]->layout(), inputs[1]->layout());
98 }
99 }();
100
101 // alloc memory
102 auto out = Tensor::make(out_layout, cn);
103
104 if (inputs[0]->layout().is_empty() || inputs[1]->layout().is_empty()) {
105 check_empty_inputs();
106 return {out};
107 }
108
109 dnn_opr.exec_fastrun(inputs[0], inputs[1], out);
110 return {out};
111 }
112
113 // calling dnn ConvBiasForward on cuda because it's faster then ConvolutionForward
114 // ConvolutionForward internally uses ConvBiasForward to calculate the result

Callers

nothing calls this directly

Calls 15

exec_fastrunMethod · 0.80
exec_with_wsMethod · 0.80
makeFunction · 0.70
comp_nodeMethod · 0.45
paramMethod · 0.45
layoutMethod · 0.45
to_stringMethod · 0.45
device_typeMethod · 0.45
deduce_layoutMethod · 0.45
is_emptyMethod · 0.45
dtypeMethod · 0.45

Tested by

no test coverage detected