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

Function apply_on_physical_tensor

imperative/src/impl/ops/pooling.cpp:37–81  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

35}
36
37SmallVector<TensorPtr> apply_on_physical_tensor(
38 const OpDef& def, const SmallVector<TensorPtr>& inputs,
39 SmallVector<LogicalTensorDesc>& output_descs, const bool& validated) {
40 mgb_assert(
41 inputs.size() == 1, "num of inputs of pooling should be 1 but you give %zu",
42 inputs.size());
43
44 auto&& pooling = def.cast_final_safe<Pooling>();
45 auto cn = inputs[0]->comp_node();
46 DnnOprCaller<megdnn::Pooling> dnn_opr(cn, pooling.param(), pooling.policy());
47 auto oup_layout = [&] {
48 if (validated) {
49 return output_descs[0].layout;
50 } else {
51 return dnn_opr.deduce_layout(inputs[0]->layout());
52 }
53 }();
54
55 auto&& src_layout = inputs[0]->layout();
56 if (src_layout.ndim == 4 && src_layout.is_empty()) {
57 if (pooling.param().format == megdnn::Pooling::Param::Format::NCHW) {
58 mgb_assert(
59 src_layout.shape[2] != 0 && src_layout.shape[3] != 0,
60 "Pooling expect input to have non-zero size for non-batch "
61 "dimensions, but the input has layout %s",
62 src_layout.to_string().c_str());
63 } else if (pooling.param().format == megdnn::Pooling::Param::Format::NHWC) {
64 mgb_assert(
65 src_layout.shape[1] != 0 && src_layout.shape[2] != 0,
66 "Pooling expect input to have non-zero size for non-batch "
67 "dimensions, but the input has layout %s",
68 src_layout.to_string().c_str());
69 } else {
70 mgb_assert(
71 false,
72 "Pooling support empty input only when the input format is NHWC or "
73 "NCHW");
74 }
75 return {Tensor::make(oup_layout, cn)};
76 }
77
78 auto out = Tensor::make(oup_layout, cn);
79 dnn_opr.exec_fastrun(inputs[0], out);
80 return {out};
81}
82
83OP_TRAIT_REG(Pooling, Pooling)
84 .apply_on_var_node(apply_on_var_node)

Callers

nothing calls this directly

Calls 9

exec_fastrunMethod · 0.80
makeFunction · 0.70
sizeMethod · 0.45
comp_nodeMethod · 0.45
paramMethod · 0.45
deduce_layoutMethod · 0.45
layoutMethod · 0.45
is_emptyMethod · 0.45
to_stringMethod · 0.45

Tested by

no test coverage detected