| 197 | } |
| 198 | |
| 199 | Maybe<const LocalTensorInferResult> LocalTensorInferCache::GetOrInfer( |
| 200 | const LocalTensorMetaInferArgs& infer_args) { |
| 201 | if (ThreadLocalEnvBool<ONEFLOW_EAGER_ENABLE_LOCAL_INFER_CACHE>()) { |
| 202 | auto iter = cache_.find(infer_args); |
| 203 | if (iter == cache_.end()) { |
| 204 | if (unlikely(cache_.size() |
| 205 | >= ThreadLocalEnvInteger<ONEFLOW_EAGER_TENSOR_INFER_CACHE_SIZE>())) { |
| 206 | cache_.clear(); |
| 207 | } |
| 208 | const auto& user_op_expr = user_op_expr_.lock(); |
| 209 | CHECK_OR_RETURN(static_cast<bool>(user_op_expr)); // NOLINT |
| 210 | const auto& output_tensor_metas = JUST(Infer(*user_op_expr, infer_args)); |
| 211 | iter = cache_.emplace(infer_args, output_tensor_metas).first; |
| 212 | } |
| 213 | return iter->second; |
| 214 | } else { |
| 215 | const auto& user_op_expr = user_op_expr_.lock(); |
| 216 | return JUST(Infer(*user_op_expr, infer_args)); |
| 217 | } |
| 218 | } |
| 219 | |
| 220 | } // namespace one |
| 221 | } // namespace oneflow |