| 128 | }; |
| 129 | |
| 130 | class UserOpInferContext final : public user_op::InferContext { |
| 131 | public: |
| 132 | using ArgVec = std::vector<std::pair<std::string, int32_t>>; |
| 133 | |
| 134 | UserOpInferContext(const UserOp* op, const ParallelContext* parallel_ctx, const JobDesc* job_desc, |
| 135 | const std::function<BlobDesc*(const std::string&)>& GetBlobDesc4BnInOp) |
| 136 | : op_(op), parallel_ctx_(parallel_ctx), job_desc_(job_desc) { |
| 137 | bn2logical_tensor_desc_.reset(new HashMap<std::string, user_op::NaiveTensorDesc>()); |
| 138 | auto InitTensorDesc = [&](const ArgVec& arg_vec, const PbRpf<std::string>& bns) { |
| 139 | CHECK_EQ(arg_vec.size(), bns.size()); |
| 140 | for (int32_t i = 0; i < arg_vec.size(); ++i) { |
| 141 | const auto& bn_i = bns.Get(i); |
| 142 | BlobDesc* blob = GetBlobDesc4BnInOp(bns.Get(i)); |
| 143 | CHECK(blob != nullptr) << bn_i; |
| 144 | arg2tensor_desc_.emplace(arg_vec.at(i), GenTensorDescFromBlobDesc(blob)); |
| 145 | } |
| 146 | }; |
| 147 | InitTensorDesc(op->inputs(), op->input_bns()); |
| 148 | InitTensorDesc(op->outputs(), op->output_bns()); |
| 149 | } |
| 150 | ~UserOpInferContext() override = default; |
| 151 | |
| 152 | const user_op::TensorDesc& InputTensorDesc(const std::string& arg_name, |
| 153 | int32_t index) const override { |
| 154 | return *TensorDesc4ArgNameAndIndex(arg_name, index); |
| 155 | } |
| 156 | const user_op::TensorDesc& OutputTensorDesc(const std::string& arg_name, |
| 157 | int32_t index) const override { |
| 158 | return *TensorDesc4ArgNameAndIndex(arg_name, index); |
| 159 | } |
| 160 | user_op::TensorDesc* MutOutputTensorDesc(const std::string& arg_name, int32_t index) override { |
| 161 | return MutTensorDesc4ArgNameAndIndex(arg_name, index); |
| 162 | } |
| 163 | const user_op::TensorDesc* TensorDesc4ArgNameAndIndex(const std::string& arg_name, |
| 164 | int32_t index) const { |
| 165 | auto it = arg2tensor_desc_.find(std::make_pair(arg_name, index)); |
| 166 | if (it == arg2tensor_desc_.end()) { return nullptr; } |
| 167 | return &it->second; |
| 168 | } |
| 169 | user_op::TensorDesc* MutTensorDesc4ArgNameAndIndex(const std::string& arg_name, int32_t index) { |
| 170 | auto it = arg2tensor_desc_.find(std::make_pair(arg_name, index)); |
| 171 | if (it == arg2tensor_desc_.end()) { return nullptr; }; |
| 172 | return &(it->second); |
| 173 | } |
| 174 | const user_op::TensorDesc* LogicalTensorDesc4ArgNameAndIndex(const std::string& arg_name, |
| 175 | int32_t index) const override { |
| 176 | const std::string bn = GenRepeatedBn(arg_name, index); |
| 177 | const auto it = bn2logical_tensor_desc_->find(bn); |
| 178 | if (it != bn2logical_tensor_desc_->end()) { |
| 179 | return &it->second; |
| 180 | } else { |
| 181 | std::shared_ptr<const BlobDesc> blob_desc = CHECK_JUST(op_->GetLogicalBlobDesc4BnInOp(bn)); |
| 182 | bn2logical_tensor_desc_->emplace(bn, GenTensorDescFromBlobDesc(blob_desc.get())); |
| 183 | return &(bn2logical_tensor_desc_->emplace(bn, GenTensorDescFromBlobDesc(blob_desc.get())) |
| 184 | .first->second); |
| 185 | } |
| 186 | } |
| 187 | const Shape& InputShape(const std::string& arg_name, int32_t index) const override { |
nothing calls this directly
no test coverage detected