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

Method Impl

src/tensorrt/impl/opr_replace.cpp:145–1267  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

143
144public:
145 Impl(const Pass& pass, OptState& opt_state)
146 : m_pass{pass},
147 m_opt_state{opt_state},
148 m_rewriter{opt_state.graph().make_rewriter()},
149 m_const_var_propogate{std::make_unique<ConstVarPropogate>(
150 ConstVarType::IMMUTABLE_AND_PARAM)} {
151#define REPLACE_FAIL_MSG_EPILOGUE \
152 { \
153 auto&& mgr = opr->owner_graph()->static_infer_manager(); \
154 auto&& shp = mgr.infer_shape_fallible(opr->output(0)); \
155 if (!shp) \
156 return "Unsupported opr, because operator shape cannot be " \
157 "inferred at compile time."; \
158 else \
159 return None; \
160 }
161 m_opr_trait[opr::Elemwise::typeinfo()].get_replace_fail_msg =
162 [](OperatorNodeBase* opr) -> Maybe<std::string> {
163 bool has_scalar = false;
164 for (auto&& inp : opr->input()) {
165 if (inp->shape().is_scalar()) {
166 has_scalar = true;
167 break;
168 }
169 }
170 if (has_scalar)
171 return "Elemwise with scalar input is not supported.";
172 if (opr->input(0)->dtype().enumv() != DTypeEnum::QuantizedS8 &&
173 opr->input(0)->dtype() != dtype::Float32()) {
174 return "Unsupported data type.";
175 }
176 using Mode = opr::Elemwise::Mode;
177 static const ThinHashSet<Mode> supported_modes {
178#if NV_TENSOR_RT_VERSION >= 5105
179 Mode::SIN, Mode::COS, Mode::ASIN, Mode::ACOS, Mode::CEIL, Mode::FLOOR,
180#endif
181 Mode::EXP, Mode::LOG, Mode::ABS,
182
183 Mode::RELU, Mode::SIGMOID, Mode::TANH, Mode::ADD, Mode::MUL,
184 Mode::MIN, Mode::MAX, Mode::SUB, Mode::TRUE_DIV, Mode::POW,
185 Mode::FUSE_ADD_RELU, Mode::FUSE_ADD_TANH, Mode::FUSE_ADD_SIGMOID
186 };
187 auto mode = opr->cast_final_safe<opr::Elemwise>().param().mode;
188 if (!supported_modes.count(mode)) {
189 return "Unsupported Elemwise mode.";
190 }
191#if NV_TENSOR_RT_VERSION >= 6001
192 if (opr->input(0)->dtype().enumv() == DTypeEnum::QuantizedS8) {
193 TensorShapeArray inps;
194 for (auto&& inp : opr->input()) {
195 inps.push_back(inp->shape());
196 }
197 TensorShape brdcast;
198 megdnn::Elemwise::deduce_shape(inps, brdcast);
199 if (brdcast.ndim < 4) {
200 return "Elemwise with QuantizedS8 data type must have more "
201 "than 4 dimensions. Less than 3 dimensions is not "
202 "supported since trt6.0.";

Callers

nothing calls this directly

Calls 15

get_valueFunction · 0.85
SymbolVarClass · 0.85
subFunction · 0.85
make_rewriterMethod · 0.80
is_constMethod · 0.80
makeFunction · 0.50
graphMethod · 0.45
inputMethod · 0.45
is_scalarMethod · 0.45
shapeMethod · 0.45
enumvMethod · 0.45
dtypeMethod · 0.45

Tested by

no test coverage detected