| 111 | MGB_DYN_TYPE_OBJ_FINAL_IMPL(ProxyGraph::InputPlaceholder); |
| 112 | |
| 113 | class ProxyGraph::StaticInferManager : public cg::static_infer::StaticInferManager { |
| 114 | public: |
| 115 | using Tag = cg::static_infer::Tag; |
| 116 | using ShapeInferDesc = cg::static_infer::ShapeInferDesc; |
| 117 | using ValueInferDesc = cg::static_infer::ValueInferDesc; |
| 118 | using InferType = cg::static_infer::InferType; |
| 119 | using DepVal = cg::static_infer::DepVal; |
| 120 | using DepElement = cg::static_infer::DepElement; |
| 121 | using DepType = cg::static_infer::DepType; |
| 122 | using InpElement = cg::static_infer::InpElement; |
| 123 | |
| 124 | struct Result { |
| 125 | TensorShape shape; |
| 126 | DeviceTensorND value; |
| 127 | }; |
| 128 | |
| 129 | ProxyGraph* owner; |
| 130 | cg::OperatorNodeBase* cur_opr = nullptr; |
| 131 | std::vector<std::optional<ShapeInferDesc>> shape_descs; |
| 132 | std::vector<std::optional<ValueInferDesc>> value_descs; |
| 133 | std::vector<Result> inferred_outputs; |
| 134 | |
| 135 | StaticInferManager(ProxyGraph* owner_) : owner(owner_) {} |
| 136 | |
| 137 | size_t locate_output(VarNode* var) { |
| 138 | mgb_assert(cur_opr); |
| 139 | auto&& output_vars = cur_opr->output(); |
| 140 | mgb_assert(shape_descs.size() == output_vars.size()); |
| 141 | auto&& it = std::find(output_vars.begin(), output_vars.end(), var); |
| 142 | mgb_assert(it != output_vars.end()); |
| 143 | return it - output_vars.begin(); |
| 144 | } |
| 145 | |
| 146 | void register_shape_infer(Tag dest, const ShapeInferDesc& desc) override { |
| 147 | auto i = locate_output(dest); |
| 148 | mgb_assert(!shape_descs[i]); |
| 149 | shape_descs[i].emplace(desc); |
| 150 | } |
| 151 | |
| 152 | void register_value_infer(Tag dest, const ValueInferDesc& desc) override { |
| 153 | auto i = locate_output(dest); |
| 154 | mgb_assert(!value_descs[i]); |
| 155 | value_descs[i].emplace(desc); |
| 156 | } |
| 157 | |
| 158 | InferType get_infer_type(Tag var) override { |
| 159 | // don't let opr apply any immediate optimization |
| 160 | return {InferType::MISSING_INP, InferType::MISSING_INP}; |
| 161 | } |
| 162 | |
| 163 | void update() { |
| 164 | if (cur_opr != owner->m_cur_opr) { |
| 165 | clear(); |
| 166 | cur_opr = owner->m_cur_opr; |
| 167 | if (cur_opr) { |
| 168 | auto nout = cur_opr->output().size(); |
| 169 | shape_descs.resize(nout); |
| 170 | value_descs.resize(nout); |