| 26 | class PluginSizeComputer : public SizeComputer { |
| 27 | public: |
| 28 | virtual bool onComputeSize(const MNN::Op* op, const std::vector<Tensor*>& inputs, |
| 29 | const std::vector<Tensor*>& outputs) const override { |
| 30 | // Plugin op should has inputs or outputs, or both. |
| 31 | MNN_CHECK(inputs.size() > 0 || outputs.size() > 0, // NOLINT |
| 32 | "Plugin op should has inputs or outputs, or both of them."); |
| 33 | |
| 34 | #ifdef MNN_WITH_PLUGIN |
| 35 | const Plugin* plugin_param = op->main_as<Plugin>(); |
| 36 | std::shared_ptr<plugin::InferShapeKernel> kernel = // NOLINT |
| 37 | getInferShapeKernel(plugin_param->type()->str()); |
| 38 | MNN_CHECK(nullptr != kernel.get(), // NOLINT |
| 39 | "Shape inference kernel has not been registered for plugin op."); |
| 40 | |
| 41 | plugin::InferShapeContext ctx(inputs, outputs); |
| 42 | for (const Attribute* attr : *(plugin_param->attr())) { |
| 43 | ctx.setAttr(attr->key()->str(), attr); |
| 44 | } |
| 45 | bool status = kernel->compute(&ctx); |
| 46 | if (!status) { |
| 47 | MNN_ERROR("Plugin op infer shape failed with false returned."); |
| 48 | } |
| 49 | return status; |
| 50 | #else |
| 51 | MNN_ERROR("Plugin is not supported. Please recompile with `MNN_WITH_PLUGIN` enabled."); |
| 52 | return false; |
| 53 | #endif // MNN_WITH_PLUGIN |
| 54 | } |
| 55 | }; |
| 56 | |
| 57 | REGISTER_SHAPE(PluginSizeComputer, OpType_Plugin); |
nothing calls this directly
no test coverage detected